#include "arguments.h"
#include "span_context.h"
#include "go_context.h"
#include "go_types.h"
#include "uprobe.h"

#define MAX_DATA_SIZE 64

#define ENTPOINT 0
#define RETPOINT 1

// bpf write event data when enter/exit a function
struct event
{
	__u64 pid;
	__u64 trace_id;
	__u64 goid;
	__u64 ip;
	__u64 bp;
	__u64 caller_ip;
	__u64 caller_bp;
	__u64 time_ns_start;
	__u64 time_ns_end;
	__u64 nid;
	__u64 fpid;
	__u64 level;
	__u8 location;
}__attribute__((packed));

struct trace_stack_entry_index_key_t {
	__u32 trace_id;
	__u64 goid;
};

struct trace_stack_entry_index_value_t {
	__u64 fun_key;
};

struct trace_stack_entry_key_t {
	__u32 trace_id;
	__u64 goid;
	__u64 fun_key;
};

struct {
    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
    __uint(key_size, sizeof(int));
    __uint(value_size, sizeof(int));
} event_queue SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct trace_stack_entry_key_t));
	__uint(value_size, sizeof(struct event));
	__uint(max_entries, 32768);
} trace_stack_entry SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct trace_stack_entry_index_key_t));
	__uint(value_size, sizeof(__u64));
	__uint(max_entries, 32768);
} trace_stack_entry_index SEC(".maps");

struct bpf_map_def SEC("maps") event_stack = {
	.type = BPF_MAP_TYPE_PERCPU_ARRAY,
	.key_size = sizeof(__u32),
	.value_size = sizeof(struct event),
	.max_entries = 1,
};

struct bpf_map_def SEC("maps") should_trace_goid = {
	.type = BPF_MAP_TYPE_HASH,
	.key_size = sizeof(__u64),
	.value_size = sizeof(bool),
	.max_entries = 10000,
};

