#define GOROUTINE(x) ((x)->r14)
#define MAX_SYSTEM_THREADS 40960
#define HASH_ENTRIES_MAX 40960

struct go_key {
    __u32 tgid;
    __u64 goid;
} __attribute__((packed));

struct bpf_map_def SEC("maps") go_ancerstor_map = {
        .type = BPF_MAP_TYPE_LRU_HASH,
        .key_size = sizeof(struct go_key),
        .value_size = sizeof(__u64),
        .max_entries = HASH_ENTRIES_MAX,
};

struct bpf_map_def SEC("maps") go_rw_ts_map = {
        .type = BPF_MAP_TYPE_LRU_HASH,
        .key_size = sizeof(struct go_key),
        .value_size = sizeof(__u64),
        .max_entries = HASH_ENTRIES_MAX,
};

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

// Pass data between coroutine entry and exit functions
struct go_newproc_caller {
    __u64 goid;
    void *sp; // stack pointer
} __attribute__((packed));

struct bpf_map_def SEC("maps") pid_tgid_callerid_map = {
        .type = BPF_MAP_TYPE_HASH,
        .key_size = sizeof(__u64),
        .value_size = sizeof(struct go_newproc_caller),
        .max_entries = MAX_SYSTEM_THREADS,
};


static __inline __u64 get_current_goroutine(void)
{
    __u64 current_thread = bpf_get_current_pid_tgid();
    __u64 *goid_ptr = bpf_map_lookup_elem(&goroutines_map, &current_thread);
    if (goid_ptr) {
        return *goid_ptr;
    }

    return 0;
}

static __inline int is_final_ancestor(__u32 tgid, __u64 goid, __u64 now,
                                      __u64 timeout)
{
    struct go_key key = { .tgid = tgid, .goid = goid };

    __u64 *ts = bpf_map_lookup_elem(&go_rw_ts_map, &key);
    if (!ts) {
        return 0;
    }

    return now < *ts + timeout;
}


static __inline __u64 get_rw_goid(__u64 timeout, int is_socket_io)
{
    __u32 tgid = (__u32)(bpf_get_current_pid_tgid() >> 32);
    __u64 ts = bpf_ktime_get_ns();
    __u64 goid = get_current_goroutine();
    if (goid == 0) {
        return 0;
    }

    __u64 ancestor = goid;

    int idx = 0;
#pragma unroll
    for (idx = 0; idx < 6; ++idx) {
        if (is_final_ancestor(tgid, ancestor, ts, timeout)) {
            return ancestor;
        }
        struct go_key key = { .tgid = tgid, .goid = ancestor };
        __u64 *newancestor =
                bpf_map_lookup_elem(&go_ancerstor_map, &key);
        if (!newancestor) {
            break;
        }
        ancestor = *newancestor;
    }

//    if (!is_socket_io) {
//        return 0;
//    }

    struct go_key key = { .tgid = tgid, .goid = goid };
    bpf_map_update_elem(&go_rw_ts_map, &key, &ts, BPF_ANY);
    return goid;
}



SEC("uprobe/runtime.execute")
int runtime_execute(struct pt_regs *ctx)// ok
{
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    __u32 tid = (__u32) pid_tgid;

//    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
//    if (!info) {
//        return 0;
//    }
//    int offset_g_goid = info->offsets[OFFSET_IDX_GOID_RUNTIME_G];
//    if (offset_g_goid < 0) {
//        return 0;
//    }

//    void *g_ptr;

//    g_ptr = (void *) PT_GO_REGS_PARM1(ctx);

    int offset_g_goid = 152;

//    __s64 goid = 0;
//    bpf_probe_read(&goid, sizeof(goid), g_ptr + offset_g_goid);

    __s64 goroutine_id = GOROUTINE(ctx);

    bpf_map_update_elem(&goroutines_map, &pid_tgid, &goroutine_id, BPF_ANY);
    // bpf_printk("[GO] [runtime.execute] thread_id:%d|goid:%d", tid, goroutine_id);

    return 0;
}

