//
// BuntDB uprobe instrumentation
// Tracks individual operations: (*Tx).Set(), (*Tx).Get(), (*Tx).Delete()
//

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

#define MAX_KEY_LEN 256       // max key length
#define MAX_VALUE_LEN 512     // max value length

#define MAX_CONCURRENT 50

// ============================================================================
// Per-operation spans (tx.Set/Get/Delete each creates its own span)
// ============================================================================

// Per-operation context (one span per tx.Set/Get/Delete)
struct buntdb_op_context_t {
    BASE_SPAN_PROPERTIES
    char key[MAX_KEY_LEN];
    char value[MAX_VALUE_LEN];
    __u64 key_size;
    __u64 value_size;
    __u32 op_type;  // 0=Set, 1=Get, 2=Delete
};

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

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

// This instrumentation attaches uprobe to the following function:
// func (tx *Tx) Set(key, value string, opts *SetOptions) (previousValue string, replaced bool, err error)
// Params: 1=tx (receiver), 2=key_ptr, 3=key_len, 4=value_ptr, 5=value_len, 6=opts
// Returns: 1=previousValue_ptr, 2=previousValue_len, 3=replaced, 4=err_type_ptr, 5=err_data_ptr
SEC("uprobe/Tx_Set")
int uprobe_Tx_Set(struct pt_regs *ctx) {
    __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;
    }

    // Get consistent key (use goroutine address as key)
    void *key = (void *)GOROUTINE(ctx);
    
    // Skip if the operation already exists
    struct buntdb_op_context_t *existing = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (existing != NULL) {
        return 0;
    }

    // Allocate storage
    u32 zero = 0;
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_storage_map, &zero);
    if (op_ctx == NULL) {
        return -1;
    }

    __builtin_memset(op_ctx, 0, sizeof(struct buntdb_op_context_t));
    op_ctx->start_time = bpf_ktime_get_ns();
    op_ctx->op_type = 0;

    // Read key params: 1=tx receiver, 2=key pointer, 3=key length
    // In Go calling convention, string is split into pointer and length
    void *key_str_ptr = get_argument(ctx, 2);
    __u64 key_len = (__u64)get_argument(ctx, 3);
    
    // Write operation prefix "set "
    int off = 4;
    char opr[4] = "set ";
    bpf_probe_read(op_ctx->key, 4, opr);
    if (key_str_ptr != 0 && key_len > 0) {
        __u64 copy_key_len = key_len < (MAX_KEY_LEN - off) ? key_len : (MAX_KEY_LEN - off);
        long res = bpf_probe_read(op_ctx->key + off, copy_key_len, key_str_ptr);
        if (res == 0) {
            op_ctx->key_size = copy_key_len + off;
        }
    }
    
    // Read value params: 4=value pointer, 5=value length
    void *value_str_ptr = get_argument(ctx, 4);
    __u64 value_len = (__u64)get_argument(ctx, 5);
    
    // Read string content directly (use bpf_probe_read)
    if (value_str_ptr != 0 && value_len > 0) {
        __u64 copy_value_len = value_len < MAX_VALUE_LEN ? value_len : MAX_VALUE_LEN;
        long res = bpf_probe_read(op_ctx->value, copy_value_len, value_str_ptr);
        if (res == 0) {
            op_ctx->value_size = copy_value_len;
        }
    }
    
    // Store in map
    bpf_map_update_elem(&buntdb_op_events, &key, op_ctx, 0);
    
    // Log SQL-like message: Set key=xxx
    if (op_ctx->key_size > 0 && op_ctx->key_size < MAX_KEY_LEN) {
        cw_bpf_debug("[buntdb] Set: %.*s", (int)op_ctx->key_size, op_ctx->key);
    } else {
        cw_bpf_debug("[buntdb] Set: key=unknown");
    }
    
    return 0;
}

// This instrumentation attaches uprobe to the return of Tx.Set
SEC("uprobe/Tx_Set")
int uprobe_Tx_Set_Returns(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    __u32 pid = pid_tgid;
    
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        return 0;
    }

    // Get consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // Lookup operation context
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (op_ctx == NULL) {
        return 0;
    }

    // Update end time
    op_ctx->end_time = bpf_ktime_get_ns();

    // Read error info
    // Returns: 1=previousValue_ptr, 2=previousValue_len, 3=replaced, 4=err_type_ptr, 5=err_data_ptr
    // Go error interface is (type_ptr, data_ptr); non-nil data_ptr means error
    void *error_data_ptr = get_argument(ctx, 5);  // error 接口的 data 指针
    __u32 has_error = 0;
    if (error_data_ptr != 0) {
        has_error = 1;
    }

    __u32 zero = 0;
    struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
    if (!e) {
        bpf_map_delete_elem(&buntdb_op_events, &key);
        return 0;
    }

    // Set event attributes
    e->protocol = PROTOCOL_BUNTDB;
    e->pid = tgid;
    e->start_at = op_ctx->start_time;
    e->end_at = op_ctx->end_time;
    e->duration = e->end_at - e->start_at;
    
    // Use op_ctx->key as payload
    __u64 payload_size = op_ctx->key_size > 0 ? op_ctx->key_size : 0;
    if (payload_size > 1024) {
        payload_size = 1024;
    }
    e->payload_size = payload_size;
    
    if (payload_size > 0) {
        COPY_PAYLOAD(e->payload, payload_size, op_ctx->key);
        // Log SQL-like message
        cw_bpf_debug("[buntdb] SQL: Set: %.*s", (int)payload_size, op_ctx->key);
    }

    // Get trace info
    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 error status
    if (has_error) {
        e->status = STATUS_FAILED;
    } else {
        e->status = STATUS_OK;
    }

    // Emit event (one span per Set)
    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("[buntdb] Set Returns: success, trace_id=%llu, key=%.*s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
    }
    
    // Cleanup
    bpf_map_delete_elem(&buntdb_op_events, &key);
    return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (tx *Tx) Get(key string, ignoreExpired ...bool) (val string, err error)
