// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef _ALLOC_H_
#define _ALLOC_H_

//#include "bpf_helpers.h"

#define MAX_ENTRIES 50
#define MAX_BUFFER_SIZE 1024
#define MIN_BUFFER_SIZE 1

// Injected in init
// volatile const u64 total_cpus;
// volatile const u64 start_addr;
// volatile const u64 end_addr;

struct
{
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __type(key, s32);
    __type(value, u64);
    __uint(max_entries, MAX_ENTRIES);
//    __uint(pinning, LIBBPF_PIN_BY_NAME);
} alloc_map SEC(".maps");

static __always_inline u64 get_area_start(u64 start_addr,u64 end_addr)
{
    u32 k0 = 0;
	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
	if (!trace_conf) {
        return 0;
	}
    s64 partition_size = (end_addr - start_addr) / trace_conf->total_cpus;
    u32 current_cpu = bpf_get_smp_processor_id();
    s32 start_index = 0;
    u64 *start = (u64 *)bpf_map_lookup_elem(&alloc_map, &start_index);
    if (start == NULL || *start == 0)
    {
        u64 current_start_addr = start_addr + (partition_size * current_cpu);
        bpf_map_update_elem(&alloc_map, &start_index, &current_start_addr, BPF_ANY);
        return current_start_addr;
    }
    else
    {
        return *start;
    }
}

static __always_inline u64 get_area_end(u64 start,u64 start_addr,u64 end_addr)
{
    u32 k0 = 0;
	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
	if (!trace_conf) {
        return 0;
	}
    s64 partition_size = (end_addr - start_addr) / trace_conf->total_cpus;
    s32 end_index = 1;
    u64 *end = (u64 *)bpf_map_lookup_elem(&alloc_map, &end_index);
    if (end == NULL || *end == 0)
    {
        u64 current_end_addr = start + partition_size;
        bpf_map_update_elem(&alloc_map, &end_index, &current_end_addr, BPF_ANY);
        return current_end_addr;
    }
    else
    {
        return *end;
    }
}

static __always_inline s32 bound_number(s32 num, s32 min, s32 max)
{
    if (num < min)
    {
        return min;
    }
    else if (num > max)
    {
        return max;
    }
    return num;
}

static __always_inline void *write_target_data(void *data, s32 size)
{
    if (!data || data == NULL)
    {
        return NULL;
    }
	u64 start_from_proc;
	u64 end_from_proc;
	__u32 tgid = (__u32)(bpf_get_current_pid_tgid() >> 32);
	struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
	if (!info) {
		return NULL;
	}
	start_from_proc = info->start_addr;
	end_from_proc = info->end_addr;

	u64 start = get_area_start(start_from_proc, end_from_proc);
    u64 end = get_area_end(start,start_from_proc,end_from_proc);
    if (end - start < size)
    {
        cw_bpf_debug("reached end of CPU memory block, going to the start again");
        s32 start_index = 0;
        bpf_map_delete_elem(&alloc_map, &start_index);
        start = get_area_start(start_from_proc, end_from_proc);
    }

    void *target = (void *)start;
    size = bound_number(size, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
    u64 page_offset = (u64)target & 0xFFF;
    u64 dist_to_next_page = 4096 - page_offset;
    if (dist_to_next_page < size)
    {
        target += dist_to_next_page;
    }
    u64 target_u = (u64)target;
    if (target_u > end_from_proc || target_u < start_from_proc) {
	    cw_bpf_debug("TARGET ADDRESS IS OUT OF BOUNDS: 0x%llx", target);
        return NULL;
    }

    long success = bpf_probe_write_user(target, data, size);
    if (success == 0)
    {
        s32 start_index = 0;
        // Update the start position of this chunk, taking into account possible adjustments
        // we made to be page aligned
        u64 updated_start = target_u + size;

        // align updated_start to 8 bytes
        if (updated_start % 8 != 0) {
            updated_start += 8 - (updated_start % 8);
        }

        bpf_map_update_elem(&alloc_map, &start_index, &updated_start, BPF_ANY);
        return target;
    }
    else
    {
	    cw_bpf_debug("failed to write to userspace, error code: %d, addr: %lx, size: %d", success, target, size);
        return NULL;
    }
}

static __always_inline void *cw_write_target_data(void *data, s32 size, struct ebpf_proc_info* info)
{
    if (!data || data == NULL)
    {
        return NULL;
    }
	u64 start_from_proc;
	u64 end_from_proc;
//	__u32 tgid = (__u32)(bpf_get_current_pid_tgid() >> 32);
//	struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
//	if (!info) {
//		return NULL;
//	}
	start_from_proc = info->start_addr;
	end_from_proc = info->end_addr;

	u64 start = get_area_start(start_from_proc, end_from_proc);
    u64 end = get_area_end(start,start_from_proc,end_from_proc);
    if (end - start < size)
    {
        cw_bpf_debug("reached end of CPU memory block, going to the start again");
        s32 start_index = 0;
        bpf_map_delete_elem(&alloc_map, &start_index);
        start = get_area_start(start_from_proc, end_from_proc);
    }

    void *target = (void *)start;
    size = bound_number(size, MIN_BUFFER_SIZE, MAX_BUFFER_SIZE);
    u64 page_offset = (u64)target & 0xFFF;
    u64 dist_to_next_page = 4096 - page_offset;
    if (dist_to_next_page < size)
    {
        target += dist_to_next_page;
    }
    u64 target_u = (u64)target;
    if (target_u > end_from_proc || target_u < start_from_proc) {
	    cw_bpf_debug("TARGET ADDRESS IS OUT OF BOUNDS: 0x%llx", target);
        return NULL;
    }

    long success = bpf_probe_write_user(target, data, size);
    if (success == 0)
    {
        s32 start_index = 0;
        // Update the start position of this chunk, taking into account possible adjustments
        // we made to be page aligned
        u64 updated_start = target_u + size;

        // align updated_start to 8 bytes
        if (updated_start % 8 != 0) {
            updated_start += 8 - (updated_start % 8);
        }

        bpf_map_update_elem(&alloc_map, &start_index, &updated_start, BPF_ANY);
        return target;
    }
    else
    {
        bpf_printk("failed to write to userspace, error code: %d, addr: %lx, size: %d", success, target, size);
        return NULL;
    }
}

#endif
