//
// 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 "go_net.h"
#include "uprobe.h"
#include "apm_trace.h"
//#include "span_output.h"
//#include "trace/start_span.h"



#define MAX_CONCURRENT 50
// https://github.com/segmentio/kafka-go/blob/main/writer.go#L118
// TODO: (this value is directly impact the map sizes as well as the verification complexity)
// limitation on map entry size: https://github.com/iovisor/bcc/issues/2519#issuecomment-534359316
// the default value is 100, but it can be changed by the user
// we must specify a limit for the verifier
#if __KERNEL_FROM >= 512
#define MAX_BATCH_SIZE 10
#else
#define MAX_BATCH_SIZE 1
#endif

// 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_BROKER_ADDR_SIZE 21  // Max broker address length (e.g., "192.168.1.1:9092")
#define MAX_HOSTNAME_SIZE 256    // Max hostname/broker address size
#define MAX_PAYLOAD_SIZE 1024    // Max payload size (must match l7.c definition)

struct message_attributes_t {
	struct apm_span_context sc;
	char topic[MAX_TOPIC_SIZE];
	char key[MAX_KEY_SIZE];
};

struct kafka_request_t {
	// common attributes to all the produced messages
	u64 start_time;
	u64 end_time;
	struct apm_span_context apm_sc;
	// attributes per message
	struct message_attributes_t msgs[MAX_BATCH_SIZE];
	char global_topic[MAX_TOPIC_SIZE];
	u64 valid_messages;
	char host[MAX_HOSTNAME_SIZE];// First broker address (e.g., "192.168.1.1:9092")
}__attribute__((packed));

struct kafka_request_t2 {
	// common attributes to all the produced messages
	u64 start_time;
	u64 end_time;
	struct apm_span_context apm_sc;
	char host[MAX_HOSTNAME_SIZE];

	// attributes per message
//	struct message_attributes_t msgs[MAX_BATCH_SIZE];
//	char global_topic[MAX_TOPIC_SIZE];
//	char broker_addr[MAX_BROKER_ADDR_SIZE];  // First broker address (e.g., "192.168.1.1:9092")
//	u64 valid_messages;

}__attribute__((packed));

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

struct kafka_header_t {
	struct go_string_ot key;
	struct go_slice_ot value;
};

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct kafka_request_t));
	__uint(max_entries, 2);
} kafka_request_storage_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct kafka_request_t2));
	__uint(max_entries, 2);
} kafka_request_storage_map2 SEC(".maps");

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

// Storage for temporary strings in build_contxet_header to avoid stack overflow
struct kafka_temp_strings_t {
	char 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_temp_strings_t));
	__uint(max_entries, 1);
} kafka_temp_strings_storage_map SEC(".maps");

// Storage for temporary Go types in broker address reading to avoid stack overflow
struct kafka_broker_temp_t {
	struct go_slice_ot brokers_slice;
	struct go_string_ot broker_str;
	struct go_iface addr_iface;  // For net.Addr interface
	struct go_slice_ot ip_slice;  // For net.TCPAddr.IP
	u8 ip_bytes[16];              // For IP bytes
};

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

// https://github.com/segmentio/kafka-go/blob/main/protocol/record.go#L48

struct tcp_addr {
	struct go_slice_ot IP;  // offset 0
	int Port;            // offset 24
	// 忽略 Zone
};


static __always_inline int build_contxet_header(struct kafka_header_t *header, struct apm_span_context *span_ctx) {
	if (header == NULL || span_ctx == NULL) {
		cw_bpf_debug("build_contxt_header: Invalid arguments");
		return -1;
	}

	// Get temporary string storage from per-cpu map to avoid stack overflow
	u32 temp_strings_id = 0;
	struct kafka_temp_strings_t *temp_strings = bpf_map_lookup_elem(&kafka_temp_strings_storage_map, &temp_strings_id);
	if (temp_strings == NULL) {
		cw_bpf_debug("build_contxt_header: Failed to get temp strings storage");
		return -1;
	}

	__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/WriteMessages: proc_info is NULL");
		return 0;
	}

	// Prepare the key string for the user - use per-cpu storage
	__builtin_memcpy(temp_strings->key, CW_HEADER_KEY_VAL, CW_HEADER_KEY_LENGTH);
	void *ptr = cw_write_target_data(temp_strings->key, CW_HEADER_KEY_LENGTH, proc_info);
	if (ptr == NULL) {
		cw_bpf_debug("build_contxt_header: Failed to write key to user");
		return -1;
	}

	// build the go string of the key
	header->key.str = ptr;
	header->key.len = CW_HEADER_KEY_LENGTH;

	// Convert span_context to apm_span_context for CW format
