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

struct jvm_trace_stack_entry_key_t {
	__u64 trace_id;
	__u64 pid;
	__u64 caller_bp;
	__u64 bp;
};

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

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

SEC("uprobe/setNodeEnter")
int setNodeEnter(struct pt_regs *ctx)
{
    cw_bpf_debug("[Go] [uprobe/setNodeEnter] 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/setNodeEnter]: 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/setNodeEnter]: get_apm_trace_id::%ld, %d, %d\n", trace_id, pid, pid_tgid);
	cw_bpf_debug("[Go] [uprobeThread/pidpidpidpid][setNodeEnter][HTTP]pid:[%d]--[%lld]--[%x]", pid_tgid, bpf_ktime_get_ns(), PT_REGS_IP(ctx));

	__u32 key = 0;
	struct event *e = bpf_map_lookup_elem(&jvmevent_stack, &key);
	if (!e)
	{
		cw_bpf_debug("[Go] [uprobe/setNodeEnter]: 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); // 当前函数的地址
	e->caller_bp = (__u64)PT_REGS_PARM5(ctx);
	e->bp = (__u64)PT_REGS_PARM6(ctx);

	e->pid = pid;
	e->trace_id = trace_id;
	e->location = ENTPOINT; // 标识入口还是出口
	e->time_ns_start = bpf_ktime_get_ns();

    // bpf_probe_read_user(&e->caller_bp, sizeof(e->caller_bp), (void *)(PT_REGS_SP(ctx) + 8));  
    // bpf_probe_read_user(&e->bp, sizeof(e->bp), (void *)(PT_REGS_SP(ctx) + 16)); 
	// cw_bpf_debug("[Go] [uprobe/setNodeEnter]: e->arg7: arg8: %d, %d", e->caller_bp, e->bp);

	__u64 arg7, arg8;  
    bpf_probe_read_user(&arg7, sizeof(arg7), (void *)(PT_REGS_SP(ctx) + 8));  
    bpf_probe_read_user(&arg8, sizeof(arg8), (void *)(PT_REGS_SP(ctx) + 16)); 

	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: e->arg7: arg8: %d, %d", arg7, arg8);

	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: e->caller_bp: %d", e->caller_bp);
	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: e->bp: %d", e->bp);
	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: e->ip: %d", e->ip);
	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: yes");

	struct jvm_trace_stack_entry_key_t trace_key = {};
	trace_key.caller_bp = e->caller_bp;
	trace_key.bp = e->bp;
	trace_key.pid = pid;
	trace_key.trace_id = trace_id;

	cw_bpf_debug("[Go] [uprobe/ent]: trace_keytrace_keytrace_key: trace_id: %llu", trace_key.trace_id);
	cw_bpf_debug("[Go] [uprobe/setNodeEnter]: trace_keytrace_keytrace_key: caller_bp: %d, bp: %d, pid: %lld", trace_key.caller_bp, trace_key.bp, trace_key.pid);

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

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

	struct event event_current = {};
	event_current.type = CodeTypeJava,
	event_current.caller_bp = e->caller_bp,
	event_current.bp = e->bp,
	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(&jvmwww_trace_stack_entry, &trace_key, &event_current, BPF_ANY);

	struct event *event_p2 = bpf_map_lookup_elem(&jvmwww_trace_stack_entry, &trace_key);

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

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

SEC("uprobe/setNodeReturn")
int setNodeReturn(struct pt_regs *ctx)
{
    cw_bpf_debug("[Go] [uprobe/setNodeReturn] 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(&jvmevent_stack, &key);
	if (!e)
	{
		cw_bpf_debug("[Go] [uprobe/setNodeReturn]: conot get event");
		return 0;
	}
	__builtin_memset(e, 0, sizeof(*e));

	// ------- 获取 class_name 和 methed_name  Start
	__u64 class_name_len;
    bpf_probe_read_user(&class_name_len, sizeof(class_name_len), (void *)(PT_REGS_SP(ctx) + 8));  
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->class_name_len: %lld", class_name_len);

	void *rdx_ptr = (void *)PT_REGS_PARM3(ctx);

	void* class_name_ptr;
    if (bpf_probe_read(&class_name_ptr, sizeof(class_name_ptr), rdx_ptr)) {
        // return 1;
    }
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->class_name_ptr: %llx", class_name_ptr);

	if (class_name_len >= 100) {
		class_name_len = 99;
	}
	if (class_name_len <= 0) {
		class_name_len = 0;
	}

    bpf_probe_read(e->class_name, sizeof(e->class_name), class_name_ptr + 16); 
	// e->class_name[class_name_len] = '\0';

	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->class_name: %s", e->class_name);

	__u64 methed_name_len;
	bpf_probe_read_user(&methed_name_len, sizeof(methed_name_len), (void *)(PT_REGS_SP(ctx) + 16)); 
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->methed_name_len: %lld", methed_name_len);

	void *rcx_ptr = (void *)PT_REGS_PARM4(ctx);

	void* methed_name_ptr;
    if (bpf_probe_read(&methed_name_ptr, sizeof(methed_name_ptr), rcx_ptr)) {
        // return 1;
    }
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->methed_name_ptr: %llx", methed_name_ptr);

	if (methed_name_len >= 100) {
		methed_name_len = 99;
	}
	if (methed_name_len <= 0) {
		methed_name_len = 0;
	}
    bpf_probe_read(e->methed_name, sizeof(e->methed_name), methed_name_ptr + 16); 
	// e->methed_name[methed_name_len] = '\0';

	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->methed_name: %s", e->methed_name);

	// ------- 获取 class_name 和 methed_name End

	e->goid = get_current_goroutine();
	e->pid = pid;
	e->type = CodeTypeJava;
	e->trace_id = trace_id;
	e->location = ENTPOINT; // 标识入口还是出口
	e->caller_bp = (__u64)PT_REGS_PARM5(ctx);
	e->bp = (__u64)PT_REGS_PARM6(ctx);
	e->time_ns_end = bpf_ktime_get_ns();

	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->caller_bp: %llu", e->caller_bp);
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->bp: %llu", e->bp);
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: e->ip: %d", e->ip);
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: yes");

	struct jvm_trace_stack_entry_key_t trace_key = {};
	trace_key.caller_bp = e->caller_bp;
	trace_key.bp = e->bp;
	trace_key.pid = pid;
	trace_key.trace_id = trace_id;

	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: trace_keytrace_keytrace_key: trace_id: %llu", trace_key.trace_id);
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: trace_keytrace_keytrace_key: caller_bp: %llu, bp: %llu, pid: %lld", trace_key.caller_bp, trace_key.bp, trace_key.pid);
	
	struct event *event_p = bpf_map_lookup_elem(&jvmwww_trace_stack_entry, &trace_key);

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

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

	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: 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/setNodeReturn]: ent:event: bp:%d,caller_bp:%d,caller_ip:%llu\n", event_p->caller_bp, event_p->bp, event_p->trace_id);

	e->time_ns_start = event_p->time_ns_start;
	
	cw_bpf_debug("[Go] [uprobe/setNodeReturn]: 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(&jvmwww_trace_stack_entry, &trace_key);
	cw_bpf_debug("[Go] [uprobe/setNodeReturn] end");
	return 1;
}
