//
// Created by Carl.Guo on 2025/11/17.
//
// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

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

#define MAX_CONCURRENT 50
// https://github.com/apache/kafka/blob/0.10.2/core/src/main/scala/kafka/common/Topic.scala#L30C3-L30C34
#define MAX_TOPIC_SIZE 256
// No constraint on the key size, but we must have a limit for the verifier
#define MAX_KEY_SIZE 256
#define MAX_CONSUMER_GROUP_SIZE 128
#define MAX_HEADERS 20

struct kafka_consumer_request_t {
	u64 start_time;
	u64 end_time;
//	struct apm_span_context sc;
	struct apm_span_context psc;
	char topic[MAX_TOPIC_SIZE];
	char key[MAX_KEY_SIZE];
	char consumer_group[MAX_CONSUMER_GROUP_SIZE];
	s64 offset;
	s64 partition;
} __attribute__((packed));

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

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

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

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

// Storage for temporary strings in extract_span_context_from_headers to avoid stack overflow
struct kafka_extract_temp_t {
	char key[CW_HEADER_KEY_LENGTH];
	char current_key[CW_HEADER_KEY_LENGTH];
	char val[CW_HEADER_VAL_LENGTH];
};

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

// Storage for parent span context to avoid stack overflow


static __always_inline long extract_span_context_from_headers(void *message, struct apm_span_context *parent_span_context, struct ebpf_proc_info *proc_info) {

	// Read the headers slice descriptor
	void *headers = (void *)(message + proc_info->kafka_message_headers_pos);
	struct go_slice headers_slice = {0};
	bpf_probe_read(&headers_slice, sizeof(headers_slice), headers);

	char key[CW_HEADER_KEY_LENGTH] = CW_HEADER_KEY_VAL;
	char current_key[CW_HEADER_KEY_LENGTH];

	for (u64 i = 0; i < headers_slice.len; i++) {
		if (i >= MAX_HEADERS) {
			break;
		}
		// Read the header
		struct kafka_header_t header = {0};
		bpf_probe_read(&header, sizeof(header), headers_slice.array + (i * sizeof(header)));
		// Check if it is the traceparent header
		if (header.key.len == CW_HEADER_KEY_LENGTH && header.value.len == CW_HEADER_VAL_LENGTH) {
			bpf_probe_read_user(current_key, sizeof(current_key), header.key.str);
			if (bpf_memcmp(key, current_key, sizeof(key))) {
				// Found the traceparent header, extract the span context
//				char val[CW_HEADER_VAL_LENGTH];
//				bpf_probe_read(val, CW_HEADER_VAL_LENGTH, header.value.array);
//				w3c_string_to_span_context(val, parent_span_context);
				cw_string_to_span_context(header.value.array, parent_span_context);
				return 0;
			}
		}
	}
	return -1;
}

