//
// Created by Carl.Guo on 2025/11/12.
//

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

#define MAX_STMT_LEN 128  // 最大语句长度

// Query 结构体中 stmt 字段的偏移量
// 如果 stmt 是 Query 结构体的第一个字段，偏移量为 0
// 可以通过 DWARF 信息动态获取，这里使用可配置的偏移量

// internalQuery 结构体中 originalQuery 字段的偏移量
// internalQuery.originalQuery 是指向 Query 的指针

// 存储 gocql 查询请求的上下文
struct gocql_query_context_t {
    BASE_SPAN_PROPERTIES
    char stmt[MAX_STMT_LEN];
    __u64 stmt_size;
};


struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct gocql_query_context_t));
	__uint(max_entries, 1);
} gocql_query_storage_map SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __type(key, void *);
    __type(value, struct gocql_query_context_t);
    __uint(max_entries, MAX_CONCURRENT);
} gocql_query_events SEC(".maps");

#define MAX_CONCURRENT 50

#define MAX_QUERY_SIZE 256

struct sql_request_t {
	BASE_SPAN_PROPERTIES
	char query[MAX_QUERY_SIZE];
};

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__type(key, void *);
	__type(value, struct sql_request_t);
	__uint(max_entries, MAX_CONCURRENT);
} sql_events SEC(".maps");

// This instrumentation attaches uprobe to the following function:
// func (s *Session) Query(stmt string, values ...interface{}) *Query
/*
SEC("uprobe/Session_Query")
int uprobe_Session_Query(struct pt_regs *ctx) {
	cw_bpf_debug("[Session_Query] 0");

    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        return 0;
    }

    // 获取 stmt string 参数
    // 在 Go 的调用约定中，string 参数被拆分成两个独立的参数：指针和长度
    // 参数顺序：1 = s *Session (receiver), 2 = stmt 指针, 3 = stmt 长度
    void *stmt_ptr = get_argument(ctx, 2);
    __u64 stmt_len = (__u64)get_argument(ctx, 3);
    
    cw_bpf_debug("[Session_Query] stmt_ptr=%llx", (__u64)stmt_ptr);
    cw_bpf_debug("[Session_Query] stmt_len=%llu", stmt_len);
    
    if (stmt_ptr == 0 || stmt_len == 0) {
        return 0;
    }
    
    if (stmt_len > MAX_STMT_LEN) {
        stmt_len = MAX_STMT_LEN;
    }

    char stmt_buf[MAX_STMT_LEN + 1] = {0};

    // 使用固定大小读取，验证器可以确定读取大小的上界
    // sizeof(stmt_buf) - 1 是编译时常量，验证器可以验证
    long res = bpf_probe_read(stmt_buf, sizeof(stmt_buf) - 1, stmt_ptr);
    if (res != 0) {
        return 0;
    }

    // 确保字符串以 null 结尾
    stmt_buf[MAX_STMT_LEN] = '\0';
    
    // 打印 stmt
    cw_bpf_debug("[Session_Query] stmt: %s", stmt_buf);

    return 0;
}
*/

// This instrumentation attaches uprobe to the following function:
// func (s *Session) executeQuery(qry *Query) (it *Iter)
SEC("uprobe/Session_executeQuery")
int uprobe_Session_executeQuery(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    
//    cw_bpf_debug("[gocql] enter executeQuery");
//    cw_bpf_debug("[gocql] tgid=%u", tgid);
    
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        cw_bpf_debug("[gocql] no proc_info");
        return 0;
    }

    // 获取第二个参数 qry *Query
    // 参数顺序：1 = s *Session (receiver), 2 = qry *Query
    void *qry_ptr = get_argument(ctx, 2);
//    cw_bpf_debug("[gocql] qry_ptr=%llx", (__u64)qry_ptr);
    if (qry_ptr == 0) {
        cw_bpf_debug("[gocql] qry_ptr is null");
        return 0;
    }

    // 获取 consistent key (使用 goroutine 地址作为 key)
    void *key = (void *)GOROUTINE(ctx);