//	struct apm_span_context apm_sc = {0};
	// Copy trace_id and span_id from span_context
//	copy_byte_arrays(span_ctx->TraceID, apm_sc.trace_id, TRACE_ID_SIZE);
//	copy_byte_arrays(span_ctx->SpanID, apm_sc.span_id, SPAN_ID_SIZE);

	// Set type_from and sample
//	apm_sc.type_from[0] = '0';
//	apm_sc.sample[0] = '0';

	// Get host_id from trace_conf_map
//	u32 k0 = 0;
//	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
//	if (trace_conf) {
//		copy_byte_arrays(trace_conf->host_id, apm_sc.host_id, APM_HOST_ID_SIZE);
//	}

	// Get app_id and instance_id from proc_info
//	copy_byte_arrays(proc_info->app_id, apm_sc.app_id, APM_APP_ID_SIZE);

	// Set assumed_app_id (for Kafka, we can use app_id as assumed_app_id)
//	copy_byte_arrays(proc_info->app_id, apm_sc.assumed_app_id, APM_ASSUMED_APP_ID_SIZE);

	// Prepare the value string for the user - use per-cpu storage
//	generate_random_bytes(apm_sc.trace_id, APM_TRACE_ID_SIZE);
//	generate_random_bytes(apm_sc.span_id, APM_SPAN_ID_SIZE);

	span_context_to_cw_string(span_ctx, temp_strings->val);

	cw_bpf_debug("%s",temp_strings->val);
	ptr = cw_write_target_data(temp_strings->val, CW_HEADER_VAL_LENGTH, proc_info);
	if (ptr == NULL) {
		cw_bpf_debug("build_contxt_header: Failed to write value to user");
		return -1;
	}

	// build the go slice of the value
	header->value.array = ptr;
	header->value.len = CW_HEADER_VAL_LENGTH;
	header->value.cap = CW_HEADER_VAL_LENGTH;
	cw_bpf_debug("build_contxt_header success");
	return 0;
}

static __always_inline int inject_kafka_header(void *message, struct kafka_header_t *header, u64 msg_headers_pos) {
//	append_item_to_slice(header, sizeof(*header), (void *) (message + msg_headers_pos));
	cw_append_item_to_slice(header, sizeof(*header), (void *) (message + msg_headers_pos));
	return 0;
}


