/**
@file     task_info_map_common.h
@brief    Header of task info map common functions
@details  Copyright (c) 2025 Acronis International GmbH
@author   Bruce Wang (bruce.wang@acronis.com)
@since    $Id: $
*/

#pragma once

#ifndef BPF_PROGRAM
#include <linux/ktime.h>
#include "task_info_map.h"
#else
#include "ebpf_programs.bpf.h"
#include "file_key_tools.h"
#endif

static void task_info_set_pid_version(task_info_t *task_info, uint64_t pid_version, uint64_t artificial_start_time_us)
{
    WRITE_ONCE(task_info->pid_version, pid_version);
    WRITE_ONCE(task_info->artificial_start_time_us, artificial_start_time_us);
}

static bool task_info_may_need_refresh(const task_info_t *task_info, const file_key_t *key)
{
    file_key_t exe_file_key;
    exe_file_key.ptr = READ_ONCE(task_info->exe_file_key.ptr);
    exe_file_key.ino = READ_ONCE(task_info->exe_file_key.ino);
    exe_file_key.gen = READ_ONCE(task_info->exe_file_key.gen);
    exe_file_key.dev = READ_ONCE(task_info->exe_file_key.dev);
    return 0 != cmp_file_key(&exe_file_key, key);
}

#ifndef BPF_PROGRAM
static task_info_map_t global_task_info_map;
static void task_info_context_invalidate(transport_task_context_t *context);
#endif

static bool task_info_refresh(task_info_t *task_info, const struct path *exe_path, bool force_advance_pid_version)
{
    file_key_t key;
    int i;

    if (!exe_path)
        return false;

    make_key(&key, exe_path);
    if (!force_advance_pid_version)
    {
        if (!task_info_may_need_refresh(task_info, &key))
            return false;
    }
#ifdef BPF_PROGRAM
    uint64_t boot_time = bpf_ktime_get_boot_ns() / 1000;
#endif
    spin_lock(&task_info->spinlock);
    if (force_advance_pid_version || 0 != cmp_file_key(&key, &task_info->exe_file_key))
    {
        // If 'exe_file_key.ptr==0', exe was not provided on creation. This means that we get to keep the same 'pid_version'
        // but 'exe_file_key' should be updated to the current 'exe'. Keep the normal logic for 'force_advance_pid_version'.
        if (force_advance_pid_version || 0 != task_info->exe_file_key.ptr)
        {
#ifndef BPF_PROGRAM
            task_info_set_pid_version(task_info, atomic64_add_return(1, &global_task_info_map.next_pid_version), ktime_to_us(ktime_get_real()));
            for (i = 0; i < MAX_TRANSPORT_SIZE; i++)
            {
                transport_task_context_t *context = &task_info->contexts[i];
                task_info_context_invalidate(context);
            }
#else
            task_info_set_pid_version(task_info, __sync_fetch_and_add(&g_pid_version, 1), boot_time);
#endif
        }
        WRITE_ONCE(task_info->exe_file_key.ptr, key.ptr);
        WRITE_ONCE(task_info->exe_file_key.ino, key.ino);
        WRITE_ONCE(task_info->exe_file_key.gen, key.gen);
        WRITE_ONCE(task_info->exe_file_key.dev, key.dev);
    }
    spin_unlock(&task_info->spinlock);
    return true;
}

#ifndef BPF_PROGRAM
static void task_info_context_init(transport_task_context_t *context, transport_id_t id);
static task_info_t *task_info_init(task_info_t *task_info, pid_t pid, uint64_t unique_pid, uint64_t pid_version, const struct path *exe_path, uint64_t start_time_us)
#else
static task_info_t *task_info_init(task_info_t *task_info, pid_t pid, uint64_t unique_pid, uint64_t pid_version, const struct path *exe_path, const file_key_t *key, uint64_t start_time_us)
#endif
{
    int i;
#ifndef BPF_PROGRAM
    DPRINTF("task_info=%p pid=%i", task_info, pid);
    RB_CLEAR_NODE(&task_info->rb_node);
    task_info->pid = pid;
    atomic_set(&task_info->ref_cnt, 1);
    spin_lock_init(&task_info->spinlock);
#endif
    task_info->pid = pid;
    task_info->unique_pid = unique_pid;
    task_info_set_pid_version(task_info, pid_version, start_time_us);
    if (exe_path)
#ifndef BPF_PROGRAM
        make_key(&task_info->exe_file_key, exe_path);
#else
        task_info->exe_file_key = *key;
#endif
    else
        task_info->exe_file_key = (file_key_t){0};

#ifndef BPF_PROGRAM
    for (i = 0; i < MAX_TRANSPORT_SIZE; i++)
    {
        transport_id_t id = transport_id_make_not_taken(i);
        transport_task_context_t *context = &task_info->contexts[i];
        task_info_context_init(context, id);
    }
#endif

    return task_info;
}