//    cw_bpf_debug("[gocql] key=%llx", (__u64)key);
    
    // 检查是否已经存在该请求
    struct gocql_query_context_t *existing = bpf_map_lookup_elem(&gocql_query_events, &key);
    if (existing != NULL) {
        cw_bpf_debug("[gocql] already exists");
        return 0;
    }



	// 从 Query 结构体中读取 stmt 字段
    // stmt 是一个 string 类型，在 Go 中是 go_string_ot 结构体
    struct go_string_ot stmt_str = {0};
    void *stmt_ptr = (void *)((char *)qry_ptr + 0);
//    cw_bpf_debug("[gocql] stmt_ptr=%llx", (__u64)stmt_ptr);

    // 读取 string 结构体
    long res = bpf_probe_read(&stmt_str, sizeof(stmt_str), stmt_ptr);
//    cw_bpf_debug("[gocql] read res=%ld", res);
//    cw_bpf_debug("[gocql] read ken=%ld", stmt_str.len);
//    cw_bpf_debug("[gocql] sql : %s", stmt_str.str);
//	return 0;

    if (res != 0) {
        return 0;
    }
//    cw_bpf_debug("[gocql] stmt len=%lld", stmt_str.len);

    // 分配存储空间
	u32 zero = 0;
	struct gocql_query_context_t *query_ctx = bpf_map_lookup_elem(&gocql_query_storage_map, &zero);
	if (query_ctx == NULL) {
		cw_bpf_debug("grpc:client:ClientConn_Invoke: failed to get storage");
		return -1;
	}

	__builtin_memset(query_ctx, 0, sizeof(struct gocql_query_context_t));

	query_ctx->start_time = bpf_ktime_get_ns();
	query_ctx->stmt_size = stmt_str.len;

    // 读取 stmt 字符串内容
//    __u64 stmt_size = stmt_str.len > MAX_STMT_LEN ? MAX_STMT_LEN : stmt_str.len;
//    cw_bpf_debug("[gocql] stmt_size=%llu", stmt_size);
    res = bpf_probe_read(query_ctx->stmt, sizeof(query_ctx->stmt), stmt_str.str);
//    cw_bpf_debug("[gocql] read stmt res=%ld", res);
    if (res != 0) {
        cw_bpf_debug("[gocql] read stmt failed");
        return 0;
    }

//    // 确保字符串以 null 结尾
//    if (stmt_size < MAX_STMT_LEN) {
//        query_ctx->stmt[stmt_size] = '\0';
//    } else {
//        query_ctx->stmt[MAX_STMT_LEN - 1] = '\0';
//    }

    // 获取 context
//    void *context_ptr_val = get_Go_context(ctx, 2, 0, true);
//    cw_bpf_debug("[gocql] context_ptr=%llx", (__u64)context_ptr_val);
//    if (context_ptr_val != 0) {
//        start_tracking_span(context_ptr_val, &query_ctx.sc);
//    }

    // 保存到 map
     bpf_map_update_elem(&gocql_query_events, &key, query_ctx, 0);
//    cw_bpf_debug("[gocql] save res=%ld", update_res);
//    cw_bpf_debug("[gocql] stmt_size=%llu", query_ctx->stmt_size);
//    cw_bpf_debug("[gocql] stmt=%s", query_ctx->stmt);

//    if (query_ctx.stmt_size > 0 && query_ctx.stmt_size <= 64) {
//        cw_bpf_debug("[gocql] stmt: %s", query_ctx.stmt);
//    }

    return 0;
}

// This instrumentation attaches uprobe to the return of executeQuery
SEC("uprobe/Session_executeQuery")
int uprobe_Session_executeQuery_Returns(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    __u32 pid = pid_tgid;
    
    cw_bpf_debug("[gocql] enter Returns");
//    cw_bpf_debug("[gocql] tgid=%u", tgid);
//    return 0;
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        cw_bpf_debug("[gocql] Returns no info");
        return 0;
    }

    // 获取 consistent key
    void *key = (void *)GOROUTINE(ctx);