SEC("uprobe/enter_runtime.newproc1")
int enter_runtime_newproc1(struct pt_regs *ctx)
{
    __s64 goroutine_id = GOROUTINE(ctx);
//    __u64 goroutine_id2 = GOROUTINE(ctx);


    __u64 pid_tgid = bpf_get_current_pid_tgid();
    bpf_printk("[GO] [runtime.newproc1] [enter] pid_tgid:%llu | goid:%d", pid_tgid, goroutine_id);
//    bpf_printk("[GO] [runtime.newproc1] [enter] pid_tgid:%llu | goid:%d", pid_tgid, goroutine_id2);
    __u32 tgid = pid_tgid >> 32;

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

    // go less than 1.15 cannot get parent-child coroutine relationship
    // ~ go1.14: func newproc1(fn *funcval, argp unsafe.Pointer, narg int32, callergp *g, callerpc uintptr)
//    if (info->version < GO_VERSION(1, 15, 0)) {
//        return 0;
//    }

    int offset_g_goid = 152;
//    if (offset_g_goid < 0)
//    {
//        return 0;
//    }

//    void *g_ptr = (void *) PT_GO_REGS_PARM2(ctx);;
//    if (is_register_based_call(info)) {
//        // https://github.com/golang/go/commit/8e5304f7298a0eef48e4796017c51b4d9aeb52b5
//        if (info->version >= GO_VERSION(1, 18, 0)) {
//            g_ptr = (void *)PT_GO_REGS_PARM2(ctx);
//        } else {
//            g_ptr = (void *)PT_GO_REGS_PARM4(ctx);
//        }
//    } else {
//        if (info->version >= GO_VERSION(1, 18, 0)) {
//            bpf_probe_read(&g_ptr, sizeof(g_ptr),
//                           (void *)(PT_REGS_SP(ctx) + 16));
//        } else {
//            bpf_probe_read(&g_ptr, sizeof(g_ptr),
//                           (void *)(PT_REGS_SP(ctx) + 32));
//        }
//    }

//    __s64 goid = 0;
//    bpf_probe_read(&goid, sizeof(goid), g_ptr + offset_g_goid);
//    bpf_printk("[GO] [runtime.newproc1] [enter] pid_tgid:%llu | goid:%llu", pid_tgid, goid);

//    if (!goid)
//    {
//        bpf_printk("[GO] [runtime.newproc1] [enter] pid_tgid:%llu | goid:%llu", pid_tgid, goid);
//        return 0;
//    }

    struct go_newproc_caller caller = {
            .goid = goroutine_id,
//            .sp = (void *) PT_REGS_SP(ctx),
    };
//    if (!caller.sp) {
    bpf_map_update_elem(&pid_tgid_callerid_map, &pid_tgid, &caller, BPF_ANY);
//    }
//    bpf_printk("[GO] [runtime.newproc1] [enter] pid_tgid:%llu | goid:%llu", pid_tgid);
    return 0;
}

//
SEC("uprobe/exit_runtime.newproc1")
int exit_runtime_newproc1(struct pt_regs *ctx)
{
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    __u32 tid = (__u32) pid_tgid;

//    struct ebpf_proc_info *info =
//            bpf_map_lookup_elem(&proc_info_map, &tgid);
//    if (!info) {
//        return 0;
//    }
//
//    if(info->version < GO_VERSION(1, 15, 0)){
//        return 0;
//    }

//    int offset_g_goid = 152;
//    if (offset_g_goid < 0)
//    {
//        return 0;
//    }

    struct go_newproc_caller *caller =
            bpf_map_lookup_elem(&pid_tgid_callerid_map, &pid_tgid);
    if (!caller)
    {
        bpf_printk("[GO] [runtime.newproc1] [exit] has no caller thread_id:%d", tid);

        return 0;
    }

//    void *g_ptr;
//    g_ptr = (void *) PT_GO_REGS_PARM1(ctx);
//    if (is_register_based_call(info)) {
//    } else {
//        if (info->version >= GO_VERSION(1, 18, 0)) {
//            bpf_probe_read(&g_ptr, sizeof(g_ptr), caller->sp + 32);
//        } else {
//            bpf_probe_read(&g_ptr, sizeof(g_ptr), caller->sp + 48);
//        }
//    }

    __s64 goid = GOROUTINE(ctx);
//    bpf_probe_read(&goid, sizeof(goid), g_ptr + offset_g_goid);
//    if (!goid)
//    {
//        bpf_map_delete_elem(&pid_tgid_callerid_map, &pid_tgid);
//        return 0;
//    }

    struct go_key key = {.tgid = tgid, .goid = goid};
    __u64 callergoid = caller->goid;
    bpf_map_update_elem(&go_ancerstor_map, &key, &callergoid, BPF_ANY);

    bpf_map_delete_elem(&pid_tgid_callerid_map, &pid_tgid);
    bpf_printk("[GO] [runtime.newproc1] [exit] thread_id:%d|current_goid:%d|caller_goid:%d\n", tid, goid, callergoid);

    return 0;
}
