#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;
	__u8 location;
}__attribute__((packed));

struct trace_stack_entry_key_t {
	__u64 trace_id;
	__u64 caller_bp;
	__u64 bp;
};

struct {
    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
    __uint(key_size, sizeof(int));
    __uint(value_size, sizeof(int));
	__uint(max_entries, 655351);
} 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 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_LRU_HASH,
	.key_size = sizeof(__u64),
	.value_size = sizeof(bool),
	.max_entries = 32768,
};

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

	struct ebpf_proc_info *info =
			bpf_map_lookup_elem(&proc_info_map, &pid);
	if (!info) {
		return 0;
	}

	cw_bpf_debug("[Go] [uprobe/ent]: proc_info_map::%ld, %d, %d\n", info->code_type);

	__u64 trace_id = get_apm_trace_id(pid, pid_tgid);

	cw_bpf_debug("[Go] [uprobe/ent]: get_apm_trace_id::%ld, %d, %d\n", trace_id, pid, pid_tgid);

	__u32 key = 0;
	struct event *e = bpf_map_lookup_elem(&event_stack, &key);
	if (!e)
	{
		cw_bpf_debug("[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);
	}
	cw_bpf_debug("[Go] [uprobe/ent]: e->goid: %llu", e->goid);
	cw_bpf_debug("[Go] [uprobe/ent]: e->ip: %d", e->ip);
	cw_bpf_debug("[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);

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

	cw_bpf_debug("[Go] [uprobe/ent]: event: location:%x,ip:%lx,time_ns_start:%llu\n", e->location, e->ip, e->time_ns_start);
	cw_bpf_debug("[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_key_t trace_key = {};
	trace_key.caller_bp = e->caller_bp % 0x800 + (e->goid << 12);
	trace_key.bp = e->bp % 0x800 + (e->goid << 12);
	trace_key.trace_id = trace_id;

	// trace_key.caller_bp = e->caller_bp;
	// trace_key.bp = e->bp;
	// cw_bpf_debug("[Go] [uprobe/ent]: trace_keytrace_keytrace_key: bp: %x", trace_key.bp);
	cw_bpf_debug("[Go] [uprobe/ent]: trace_keytrace_keytrace_key: caller_bp: %x, bp: %x, %lld", trace_key.caller_bp, trace_key.bp, trace_key.trace_id);

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

	if (event_p) {
		cw_bpf_debug("[Go] [uprobe/ent]:Error get alike funEntry %x, %x, %x", event_p->ip, event_p->caller_bp, event_p->bp);
		// return 0;
	}

	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,

	bpf_map_update_elem(&trace_stack_entry, &trace_key, &event_current, 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));

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

SEC("uprobe/ret")
int ret(struct pt_regs *ctx)
{
    cw_bpf_debug("[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)
	{
		cw_bpf_debug("[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))
	{
		cw_bpf_debug("[Go] [uprobe/ret]: conot get event: %llu", e->goid);
		return 0;
	}
	cw_bpf_debug("[Go] [uprobe/ret]: yes");
	e->pid = pid;
	e->trace_id = trace_id;
	e->location = ENTPOINT;
	e->ip = PT_REGS_IP(ctx);
	e->time_ns_end = bpf_ktime_get_ns();
	e->bp = PT_REGS_SP(ctx) - 8;
	e->caller_bp = PT_REGS_FP(ctx);

	// __u64 caller_bp = PT_REGS_FP(ctx);
	cw_bpf_debug("[Go] [uprobe/ret]: e->ip:%lx,bp:%x,caller_bp:%x", e->ip, e->bp, e->caller_bp);
	cw_bpf_debug("[Go] [uprobe/ret]: sp:%x, goid: %llu, goid:0x:%x", PT_REGS_SP(ctx), e->goid, e->goid);

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

	struct trace_stack_entry_key_t trace_key = {};
	trace_key.caller_bp = e->caller_bp % 0x800 + (e->goid << 12);
	trace_key.bp = e->bp % 0x800 + (e->goid << 12);
	trace_key.trace_id = trace_id;

	// trace_key.caller_bp = caller_bp;
	// trace_key.bp = e->bp;

	cw_bpf_debug("[Go] [uprobe/ret]: trace_keytrace_keytrace_key: caller_bp: %x, bp: %x", trace_key.caller_bp, trace_key.bp);
	// cw_bpf_debug("[Go] [uprobe/ret]: trace_keytrace_keytrace_key: bp: %x", trace_key.bp);

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

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

	event_p->time_ns_end = e->time_ns_end;

	cw_bpf_debug("[Go] [uprobe/ret]: ent:event: location:%d,ip:%x,time_ns_start:%llu\n", event_p->location, event_p->ip, event_p->time_ns_start);
	cw_bpf_debug("[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);

	e->time_ns_start = event_p->time_ns_start;
	// e->caller_ip = event_p->caller_ip;
	// e->ip = event_p->ip;
	
	// return bpf_map_push_elem(&event_queue, e, BPF_EXIST);
	cw_bpf_debug("[Go] [uprobe/ret]: ent:event_ret push: event_p->ip:%llx,ip:%llx\n", event_p->ip, e->ip);
	long err =  bpf_perf_event_output(ctx, &event_queue, BPF_F_CURRENT_CPU, e, sizeof(*e));
	if (err == 0) {
		__u32 count = cw_add_event_count(trace_id);
		if (event_p->trace_id != trace_id) {
			cw_bpf_debug("[pref] fuck err<%d> ip->%x->%x", err, PT_REGS_IP(ctx), e->ip);
		}
		cw_bpf_debug("[pref] err<%d> count->%d trace_id->%llu", err, count, trace_id);
	} else {
		cw_bpf_debug("[pref] err<%d> trace_id->%d", err, trace_id);
	}
	bpf_map_delete_elem(&trace_stack_entry, &trace_key);
	cw_bpf_debug("[Go] [uprobe/ret] end");
	return 1;
}



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

	struct ebpf_proc_info *info =
			bpf_map_lookup_elem(&proc_info_map, &pid);
	if (!info) {
		return 0;
	}

	cw_bpf_debug("[Go] [uprobe/ent]: proc_info_map::%ld, %d, %d\n", info->code_type);

	__u64 trace_id = get_apm_trace_id(pid, pid_tgid);
	
	cw_bpf_debug("[Go] [uprobe/ent]: get_apm_trace_id::%ld, %d, %d\n", trace_id, pid, pid_tgid);

	__u32 key = 0;
	struct event *e = bpf_map_lookup_elem(&event_stack, &key);
	if (!e)
	{
		cw_bpf_debug("[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);
	}
	cw_bpf_debug("[Go] [uprobe/ent]: e->goid: %llu", e->goid);
	cw_bpf_debug("[Go] [uprobe/ent]: e->ip: %d", e->ip);
	cw_bpf_debug("[Go] [uprobe/ent]: yes");

	__u64 asm_;
	bpf_probe_read_user(&asm_, sizeof(asm_), (void*)e->ip);

	asm_ = asm_ >> 32;
	cw_bpf_debug("[Go] [uprobe/ent]: e->asm: %x", asm_);

	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) + asm_; // 因为 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);

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

	cw_bpf_debug("[Go] [uprobe/ent]: event: location:%x,ip:%lx,time_ns_start:%lld\n", e->location, e->ip, e->time_ns_start);
	cw_bpf_debug("[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_key_t trace_key = {};
	trace_key.caller_bp = e->caller_bp % 0x800 + (e->goid << 12);
	trace_key.bp = e->bp % 0x800 + (e->goid << 12);
	trace_key.trace_id = trace_id;

	// trace_key.caller_bp = e->caller_bp;
	// trace_key.bp = e->bp;
	// cw_bpf_debug("[Go] [uprobe/ent]: trace_keytrace_keytrace_key: bp: %x", trace_key.bp);
	cw_bpf_debug("[Go] [uprobe/ent]: trace_keytrace_keytrace_key: caller_bp: %x, bp: %x", trace_key.caller_bp, trace_key.bp);

	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,

	bpf_map_update_elem(&trace_stack_entry, &trace_key, &event_current, 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));

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