//
// LevelDB uprobe instrumentation
// Tracks individual operations: (*DB).Get(), (*DB).putRec() (Put/Delete)
// Put and Delete both call putRec; keyType distinguishes the operation
//

#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 (DB.Get/Put/Delete each creates its own span)
// ============================================================================

// Per-operation context (one span per DB.Get/Put/Delete)
struct leveldb_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=Put, 1=Get, 2=Delete
};

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

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

// This instrumentation attaches uprobe to the following function:
// func (db *DB) putRec(kt keyType, key, value []byte, wo *opt.WriteOptions) error
// Put call: db.putRec(keyTypeVal=1, key, value, wo)
// Delete call: db.putRec(keyTypeDel=0, key, nil, wo)
// Params: 1=db (receiver), 2=kt (keyType int), 3=key_ptr, 4=key_len, 5=value_ptr, 6=value_len, 7=wo
// Returns: 1=err_type_ptr, 2=err_data_ptr
SEC("uprobe/DB_putRec")
int uprobe_DB_putRec(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 leveldb_op_context_t *existing = bpf_map_lookup_elem(&leveldb_op_events, &key);
    if (existing != NULL) {
        return 0;  // 已存在，跳过
    }

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

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

    // Read kt (keyType) - param 2 determines Put vs Delete
    // keyType is an integer: keyTypeDel=0 (Delete), keyTypeVal=1 (Put)
    __u64 kt = (__u64)get_argument(ctx, 2);
    if (kt == 0) {
        op_ctx->op_type = 2;  // Delete (keyTypeDel)
    } else {
        op_ctx->op_type = 0;  // Put (keyTypeVal)
    }

    // Read key params: 3=key pointer, 4=key length
    // In Go calling convention, []byte is split into pointer and length
    void *key_ptr = get_argument(ctx, 3);
    __u64 key_len = (__u64)get_argument(ctx, 4);
    
    // Write prefix based on op type
    int off;
    if (op_ctx->op_type == 0) {
        // Put operation (keyTypeVal=1)
        off = 4;
        char opr[4] = "put ";
        bpf_probe_read(op_ctx->key, 4, opr);
    } else {
        // Delete operation (keyTypeDel=0)
        off = 7;
        char opr[7] = "delete ";
        bpf_probe_read(op_ctx->key, 7, opr);
    }
    
    if (key_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_ptr);
        if (res == 0) {
            op_ctx->key_size = copy_key_len + off;
        }
    }
    
    // Read value params: 5=value pointer, 6=value length
    // Note: value is nil on Delete
    void *value_ptr = get_argument(ctx, 5);
    __u64 value_len = (__u64)get_argument(ctx, 6);
    
    // Read []byte content directly (only Put has value)
    if (value_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_ptr);
        if (res == 0) {
            op_ctx->value_size = copy_value_len;
        }
    }
    
    // Store in map
    bpf_map_update_elem(&leveldb_op_events, &key, op_ctx, 0);
    
    // Log SQL-like message
    if (op_ctx->op_type == 0) {
        if (op_ctx->key_size > 0 && op_ctx->key_size < MAX_KEY_LEN) {
            cw_bpf_debug("[leveldb] Put: %d %s", (int)op_ctx->key_size, op_ctx->key);
        } else {
            cw_bpf_debug("[leveldb] Put: key=unknown");
        }
    } else {
        if (op_ctx->key_size > 0 && op_ctx->key_size < MAX_KEY_LEN) {
            cw_bpf_debug("[leveldb] Delete: %d %s", (int)op_ctx->key_size, op_ctx->key);
        } else {
            cw_bpf_debug("[leveldb] Delete: key=unknown");
        }
    }
    
    return 0;
}

// This instrumentation attaches uprobe to the return of DB.putRec
SEC("uprobe/DB_putRec")
int uprobe_DB_putRec_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 leveldb_op_context_t *op_ctx = bpf_map_lookup_elem(&leveldb_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=err_type_ptr, 2=err_data_ptr
    // Go error interface is (type_ptr, data_ptr); non-nil data_ptr means error
    void *error_data_ptr = get_argument(ctx, 2);  // 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(&leveldb_op_events, &key);
        return 0;
    }

    // Set event attributes
    e->protocol = PROTOCOL_LEVELDB;
    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
        if (op_ctx->op_type == 0) {
            cw_bpf_debug("[leveldb] SQL: Put: %d %s", (int)payload_size, op_ctx->key);
        } else {
            cw_bpf_debug("[leveldb] SQL: Delete: %d %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 putRec)
    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);
        if (op_ctx->op_type == 0) {
            cw_bpf_debug("[leveldb] Put Returns: success, trace_id=%llu, key=%d %s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
        } else {
            cw_bpf_debug("[leveldb] Delete Returns: success, trace_id=%llu, key=%d %s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
        }
    }
    
    // Cleanup
    bpf_map_delete_elem(&leveldb_op_events, &key);
    return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (db *DB) Get(key []byte, ro *opt.ReadOptions) (value []byte, err error)
// Params: 1=db (receiver), 2=key_ptr, 3=key_len, 4=ro
// Returns: 1=value_ptr, 2=value_len, 3=err_type_ptr, 4=err_data_ptr
SEC("uprobe/DB_Get")
int uprobe_DB_Get(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 leveldb_op_context_t *existing = bpf_map_lookup_elem(&leveldb_op_events, &key);
    if (existing != NULL) {
        return 0;
    }

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

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

    // Read key params: 2=key pointer, 3=key length
    // In Go calling convention, []byte is split into pointer and length
    // Note: string and []byte share the same in-memory layout
    void *key_ptr = get_argument(ctx, 2);
    __u64 key_len = (__u64)get_argument(ctx, 3);
    
    // Write operation prefix "get "
    int off = 4;
    char opr[4] = "get ";
    bpf_probe_read(op_ctx->key, 4, opr);
    if (key_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_ptr);
        if (res == 0) {
            op_ctx->key_size = copy_key_len + off;
        }
    }

    // Store in map
    bpf_map_update_elem(&leveldb_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("[leveldb] Get: %d %s", (int)op_ctx->key_size, op_ctx->key);
    } else {
        cw_bpf_debug("[leveldb] Get: key=unknown");
    }
    return 0;
}

// This instrumentation attaches uprobe to the return of DB.Get
SEC("uprobe/DB_Get")
int uprobe_DB_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 leveldb_op_context_t *op_ctx = bpf_map_lookup_elem(&leveldb_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=value_ptr, 2=value_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(&leveldb_op_events, &key);
        return 0;
    }

    // Set event attributes
    e->protocol = PROTOCOL_LEVELDB;
    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("[leveldb] SQL: Get: %d %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("[leveldb] Get Returns: success, trace_id=%llu, key=%d %s", e->trace_id, (int)op_ctx->key_size, op_ctx->key);
    }
    
    // Cleanup
    bpf_map_delete_elem(&leveldb_op_events, &key);
    return 0;
}