// Params: 1=tx (receiver), 2=key_ptr, 3=key_len
// Note: ignoreExpired is variadic (...bool) and may be omitted
// Returns: 1=val_ptr, 2=val_len, 3=err_type_ptr, 4=err_data_ptr
SEC("uprobe/Tx_Get")
int uprobe_Tx_Get(struct pt_regs *ctx) {
	cw_bpf_debug("get");
    __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;
    }

    // Get consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // Skip if the operation already exists
    struct buntdb_op_context_t *existing = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (existing != NULL) {
        return 0;
    }

    // Allocate storage
    u32 zero = 0;
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_storage_map, &zero);
    if (op_ctx == NULL) {
        return -1;
    }

    __builtin_memset(op_ctx, 0, sizeof(struct buntdb_op_context_t));
    op_ctx->start_time = bpf_ktime_get_ns();
    op_ctx->op_type = 1;  // Get

    // Read key params: 1=tx receiver, 2=key pointer, 3=key length
    // In Go calling convention, string is split into pointer and length
    void *key_str_ptr = get_argument(ctx, 2);
    __u64 key_len = (__u64)get_argument(ctx, 3);
    
    // Read string content directly (use bpf_probe_read)
	int off = 4;
	char opr[4] = "get ";
	bpf_probe_read(op_ctx->key, 4, opr);
	if (key_str_ptr != 0 && key_len > 0) {
		__u64 copy_key_len = key_len < MAX_KEY_LEN ? key_len : MAX_KEY_LEN;
		long res = bpf_probe_read(op_ctx->key + off, copy_key_len, key_str_ptr);
		if (res == 0) {
			op_ctx->key_size = copy_key_len + off;
		}
	}


    // Store in map
    bpf_map_update_elem(&buntdb_op_events, &key, op_ctx, 0);
    
    // Log SQL-like message: Get key=xxx
    if (op_ctx->key_size > 0 && op_ctx->key_size < MAX_KEY_LEN) {
        cw_bpf_debug("[buntdb] Get: %d,key=%s", (int)op_ctx->key_size, op_ctx->key);
    } else {
        cw_bpf_debug("[buntdb] Get: key=unknown");
    }
    return 0;
}

// This instrumentation attaches uprobe to the return of Tx.Get
SEC("uprobe/Tx_Get")
int uprobe_Tx_Get_Returns(struct pt_regs *ctx) {
    __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;
    }

    // Get consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // Lookup operation context
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (op_ctx == NULL) {
        return 0;
    }

    // Update end time
    op_ctx->end_time = bpf_ktime_get_ns();

    // Read error info
    // Returns: 1=val_ptr, 2=val_len, 3=err_type_ptr, 4=err_data_ptr
    // Go error interface is (type_ptr, data_ptr); non-nil data_ptr means error
    void *error_data_ptr = get_argument(ctx, 4);  // error 接口的 data 指针
    __u32 has_error = 0;
    if (error_data_ptr != 0) {
        has_error = 1;
    }

    __u32 zero = 0;
    struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
    if (!e) {
        bpf_map_delete_elem(&buntdb_op_events, &key);
        return 0;
    }

    // Set event attributes
    e->protocol = PROTOCOL_BUNTDB;
    e->pid = tgid;
    e->start_at = op_ctx->start_time;
    e->end_at = op_ctx->end_time;
    e->duration = e->end_at - e->start_at;
    
    // Use op_ctx->key as payload
    __u64 payload_size = op_ctx->key_size > 0 ? op_ctx->key_size : 0;
    if (payload_size > 1024) {
        payload_size = 1024;
    }
    e->payload_size = payload_size;
    
    if (payload_size > 0) {
        COPY_PAYLOAD(e->payload, payload_size, op_ctx->key);
        // Log SQL-like message
        cw_bpf_debug("[buntdb] SQL: Get: %.*s", (int)payload_size, op_ctx->key);
    }

    // Get trace info
    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 error status
    if (has_error) {
        e->status = STATUS_FAILED;
    } else {
        e->status = STATUS_OK;
    }

    // Emit event (one span per Get)
    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("[buntdb] Get Returns: success, trace_id=%llu, key=%.*s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
    }
    
    // Cleanup
    bpf_map_delete_elem(&buntdb_op_events, &key);
    return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (tx *Tx) Delete(key string) (val string, err error)