SEC("uprobe/ent")
int ent(struct pt_regs *ctx)
{
    bpf_printk("[Go] [uprobe/ent] enter");
	__u64 pid_tgid = bpf_get_current_pid_tgid();
    __u64 pid = pid_tgid >> 32;

	__u64 trace_id = get_apm_trace_id(pid, pid_tgid);

	__u32 key = 0;
	struct event *e = bpf_map_lookup_elem(&event_stack, &key);
	if (!e)
	{
		bpf_printk("[Go] [uprobe/ret]: conot get event");
		return 0;
	}
	__builtin_memset(e, 0, sizeof(*e));

	// e->goid = get_goid();
	e->goid = get_current_goroutine();

	// e->ip = ctx->ip;
	e->ip = PT_REGS_IP(ctx); // 当前函数的地址
	if (!bpf_map_lookup_elem(&should_trace_goid, &e->goid))
	{
		__u64 should_trace = true;
		bpf_map_update_elem(&should_trace_goid, &e->goid, &should_trace, BPF_ANY);
	}
	bpf_printk("[Go] [uprobe/ent]: e->goid: %llu", e->goid);
	bpf_printk("[Go] [uprobe/ent]: e->ip: %d", e->ip);
	bpf_printk("[Go] [uprobe/ent]: yes");

	e->pid = pid;
	e->trace_id = trace_id;
	e->location = ENTPOINT; // 标识入口还是出口
	e->time_ns_start = bpf_ktime_get_ns();
	e->bp = PT_REGS_SP(ctx) - 8; // 因为 ebpf 程序执行时实际上还没进行堆栈的切换，所以此处是需要在父函数中获取子函数的基地址，新的堆栈会将父函数的 rbp 入栈，一个 rbp 在 64位机器上占用 8 个字节，所以需要减 8，获取到子函数的基地址，低地址的需要减8，如果到高地址的cpu 上就得+8
	e->caller_bp = PT_REGS_FP(ctx); // 父函数的基地址

	void *ra;
	ra = (void *)PT_REGS_SP(ctx); // 父函数的 sp，此时 sp 指向父函数的最后一个地址
	bpf_probe_read_user(&e->caller_ip, sizeof(e->caller_ip), ra);

	bpf_printk("[Go] [uprobe/ent]: goid: %llu", e->goid);

	bpf_printk("[Go] [uprobe/ent]: event: location:%x,ip:%x,time_ns_start:%lld\n", e->location, e->ip, e->time_ns_start);
	bpf_printk("[Go] [uprobe/ent]: event: bp:%x,caller_bp:%x,caller_ip:%x\n", e->bp,e->caller_bp,e->caller_ip);

	struct trace_stack_entry_index_key_t trace_index_key = {};
	trace_index_key.goid = e->goid;
	trace_index_key.trace_id = trace_id;

	bpf_printk("[Go] [uprobe/ent]: trace_stack_entry_index_key_t: goid:%d,trace_id:%lld\n", trace_index_key.goid,trace_index_key.trace_id);
	__u64 *fun_key_p = bpf_map_lookup_elem(&trace_stack_entry_index, &trace_index_key);
	__u64 fun_key = 1;
	__u64 fpid = e->goid;
	if (fun_key_p) {
		bpf_printk("[Go] [uprobe/ent]: get fun_key: %d", *fun_key_p);
		struct trace_stack_entry_key_t trace_key_parent = {};
		trace_key_parent.goid = e->goid;
		trace_key_parent.trace_id = trace_id;
		trace_key_parent.fun_key = *fun_key_p;

		struct event *event_parent = bpf_map_lookup_elem(&trace_stack_entry, &trace_key_parent);
		if (!event_parent) {
			bpf_printk("[Go] [uprobe/ent]: get event_parent Error Error: %d", *fun_key_p);
		} else {
			fpid = event_parent->nid;
		}
		fun_key = *fun_key_p + 1;
	}

	struct trace_stack_entry_key_t trace_key = {};
	trace_key.goid = e->goid;
	trace_key.trace_id = trace_id;
	trace_key.fun_key = fun_key;

	struct event event_current = {};
	event_current.bp = e->bp;
	event_current.caller_bp = e->caller_bp;
	event_current.caller_ip = e->caller_ip;
	event_current.pid = e->pid;
	event_current.trace_id = e->trace_id;
	event_current.goid = e->goid;
	event_current.ip = e->ip;
	event_current.time_ns_start = e->time_ns_start;
	event_current.location = e->location;
	event_current.nid = bpf_get_prandom_u32() & 0xFF;
	event_current.fpid = fpid;
	// 第一级是 appliction，二级是 协程
	event_current.level = fun_key + 2;

	bpf_printk("[Go] [uprobe/ent]: uodate fun_key: %d", fun_key);
	bpf_map_update_elem(&trace_stack_entry, &trace_key, &event_current, BPF_ANY);
	bpf_map_update_elem(&trace_stack_entry_index, &trace_index_key, &fun_key, BPF_ANY);

	// return bpf_map_push_elem(&event_queue, e, BPF_EXIST);
	// bpf_perf_event_output(ctx, &event_queue, BPF_F_CURRENT_CPU, e, sizeof(*e));

	bpf_printk("[Go] [uprobe/ent] end");
	return 1;
}