static __always_inline long collect_kafka_attributes(void *message, struct message_attributes_t *attrs, bool collect_topic, u64 msg_key_pos, u64 msg_topic_pos) {
	if (collect_topic) {
		// Topic might be globally set for a writer, or per message
		get_go_string_from_user_ptr((void *) (message + msg_topic_pos), attrs->topic, sizeof(attrs->topic));
	}

	// Key is a byte slice, first read the slice
	struct go_slice_ot key_slice = {0};
	bpf_probe_read(&key_slice, sizeof(key_slice), (void *) (message + msg_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
	return bpf_probe_read(attrs->key, size_to_read, key_slice.array);
}

// This instrumentation attaches uprobe to the following function:
// func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error
SEC("uprobe/WriteMessages")
int uprobe_WriteMessages(struct pt_regs *ctx) {
	// In Go, "..." is equivalent to passing a slice: https://go.dev/ref/spec#Passing_arguments_to_..._parameters
	void *writer = get_argument(ctx, 1);
	void *msgs_array = get_argument(ctx, 4);
	u64 msgs_array_len = (u64) get_argument(ctx, 5);

	void *key = (void *) GOROUTINE(ctx);

	void *kafka_request_ptr = bpf_map_lookup_elem(&kafka_events, &key);
	if (kafka_request_ptr != NULL) {
		cw_bpf_debug("uprobe/WriteMessages already tracked with the current context");
		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/WriteMessages: proc_info is NULL");
		return 0;
	}


	u32 zero_id = 0;
	struct kafka_request_t *zero_kafka_request = bpf_map_lookup_elem(&kafka_request_storage_map, &zero_id);
	if (zero_kafka_request == NULL) {
		cw_bpf_debug("uuprobe/WriteMessages: zero_kafka_request is NULL");
		return 0;
	}

	u32 actual_id = 1;
	// Zero the span we are about to build, eBPF doesn't support memset of large structs (more than 1024 bytes)
	bpf_map_update_elem(&kafka_request_storage_map, &actual_id, zero_kafka_request, BPF_ANY);
	// Get a pointer to the zeroed span
	struct kafka_request_t *kafka_request = bpf_map_lookup_elem(&kafka_request_storage_map, &actual_id);
	if (kafka_request == NULL) {
		cw_bpf_debug("uprobe/WriteMessages: Failed to get kafka_request");
		return 0;
	}

	kafka_request->start_time = bpf_ktime_get_ns();

	void *writer_addr_iface_ptr = writer + 0;
	void *na_ptr = NULL;
	bpf_probe_read(&na_ptr, sizeof(void *), get_go_interface_instance(writer_addr_iface_ptr));
	if (!na_ptr)
		return 0;

	// Zero the host
	__builtin_memset(kafka_request->host, 0, sizeof(kafka_request->host));


//	void * user_str_ptr = na_ptr + 0x10;
//	if (user_str_ptr == NULL)
//	{
//		return false;
//	}
//
//	struct go_string_ot user_str = {0};
//	long success = 0;
//	success = bpf_probe_read(&user_str, sizeof(struct go_string_ot), user_str_ptr);
//	if (success != 0 || user_str.len < 1)
//	{
//		return 0;
//	}

//	u64 size_to_read = user_str.len > 22 ? 22 : user_str.len;
//	cw_bpf_debug("%s",user_str.str);
//	return 0;
//	success = bpf_probe_read(dst, size_to_read, user_str.str);

	// 设置 appid
	copy_byte_arrays(proc_info->app_id, kafka_request->msgs[0].sc.app_id, APM_APP_ID_SIZE);

	// setting assumed_app_id
	if (!get_go_string_from_user_ptr((void *) (na_ptr + 0x10), kafka_request->host, sizeof(kafka_request->host))) {
		cw_bpf_debug("target write failed, aborting ebpf probe");
		return 0;
	}
	set_assumed_app_id_arrays(kafka_request->host, kafka_request->msgs[0].sc.assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);

//	copy_byte_arrays(kafka_request->apm_sc.assumed_app_id, kafka_request->msgs[0].sc.assumed_app_id, APM_ASSUMED_APP_ID_SIZE);

//	if (!kafka_request->broker_addr){
//		return 0;
//	}
//	unsigned char assumed_app_id[APM_TRACE_ID_SIZE];
//	__builtin_memset(assumed_app_id, 0, sizeof(assumed_app_id));
//
//	set_app_id_numeric16(kafka_request->broker_addr, assumed_app_id);
//	cw_bpf_debug("digit[%d]=%c", 0, assumed_app_id[0]);
//	return 0;

//	set_assumed_app_id_arrays2(kafka_request->broker_addr, assumed_app_id);
	/*临时map绕一下*/
//	struct kafka_request_t2 *kafka_request2 = bpf_map_lookup_elem(&kafka_request_storage_map2, &actual_id);
//	if (kafka_request2 == NULL) {
//		cw_bpf_debug("uprobe/WriteMessages: Failed to get kafka_request");
//		return 0;
//	}
//	__builtin_memset(kafka_request2, 0, sizeof(struct kafka_request_t2));
//
//
//	if (!get_go_string_from_user_ptr((void *) (na_ptr + 0x10), kafka_request2->host, sizeof(kafka_request2->host))) {
//		cw_bpf_debug("target write failed, aborting ebpf probe");
//		return 0;
//	}
//	set_assumed_app_id_arrays(kafka_request2->host, kafka_request2->apm_sc.assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);
//	copy_byte_arrays(kafka_request2->apm_sc.assumed_app_id, kafka_request->msgs[0].sc.assumed_app_id, APM_ASSUMED_APP_ID_SIZE);
	/*临时map绕一下*/



//	__builtin_memcpy(psc.assumed_app_id, psc.assumed_app_id, APM_ASSUMED_APP_ID_SIZE);



	cw_bpf_debug("Kafka BrokerAddr = %s", kafka_request->host);
	struct apm_span_context *cw_psc = cw_get_parent_tracking_span();

	if(cw_psc){
		__builtin_memcpy(kafka_request->msgs[0].sc.trace_id, cw_psc->trace_id, APM_TRACE_ID_SIZE);
	} else {
		generate_random_bytes(kafka_request->msgs[0].sc.trace_id, APM_TRACE_ID_SIZE);
	}

//	generate_random_bytes(kafka_request->msgs[0].sc.span_id, APM_SPAN_ID_SIZE);


	// Try to get a global topic from Writer
	bool global_topic = get_go_string_from_user_ptr((void *) (writer + proc_info->kafka_writer_topic_pos), kafka_request->global_topic,
	                                                sizeof(kafka_request->global_topic));

	cw_bpf_debug("global_topic %d %s",global_topic,kafka_request->global_topic);


	void *msg_ptr = msgs_array;
	// Get header from per-cpu storage instead of stack to avoid stack limit
	u32 header_storage_id = 0;
	struct kafka_header_t *header = bpf_map_lookup_elem(&kafka_header_storage_map, &header_storage_id);
	if (header == NULL) {
		cw_bpf_debug("uprobe/WriteMessages: Failed to get header storage");
		return 0;
	}
	// Zero the header
	__builtin_memset(header, 0, sizeof(*header));

	// This is hack to get the message size. This calculation is based on the following assumptions:
	// 1. "Time" is the last field in the message struct. This looks to be correct for all the versions according to
	//      https://github.com/segmentio/kafka-go/blob/v0.2.3/message.go#L24C2-L24C6
	// 2. the time.Time struct is 24 bytes. This looks to be correct for all the reasnobaly latest versions according to
	//      https://github.com/golang/go/blame/master/src/time/time.go#L135
	// In the future if more libraries will need to get structs sizes we probably want to have similar
	// mechanism to the one we have for the offsets
	u16 msg_size = proc_info->kafka_message_time_pos + 8 + 8 + 8;
	kafka_request->valid_messages = 0;
	// Iterate over the messages
	for (u8 i = 0; i < MAX_BATCH_SIZE; i++) {
		if (i >= msgs_array_len) {
			break;
		}
		// Optionally collect the topic, and always collect key
		collect_kafka_attributes(msg_ptr, &kafka_request->msgs[i], !global_topic, proc_info->kafka_message_key_pos, proc_info->kafka_message_topic_pos);
		if (global_topic) {
			__builtin_memcpy(kafka_request->msgs[i].topic, kafka_request->global_topic, MAX_TOPIC_SIZE);
		}
		// Generate span id for each message
		generate_random_bytes(kafka_request->msgs[i].sc.span_id, APM_SPAN_ID_SIZE);
		if (i > 0) {
			// Copy the trace id and trace flags from the first message. This means the sampling decision is done on the first message,
			// and all the messages in the batch will have the same trace id and trace flags.
//			kafka_request->msgs[i].sc.TraceFlags = kafka_request->msgs[0].sc.TraceFlags;
			__builtin_memcpy(kafka_request->msgs[i].sc.trace_id, kafka_request->msgs[0].sc.trace_id, APM_TRACE_ID_SIZE);
			__builtin_memcpy(kafka_request->msgs[i].sc.app_id, kafka_request->msgs[0].sc.app_id, APM_APP_ID_SIZE);
			__builtin_memcpy(kafka_request->msgs[i].sc.assumed_app_id, kafka_request->msgs[0].sc.assumed_app_id, APM_APP_ID_SIZE);
		}

		// Build the header
		if (build_contxet_header(header, &kafka_request->msgs[i].sc) != 0) {
			cw_bpf_debug("uprobe/WriteMessages: Failed to build header");
			return 0;
		}
		// Inject the header
		inject_kafka_header(msg_ptr, header, proc_info->kafka_message_headers_pos);
		// Zero the header for next iteration to avoid stale data
		__builtin_memset(header, 0, sizeof(*header));
		kafka_request->valid_messages++;
		msg_ptr = msg_ptr + msg_size;
	}

	bpf_map_update_elem(&kafka_events, &key, kafka_request, 0);
	// don't need to start tracking the span, as we don't have a context to propagate locally
	return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error
SEC("uprobe/WriteMessages")
int uprobe_WriteMessages_Returns(struct pt_regs *ctx) {
	void *key = (void *) GOROUTINE(ctx);

	struct kafka_request_t *kafka_request = bpf_map_lookup_elem(&kafka_events, &key);
	if (kafka_request == NULL) {
		cw_bpf_debug("kafka_request is null\n");
		return 0;
	}
	kafka_request->end_time = bpf_ktime_get_ns();


	u32 zero = 0;
	struct l7_event *e = bpf_map_lookup_elem(&l7_event_heap, &zero);
	if (!e) {
		return 0;
	}
	__u64 pid_tgid = bpf_get_current_pid_tgid();
	__u32 tgid = pid_tgid >> 32;

	e->protocol = PROTOCOL_KAFKA;
	e->method = METHOD_PRODUCE;
	e->trace_end = 0;
	e->trace_start = 0;
	e->start_at = kafka_request->start_time;
	e->end_at = kafka_request->end_time;
	e->duration = e->end_at - e->start_at;


	bpf_probe_read(&e->target_addr, sizeof(kafka_request->host), kafka_request->host);

	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;
	}

	for (u8 i = 0; i < MAX_BATCH_SIZE; i++) {
		if (i >= kafka_request->valid_messages) {
			break;
		}
		// Build payload from host, topic, and key
//		build_kafka_payload(e, kafka_request, i);
//		bpf_probe_read(&e->payload, 22, kafka_request->host);
//		bpf_probe_read(&e->payload , 50, kafka_request->msgs[i].topic);

//		cw_bpf_debug("tttttttttttttt %s", kafka_request->msgs[i].topic);
//		cw_bpf_debug("kkkkkkkkkk %s", kafka_request->msgs[i].key);

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

//		// 复制 topic（找到实际长度，避免复制过多）
//		u32 topic_len = 0;
//		for (u32 j = 0; j < MAX_TOPIC_SIZE && j < sizeof(e->mq_topic) - 1; j++) {
//			if (kafka_request->msgs[i].topic[j] == '\0') {
//				topic_len = j;
//				break;
//			}
//		}
//		if (topic_len == 0) {
//			topic_len = (MAX_TOPIC_SIZE < sizeof(e->mq_topic) - 1) ? MAX_TOPIC_SIZE : (sizeof(e->mq_topic) - 1);
//		}
//		if (topic_len > 0) {
//			__builtin_memcpy(e->mq_topic, kafka_request->msgs[i].topic, topic_len);
//			e->mq_topic[topic_len] = '\0';
//		}
//
//		// 复制 key（找到实际长度，避免复制过多）
//		u32 key_len = 0;
//		for (u32 j = 0; j < MAX_KEY_SIZE && j < sizeof(e->mq_key) - 1; j++) {
//			if (kafka_request->msgs[i].key[j] == '\0') {
//				key_len = j;
//				break;
//			}
//		}
//		if (key_len == 0) {
//			key_len = (MAX_KEY_SIZE < sizeof(e->mq_key) - 1) ? MAX_KEY_SIZE : (sizeof(e->mq_key) - 1);
//		}
//		__builtin_memcpy(e->mq_key, kafka_request->msgs[i].key, key_len);
//		e->mq_key[key_len] = '\0';

		cw_copy_byte_arrays(kafka_request->msgs[i].sc.assumed_app_id, e->assumed_app_id, APM_ASSUMED_APP_ID_SIZE);
		cw_copy_byte_arrays(kafka_request->msgs[i].sc.span_id, e->span_id, APM_SPAN_ID_SIZE);
		// 发送事件
		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("[kafka] send success trace %llu",e->trace_id);
		}

//		kafka_request->msgs[i].sc.span_id;
//		for (int j = 0; j < 1; ++j) {
//			cw_bpf_debug("valid_messages %02x",kafka_request->msgs[i].sc.span_id[j]);
//		}
	}



//	output_span_event(ctx, kafka_request, sizeof(*kafka_request), &kafka_request->msgs[0].sc);
	bpf_map_delete_elem(&kafka_events, &key);
	// don't need to stop tracking the span, as we don't have a context to propagate locally
	return 0;
}