// This instrumentation attaches uprobe to the following function:
// func (r *Reader) FetchMessage(ctx context.Context) (Message, error)
// Entry probe: End the span that was started when the last call returned
SEC("uprobe/FetchMessage")
int uprobe_FetchMessage(struct pt_regs *ctx) {
	/* FetchMessage is a blocking function, hence its execution time is not a good indication for the time it took to handle the message.
	Instead, we use the entry to this function to end the span which was started when it's last call returned. (A typical consumer calls FetchMessage in a loop)
	A less confusing way of looking at it is as follows 
	1. Entry to FetchMessage
	2. internal kafka code before blocking
	3. Blocking wait for message
	4. internal kafka code after blocking
	5. Return from FetchMessage
	Steps 2-4 are executed in a separate goroutine from the one the user of the library.
	*/
	void *reader = get_argument(ctx, 1);
//	struct go_iface go_context = {0};
//	get_Go_context(ctx, 2, 0, true, &go_context);
	void *goroutine = (void *)GOROUTINE(ctx);

	// Save reader to map for use in return probe
	if (reader != NULL) {
		void *reader_ptr = reader;
		bpf_map_update_elem(&kafka_reader_to_conn, &goroutine, &reader_ptr, BPF_ANY);
	}
	
	struct kafka_consumer_request_t *kafka_request = bpf_map_lookup_elem(&kafka_consumer_events, &goroutine);
	if (kafka_request == NULL || reader == NULL) {
//		cw_bpf_debug("kafka goto end");
		// The current goroutine has no kafka request,
		// this can happen in the first time FetchMessage is called
		// Save the context for the return probe for in-process context propagation
		goto save_context;
	}

	// Get consumer group from reader config
	__u64 pid_tgid = bpf_get_current_pid_tgid();
	__u32 tgid = pid_tgid >> 32;
	struct ebpf_proc_info *proc_info = bpf_map_lookup_elem(&proc_info_map, &tgid);
	if (!proc_info) {
		return 0;
	}
	get_go_string_from_user_ptr((void *)(reader + proc_info->kafka_reader_config_pos + proc_info->kafka_reader_config_group_id_pos),
	                            kafka_request->consumer_group,
	                            sizeof(kafka_request->consumer_group));


	kafka_request->end_time = bpf_ktime_get_ns();

	// Output span event
//	__u64 pid_tgid = bpf_get_current_pid_tgid();
//	__u32 tgid = pid_tgid >> 32;

	u32 zero = 0;
	struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
	if (e) {
		struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
		struct apm_trace_info_t * start_trace_info = get_apm_trace_info_by_trace_key(trace_key);
		if (!start_trace_info) {
			return -1;
		}

		cw_copy_byte_arrays(kafka_request->psc.trace_id, e->trace_id_from, APM_TRACE_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->psc.assumed_app_id, e->called_id, APM_ASSUMED_APP_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->psc.instance_id, e->instance_id_from, APM_INSTANCE_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->psc.app_id, e->app_id_from, APM_APP_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->psc.span_id, e->span_id_from, APM_SPAN_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->psc.type_from, e->type_from, APM_TYPE_FROM_SIZE);

		__u32 event_count = cw_get_event_count(start_trace_info->trace_id);

		e->method = METHOD_CONSUME;
		SET_TRACE_END(e, PROTOCOL_KAFKA);
		e->start_at = kafka_request->start_time;
		e->end_at = kafka_request->end_time;
		e->duration = e->end_at - e->start_at;
		e->trace_id = start_trace_info->trace_id;
		e->event_count = event_count;

		// 填充 MQ 相关信息
		bpf_probe_read(&e->mq.topic, sizeof(e->mq.topic), kafka_request->topic);
		bpf_probe_read(&e->mq.key, sizeof(e->mq.key), kafka_request->key);
		cw_bpf_debug("kafka_request->key %s",kafka_request->key);
		cw_bpf_debug("kafka_request->topic %s",kafka_request->topic);

		bpf_map_delete_elem(&trace_event_count_heap, &e->trace_id);
		// 清除业务层trace信息
		clear_parent_span_context_by_trace_key(start_trace_info->trace_key);
		// 清除trace信息
		__u32 tid =  (__u32)pid_tgid;
		cw_clear_trace(tgid, tid, 0);

		bpf_perf_event_output(ctx, &l7_events, BPF_F_CURRENT_CPU, e, sizeof(*e));

	}

//	stop_tracking_span(&kafka_request->sc, &kafka_request->psc);
	bpf_map_delete_elem(&kafka_consumer_events, &goroutine);

save_context:
	// Save the context for the return probe
//	bpf_map_update_elem(&goroutine_to_go_context, &goroutine, &go_context.data, 0);
	return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (r *Reader) FetchMessage(ctx context.Context) (Message, error)
// Return probe: Start tracking the span for the message that was just fetched
SEC("uprobe/FetchMessage")
int uprobe_FetchMessage_Returns(struct pt_regs *ctx) {
	/* The FetchMessage function returns a message to the user after it read it from a channel.
	The user consuming this message will handle it after this probe,
	thus it is a good place to start track the span corresponds to this message. In addition we save the message
	in a hash map to be read by the entry probe of FetchMessage, which will end this span */
	void *goroutine = (void *)GOROUTINE(ctx);
	u32 map_id = 0;
	struct kafka_consumer_request_t *kafka_request = bpf_map_lookup_elem(&kafka_consumer_request_storage_map, &map_id);
	if (kafka_request == NULL) {
		cw_bpf_debug("uprobe/FetchMessage_Returns: kafka_request is NULL");
		return 0;
	}

	// Get proc_info to access Kafka offsets
	__u64 pid_tgid = bpf_get_current_pid_tgid();
	__u32 tgid = pid_tgid >> 32;
	struct ebpf_proc_info *proc_info = bpf_map_lookup_elem(&proc_info_map, &tgid);
	if (proc_info == NULL) {
		cw_bpf_debug("uprobe/FetchMessage_Returns: proc_info is NULL");
		return 0;
	}

	// The message returned on the stack since it returned as a struct and not a pointer
	void *message = (void *)(PT_REGS_SP(ctx) + 8);
	if (proc_info->kafka_message_value_pos != 0) {
		struct go_slice_ot value_slice = {0};
		bpf_probe_read(&value_slice, sizeof(value_slice), (void *)(message + proc_info->kafka_message_value_pos));
//		cw_bpf_debug("len %d",value_slice.len);
		if (value_slice.array == NULL) {
//			cw_bpf_debug("no val");
			return 0;
		} else {
			cw_bpf_debug("has val %llu", goroutine);
		}
	} else {
		return 0;
	}

	// Zero the request
	u32 zero_id = 0;
	struct kafka_consumer_request_t *zero_kafka_request = bpf_map_lookup_elem(&kafka_consumer_request_storage_map, &zero_id);
	if (zero_kafka_request != NULL) {
		__builtin_memcpy(kafka_request, zero_kafka_request, sizeof(struct kafka_consumer_request_t));
	}

	kafka_request->start_time = bpf_ktime_get_ns();

//	cw_bpf_debug("start %llu",message);

//	struct go_iface go_context = {0};
//	get_Go_context(ctx, 2, 0, true, &go_context);

	// Get the parent span context from the message headers
	// Use per-cpu map to avoid stack overflow
//	u32 parent_sc_storage_id = 0;
//	struct apm_span_context *parent_span_context = bpf_map_lookup_elem(&cw_parent_span_context_storage_map, &parent_sc_storage_id);
//	if (parent_span_context == NULL) {
//		cw_bpf_debug("uprobe/FetchMessage_Returns: Failed to get parent_span_context storage");
//		return 0;
//	}
	// Manual zeroing to avoid memset issues
//	parent_span_context->TraceFlags = 0;
//	__builtin_memset(parent_span_context->trace_id, 0, TRACE_ID_SIZE);
//	__builtin_memset(parent_span_context->span_id, 0, SPAN_ID_SIZE);

	// Print message Value field before extracting span context from headers
	// Value is a byte slice, similar to Key

	long extract_result = extract_span_context_from_headers(message, &kafka_request->psc, proc_info);
//	if (extract_result != 0) {
//		generate_random_bytes(kafka_request->psc.trace_id, TRACE_ID_SIZE);
//	}
	if (extract_result == 0) {
		cw_bpf_debug("find");
		for (int i = 0; i < APM_TRACE_ID_SIZE; i++) {
			cw_bpf_debug("cw_get_current_tracking_span-trace_id[%d] = %02x", i, kafka_request->psc.trace_id[i]);
		}
		// Successfully extracted parent span context from headers
//		kafka_request->psc = *parent_span_context;
	} else {
		// No parent span context in headers, try to get from Go context
		generate_random_bytes(kafka_request->psc.trace_id, TRACE_ID_SIZE);
//		struct span_context *parent_sc = get_parent_span_context(go_context.data);
//		if (parent_sc != NULL) {
//			kafka_request->psc = *parent_sc;
//		}
		cw_bpf_debug("no header");
	}

	// Collecting message attributes
	// topic
	get_go_string_from_user_ptr(
			(void *)(message + proc_info->kafka_message_topic_pos), kafka_request->topic, sizeof(kafka_request->topic));
	cw_bpf_debug("topic %s",kafka_request->topic);
	// Key is a byte slice, first read the slice
	if (proc_info->kafka_message_key_pos != 0) {
		struct go_slice_ot key_slice = {0};
		bpf_probe_read(&key_slice, sizeof(key_slice), (void *)(message + proc_info->kafka_message_key_pos));
		u64 size_to_read = key_slice.len > MAX_KEY_SIZE ? MAX_KEY_SIZE : key_slice.len;
		size_to_read &= 0xFF;
		// Then read the actual key
		if (size_to_read > 0 && key_slice.array != NULL) {
			bpf_probe_read_user(kafka_request->key, size_to_read, key_slice.array);
		}
	}

	u32 zero = 0;
	struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);

	if (e) {
		e->method = METHOD_CONSUME;

		SET_TRACE_START(e, PROTOCOL_KAFKA);
		e->start_at = kafka_request->start_time;
		e->end_at = kafka_request->end_time;
		e->duration = e->end_at - e->start_at;

		struct apm_trace_info_t trace_info = cw_save_trace_info(pid_tgid, tgid, 0);
		e->trace_id = trace_info.trace_id;

		// 填充 MQ 相关信息
		bpf_probe_read(&e->mq.topic, sizeof(e->mq.topic), kafka_request->topic);
		bpf_probe_read(&e->mq.key, sizeof(e->mq.key), kafka_request->key);

		// Send event
		long error = bpf_perf_event_output(ctx, &l7_events, BPF_F_CURRENT_CPU, e, sizeof(*e));
		if (error == 0) {
			cw_bpf_debug("send start ok");
//			cw_add_event_count(e->trace_id);
//			cw_bpf_debug("[kafka consumer] send success trace %llu", e->trace_id);
		}
	} else {
		cw_bpf_debug("no e");
	}
//	return 0;

	// Initialize child span context
//	generate_random_bytes(kafka_request->sc.trace_id, TRACE_ID_SIZE);
//	generate_random_bytes(kafka_request->sc.span_id, SPAN_ID_SIZE);
//
//	// If we have a parent span, use its trace ID
//	if (extract_result == 0 && parent_span_context != NULL) {
//		copy_byte_arrays(parent_span_context->trace_id, kafka_request->sc.trace_id, TRACE_ID_SIZE);
//	}
//
//	// Collecting message attributes
//	// topic
//	get_go_string_from_user_ptr((void *)(message + proc_info->kafka_message_topic_pos), kafka_request->topic, sizeof(kafka_request->topic));
//	// partition
//	if (proc_info->kafka_message_partition_pos != 0) {
//		bpf_probe_read(&kafka_request->partition, sizeof(kafka_request->partition), (void *)(message + proc_info->kafka_message_partition_pos));
//	}
//	// offset
//	if (proc_info->kafka_message_offset_pos != 0) {
//		bpf_probe_read(&kafka_request->offset, sizeof(kafka_request->offset), (void *)(message + proc_info->kafka_message_offset_pos));
//	}
//	// Key is a byte slice, first read the slice descriptor
//	struct go_slice_ot key_slice = {0};
//	bpf_probe_read(&key_slice, sizeof(key_slice), (void *)(message + proc_info->kafka_message_key_pos));
//	u64 size_to_read = key_slice.len > MAX_KEY_SIZE ? MAX_KEY_SIZE : key_slice.len;
//	size_to_read &= 0xFF;
//	// Then read the actual key
//	if (size_to_read > 0 && key_slice.array != NULL) {
//		bpf_probe_read_user(kafka_request->key, size_to_read, key_slice.array);
//	}
//
//	// Set app_id and instance_id from proc_info
//	// proc_info is already obtained above
//	if (proc_info) {
//		copy_byte_arrays(proc_info->app_id, kafka_request->sc.app_id, APM_APP_ID_SIZE);
//		copy_byte_arrays(proc_info->instance_id, kafka_request->sc.instance_id, APM_INSTANCE_ID_SIZE);
//	}

	bpf_map_update_elem(&kafka_consumer_events, &goroutine, kafka_request, 0);
	// We are start tracking the consumer span in the return probe,
	// hence we can't read Go's context directly from the registers as we usually do.
	// Using the goroutine address as a key to the map that contains the context.
//	void *context_data_ptr = bpf_map_lookup_elem(&goroutine_to_go_context, &goroutine);
//	if (context_data_ptr != NULL) {
//		bpf_probe_read_kernel(&context_data_ptr, sizeof(context_data_ptr), context_data_ptr);
//		start_tracking_span(context_data_ptr, &kafka_request->sc);
//		bpf_map_delete_elem(&goroutine_to_go_context, &goroutine);
//	}

	return 0;
}