//    cw_bpf_debug("[gocql] Returns key=%llx", (__u64)key);
    
    // 查找对应的请求
    struct gocql_query_context_t *query_ctx = bpf_map_lookup_elem(&gocql_query_events, &key);
    if (query_ctx == 0) {
        cw_bpf_debug("[gocql] Returns no ctx");
        return 0;
    }
//    cw_bpf_debug("[gocql] Returns found ctx");

    // 更新结束时间
    query_ctx->end_time = bpf_ktime_get_ns();
//    cw_bpf_debug("[gocql] Returns stmt_size=%llu", query_ctx->stmt_size);
//    cw_bpf_debug("[gocql] Returns duration=%llu", query_ctx->end_time - query_ctx->start_time);
//
//	cw_bpf_debug("[gocql] Returns stmt=%s", query_ctx->stmt);

    // 停止跟踪 span
//    stop_tracking_span(&query_ctx->sc, &query_ctx->psc);

    __u32 zero = 0;
    struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
    if (!e) {
        cw_bpf_debug("[gocql] Returns no event");
        bpf_map_delete_elem(&gocql_query_events, &key);
        return 0;
    }

    // 设置事件属性
    e->protocol = PROTOCOL_CASSANDRA;
    e->pid = tgid;
    e->start_at = query_ctx->start_time;
    e->end_at = query_ctx->end_time;
    e->duration = e->end_at - e->start_at;
    e->payload_size = query_ctx->stmt_size;
//    cw_bpf_debug("[gocql] Returns payload_size=%llu", e->payload_size);
    
    COPY_PAYLOAD(e->payload, query_ctx->stmt_size, query_ctx->stmt);

    struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
    struct apm_trace_info_t *trace_info = get_apm_trace_info_by_trace_key(trace_key);
    
    if (trace_info == 0) {
        trace_info = get_apm_trace_info_v3(trace_key, pid_tgid, tgid, pid_tgid);
    }
    
    if (trace_info) {
        e->trace_id = trace_info->trace_id;
    }
    
    if (e->trace_id == 0) {
        e->trace_id = get_apm_trace_id(tgid, pid_tgid);
    }
    // 发送事件
	SET_TRACE_METHOD(e);
    long error = bpf_perf_event_output(ctx, &l7_events, BPF_F_CURRENT_CPU, e, sizeof(*e));
    if (error == 0) {
        cw_add_event_count(e->trace_id);
        cw_bpf_debug("[gocql] Returns success");
    }
    // 清理
    bpf_map_delete_elem(&gocql_query_events, &key);
    cw_bpf_debug("[gocql] Returns done %llu",e->trace_id);
    return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (s *Session) executeQuery(qry *internalQuery) (it *Iter)