SEC("uprobe/ret")
int ret(struct pt_regs *ctx)
{
    bpf_printk("[Go] [uprobe/ret] enter");
	__u64 pid_tgid = bpf_get_current_pid_tgid();
    __u64 pid = pid_tgid >> 32;

	__u64 trace_id = get_apm_trace_id(pid, pid_tgid);

	__u32 key = 0;
	struct event *e = bpf_map_lookup_elem(&event_stack, &key);
	if (!e)
	{
		bpf_printk("[Go] [uprobe/ret]: conot get event");
		return 0;
	}
	__builtin_memset(e, 0, sizeof(*e));

	// e->goid = get_goid();
	e->goid = get_current_goroutine();
	if (!bpf_map_lookup_elem(&should_trace_goid, &e->goid))
	{
		bpf_printk("[Go] [uprobe/ret]: conot get event: %llu", e->goid);
		return 0;
	}
	bpf_printk("[Go] [uprobe/ret]: yes");
	e->pid = pid;
	e->trace_id = trace_id;
	e->location = RETPOINT;
	e->ip = PT_REGS_IP(ctx);
	e->time_ns_end = bpf_ktime_get_ns();
	e->bp = PT_REGS_FP(ctx);
	bpf_printk("[Go] [uprobe/ret]: e->ip:%x,bp:%x", e->ip, e->bp);
	bpf_printk("[Go] [uprobe/ret]: goid: %llu", e->goid);

	bpf_printk("[Go] [uprobe/ret]: event: location:%x,ip:%x,time_ns_end:%lld\n", e->location, e->ip, e->time_ns_end);

	struct trace_stack_entry_index_key_t trace_index_key = {};
	trace_index_key.goid = e->goid;
	trace_index_key.trace_id = trace_id;

	__u64 *fun_key_p = bpf_map_lookup_elem(&trace_stack_entry_index, &trace_index_key);
	__u64 fun_key = 1;
	if (fun_key_p) {
		bpf_printk("[Go] [uprobe/ret]: get fun_key: %d", *fun_key_p);
		fun_key = *fun_key_p;
	} else {
		bpf_printk("[Go] [uprobe/ret]:ErrorErrorErrorError Not get fun_key");
	}
	struct trace_stack_entry_key_t trace_key = {};
	trace_key.goid = e->goid;
	trace_key.trace_id = trace_id;
	trace_key.fun_key = fun_key;

	struct event *event_p = bpf_map_lookup_elem(&trace_stack_entry, &trace_key);

	if (!event_p) {
		bpf_printk("[Go] [uprobe/ret]:ErrorErrorErrorError Not get funEntry");
		return 0;
	}

	event_p->time_ns_end = e->time_ns_end;

	bpf_printk("[Go] [uprobe/ret]: ent:event: location:%d,ip:%x,time_ns_start:%lld\n", event_p->location, event_p->ip, event_p->time_ns_start);
	bpf_printk("[Go] [uprobe/ret]: ent:event: bp:%x,caller_bp:%x,caller_ip:%x\n", event_p->bp,event_p->caller_bp,event_p->caller_ip);
	bpf_printk("[Go] [uprobe/ret]: ent:event: nid:%d,fpid:%d,level:%d\n", event_p->nid,event_p->fpid,event_p->level);

	bpf_map_delete_elem(&trace_stack_entry, &trace_key);
	__u64 fun_key_new = fun_key - 1;
	bpf_map_update_elem(&trace_stack_entry_index, &trace_index_key, &fun_key_new, BPF_ANY);

	// return bpf_map_push_elem(&event_queue, e, BPF_EXIST);
	// 最后一个函数结束，同时推送协程信息
	if (fun_key == 1) {
		struct event event_coroutine = {};
		event_coroutine.bp = e->bp;
		event_coroutine.caller_bp = e->caller_bp;
		event_coroutine.caller_ip = e->caller_ip;
		event_coroutine.pid = e->pid;
		event_coroutine.trace_id = e->trace_id;
		event_coroutine.goid = e->goid;
		event_coroutine.ip = e->ip;
		event_coroutine.time_ns_start = e->time_ns_start;
		event_coroutine.location = 2;
		event_coroutine.nid = e->goid;
		event_coroutine.fpid = 1;
		event_coroutine.level = 2;
		bpf_printk("[Go] [uprobe/ret]: ent:event_coroutine push: nid:%d,fpid:%d,level:%d\n", event_coroutine.nid,event_coroutine.fpid,event_coroutine.level);
		bpf_perf_event_output(ctx, &event_queue, BPF_F_CURRENT_CPU, &event_coroutine, sizeof(event_coroutine));
	}
	bpf_printk("[Go] [uprobe/ret]: ent:event_ret push: nid:%d,fpid:%d,level:%d\n", event_p->nid,event_p->fpid,event_p->level);
	bpf_perf_event_output(ctx, &event_queue, BPF_F_CURRENT_CPU, event_p, sizeof(*e));

	bpf_printk("[Go] [uprobe/ret] end");
	return 1;
}