// Params: 1=tx (receiver), 2=key_ptr, 3=key_len
// Returns: 1=val_ptr, 2=val_len, 3=err_type_ptr, 4=err_data_ptr
SEC("uprobe/Tx_Delete")
int uprobe_Tx_Delete(struct pt_regs *ctx) {
    __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;
    }

    // Get consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // Skip if the operation already exists
    struct buntdb_op_context_t *existing = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (existing != NULL) {
        return 0;
    }

    // Allocate storage
    u32 zero = 0;
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_storage_map, &zero);
    if (op_ctx == NULL) {
        return -1;
    }

    __builtin_memset(op_ctx, 0, sizeof(struct buntdb_op_context_t));
    op_ctx->start_time = bpf_ktime_get_ns();
    op_ctx->op_type = 2;  // Delete

    // Read key params: 1=tx receiver, 2=key pointer, 3=key length
    // In Go calling convention, string is split into pointer and length
    void *key_str_ptr = get_argument(ctx, 2);
    __u64 key_len = (__u64)get_argument(ctx, 3);
    
    // Write operation prefix "delete "
    int off = 7;
    char opr[7] = "delete ";
    bpf_probe_read(op_ctx->key, 7, opr);
    if (key_str_ptr != 0 && key_len > 0) {
        __u64 copy_key_len = key_len < (MAX_KEY_LEN - off) ? key_len : (MAX_KEY_LEN - off);
        long res = bpf_probe_read(op_ctx->key + off, copy_key_len, key_str_ptr);
        if (res == 0) {
            op_ctx->key_size = copy_key_len + off;
        }
    }
    
    // Store in map
    bpf_map_update_elem(&buntdb_op_events, &key, op_ctx, 0);
    
    // Log SQL-like message: Delete key=xxx
    if (op_ctx->key_size > 0 && op_ctx->key_size < MAX_KEY_LEN) {
        cw_bpf_debug("[buntdb] Delete: %.*s", (int)op_ctx->key_size, op_ctx->key);
    } else {
        cw_bpf_debug("[buntdb] Delete: key=unknown");
    }
    
    return 0;
}

// This instrumentation attaches uprobe to the return of Tx.Delete
SEC("uprobe/Tx_Delete")
int uprobe_Tx_Delete_Returns(struct pt_regs *ctx) {
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 tgid = pid_tgid >> 32;
    __u32 pid = pid_tgid;
    
    struct ebpf_proc_info *info = bpf_map_lookup_elem(&proc_info_map, &tgid);
    if (!info) {
        return 0;
    }

    // Get consistent key
    void *key = (void *)GOROUTINE(ctx);
    
    // Lookup operation context
    struct buntdb_op_context_t *op_ctx = bpf_map_lookup_elem(&buntdb_op_events, &key);
    if (op_ctx == NULL) {
        return 0;
    }

    // Update end time
    op_ctx->end_time = bpf_ktime_get_ns();

    // Read error info
    // Returns: 1=val_ptr, 2=val_len, 3=err_type_ptr, 4=err_data_ptr
    // Go error interface is (type_ptr, data_ptr); non-nil data_ptr means error
    void *error_data_ptr = get_argument(ctx, 4);  // error 接口的 data 指针
    __u32 has_error = 0;
    if (error_data_ptr != 0) {
        has_error = 1;
    }

    __u32 zero = 0;
    struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
    if (!e) {
        bpf_map_delete_elem(&buntdb_op_events, &key);
        return 0;
    }

    // Set event attributes
    e->protocol = PROTOCOL_BUNTDB;
    e->pid = tgid;
    e->start_at = op_ctx->start_time;
    e->end_at = op_ctx->end_time;
    e->duration = e->end_at - e->start_at;
    
    // Use op_ctx->key as payload
    __u64 payload_size = op_ctx->key_size > 0 ? op_ctx->key_size : 0;
    if (payload_size > 1024) {
        payload_size = 1024;
    }
    e->payload_size = payload_size;
    
    if (payload_size > 0) {
        COPY_PAYLOAD(e->payload, payload_size, op_ctx->key);
        // Log SQL-like message
        cw_bpf_debug("[buntdb] SQL: Delete: %.*s", (int)payload_size, op_ctx->key);
    }

    // Get trace info
    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 error status
    if (has_error) {
        e->status = STATUS_FAILED;
    } else {
        e->status = STATUS_OK;
    }

    // Emit event (one span per Delete)
    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("[buntdb] Delete Returns: success, trace_id=%llu, key=%.*s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
    }
    
    // Cleanup
    bpf_map_delete_elem(&buntdb_op_events, &key);
    return 0;
}