// internalQuery 包含 originalQuery *Query 字段
SEC("uprobe/Session_executeQuery_cassandra")
int uprobe_Session_executeQuery_cassandra(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    
//    bpf_printk("[cassandra] enter Session.executeQuery");
    
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        cw_bpf_debug("[cassandra] no proc_info");
        return 0;
    }

    // 获取第二个参数 qry *internalQuery
    // 参数顺序：1 = s *Session (receiver), 2 = qry *internalQuery
    void *internal_qry_ptr = get_argument(ctx, 2);
    if (internal_qry_ptr == 0) {
        cw_bpf_debug("[cassandra] internal_qry_ptr is null");
        return 0;
    }

    // 从 internalQuery 中读取 originalQuery *Query 指针
    void *original_query_ptr = 0;
    void *original_query_ptr_addr = (void *)((char *)internal_qry_ptr + 0);
    long res = bpf_probe_read(&original_query_ptr, sizeof(original_query_ptr), original_query_ptr_addr);
    if (res != 0 || original_query_ptr == 0) {
        cw_bpf_debug("[cassandra] failed to read originalQuery pointer");
        return 0;
    }

    // 获取 consistent key (使用 goroutine 地址作为 key)
    void *key = (void *)GOROUTINE(ctx);
    
    // 检查是否已经存在该请求
    struct gocql_query_context_t *existing = bpf_map_lookup_elem(&gocql_query_events, &key);
    if (existing != NULL) {
        cw_bpf_debug("[cassandra] already exists");
        return 0;
    }

    // 从 Query 结构体中读取 stmt 字段
    struct go_string_ot stmt_str = {0};
    void *stmt_ptr = (void *)((char *)original_query_ptr + 0);
    
    // 读取 string 结构体
    res = bpf_probe_read(&stmt_str, sizeof(stmt_str), stmt_ptr);
    if (res != 0) {
        cw_bpf_debug("[cassandra] failed to read stmt_str");
        return 0;
    }
    // 分配存储空间
    u32 zero = 0;
    struct gocql_query_context_t *query_ctx = bpf_map_lookup_elem(&gocql_query_storage_map, &zero);
    if (query_ctx == NULL) {
        cw_bpf_debug("[cassandra] failed to get storage");
        return -1;
    }

    __builtin_memset(query_ctx, 0, sizeof(struct gocql_query_context_t));

    query_ctx->start_time = bpf_ktime_get_ns();
    query_ctx->stmt_size = stmt_str.len;

    // 读取 stmt 字符串内容
    res = bpf_probe_read(query_ctx->stmt, sizeof(query_ctx->stmt), stmt_str.str);
    if (res != 0) {
        cw_bpf_debug("[cassandra] read stmt failed");
        return 0;
    }

    // 保存到 map
    bpf_map_update_elem(&gocql_query_events, &key, query_ctx, 0);

    return 0;
}

// This instrumentation attaches uprobe to the return of Session.executeQuery
SEC("uprobe/Session_executeQuery_cassandra")
int uprobe_Session_executeQuery_cassandra_Returns(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;

    cw_bpf_debug("[cassandra] enter Session.executeQuery Returns");
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        cw_bpf_debug("[cassandra] Returns no info");
        return 0;
    }

    // 获取 consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // 查找对应的请求
    struct gocql_query_context_t *query_ctx = bpf_map_lookup_elem(&gocql_query_events, &key);
    if (query_ctx == 0) {
        cw_bpf_debug("[cassandra] Returns no ctx");
        return 0;
    }

    // 更新结束时间
    query_ctx->end_time = bpf_ktime_get_ns();

    __u32 zero = 0;
    struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
    if (!e) {
        cw_bpf_debug("[cassandra] Returns no event");
        bpf_map_delete_elem(&gocql_query_events, &key);
        return 0;
    }

    // 设置事件属性
    e->protocol = PROTOCOL_CASSANDRA;
    e->pid = tgid;
    e->start_at = query_ctx->start_time;
    e->end_at = query_ctx->end_time;
    e->duration = e->end_at - e->start_at;
    e->payload_size = query_ctx->stmt_size;
    
    COPY_PAYLOAD(e->payload, query_ctx->stmt_size, query_ctx->stmt);

    struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
    struct apm_trace_info_t *trace_info = get_apm_trace_info_by_trace_key(trace_key);
    
    if (trace_info == 0) {
        trace_info = get_apm_trace_info_v3(trace_key, pid_tgid, tgid, pid_tgid);
    }
    
    if (trace_info) {
        e->trace_id = trace_info->trace_id;
    }

    // 发送事件
	SET_TRACE_METHOD(e);
    long error = bpf_perf_event_output(ctx, &l7_events, BPF_F_CURRENT_CPU, e, sizeof(*e));
    if (error == 0) {
        cw_add_event_count(e->trace_id);
        cw_bpf_debug("[cassandra] Returns success");
    }
    // 清理
    bpf_map_delete_elem(&gocql_query_events, &key);
    cw_bpf_debug("[cassandra] Returns done %llu", e->trace_id);
    return 0;
}
