//
// Created by Carl.Guo on 2024/4/1.
//

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct apm_trace_key_t));
	__uint(value_size, sizeof(struct apm_span_context));
	__uint(max_entries, 32768);
} apm_parent_span_context_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct apm_trace_key_t));
	__uint(value_size, sizeof(struct apm_span_context));
	__uint(max_entries, 32768);
} apm_current_span_context_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct apm_trace_key_t));
	__uint(value_size, sizeof(struct apm_trace_info_t));
	__uint(max_entries, 32768);
} trace_info_heap SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct fd_trace_key_t));
	__uint(value_size, sizeof(struct apm_trace_info_t));
	__uint(max_entries, 32768);
} fd_trace_info_heap SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(struct thread_trace_key_t));
	__uint(value_size, sizeof(struct apm_trace_info_t));
	__uint(max_entries, 32768);
} thread_trace_info_heap SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(key_size, sizeof(__u64));
	__uint(value_size, sizeof(__u32));
	__uint(max_entries, 32768);
} trace_event_count_heap SEC(".maps");

static __inline __attribute__((__always_inline__))
struct apm_trace_key_t get_apm_trace_key(__u64 timeout, bool is_socket_io) {
	__u64 pid_tgid = bpf_get_current_pid_tgid();
	__u64 goid = 0;

	if (timeout) {
		goid = get_rw_goid(timeout * NS_PER_SEC, is_socket_io);
	}

	struct apm_trace_key_t key = {};

	key.tgid = (__u32) (pid_tgid >> 32);

	if (goid) {
		key.goid = goid;
	} else {
		key.pid = (__u32) pid_tgid;
	}

//    debug("key.tgid:%llu|%llu",key.tgid,key.goid);
//    debug("apm key.pid:%d",(__u32)pid_tgid);
//    debug("apm key.goid:%llu",key.goid);
//    cw_bpf_debug("[Trace ID] thread id:%d",(__u32)pid_tgid);
//	bpf_printk("[get_apm_trace_key] pid:[%d]|GOID:[%llu]", (__u32) pid_tgid, key.goid);
	return key;
}

static __inline __attribute__((__always_inline__))
__u64 get_apm_trace_id(__u32 pid, __u32 tid) {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);

	struct apm_trace_info_t *trace_info = bpf_map_lookup_elem(&trace_info_heap, &trace_key);
	if (trace_info) {
		cw_bpf_debug("info_trace_id:%llu", trace_info->trace_id);
		return trace_info->trace_id;
	}
	return 0;
}

static __inline __attribute__((__always_inline__))
struct apm_trace_info_t *get_apm_trace_info(__u32 pid, __u32 tid) {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	struct apm_trace_info_t *trace_info = bpf_map_lookup_elem(&trace_info_heap, &trace_key);
	if (trace_info) {
		return trace_info;
	}
	return 0;
}

static __inline __attribute__((__always_inline__))
struct fd_trace_key_t get_fd_trace_key(__u32 pid, __u64 fd) {
	struct fd_trace_key_t trace_key = {};
	trace_key.tgid = pid;
	trace_key.fd = fd;
	return trace_key;
}

static __inline __attribute__((__always_inline__))
struct thread_trace_key_t get_thread_trace_key(__u64 id) {
	struct thread_trace_key_t trace_key = {};
	trace_key.pid_tgid = id;
	return trace_key;
}


static __inline __attribute__((__always_inline__))
__u64 get_fd_trace_id(__u32 pid, __u32 fd) {
	struct fd_trace_key_t trace_key = get_fd_trace_key(pid, fd);
	struct apm_trace_info_t *trace_info = bpf_map_lookup_elem(&fd_trace_info_heap, &trace_key);
	if (trace_info) {
		cw_bpf_debug("fd_trace_id:%llu", trace_info->trace_id);
		return trace_info->trace_id;
	}
	return 0;
}

static __inline __attribute__((__always_inline__))
void cw_save_parent_tracking_span(struct apm_span_context *sc) {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	long err = 0;
	err = bpf_map_update_elem(&apm_parent_span_context_map, &trace_key, sc, BPF_ANY);
	if (err != 0) {
		cw_bpf_debug("Failed to update tracked_spans map: %ld", err);
		return;
	}
}

static __inline __attribute__((__always_inline__))
struct apm_span_context *cw_get_parent_tracking_span() {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	struct apm_span_context *apm_sc = {0};
	struct apm_span_context *span_contexts = bpf_map_lookup_elem(&apm_parent_span_context_map, &trace_key);
//	bpf_printk("-------");

	if (span_contexts) {
		/*for (int i = 0; i < APM_TYPE_FROM_SIZE; i++) {
			bpf_printk("type_from[%d] = %02x", i, span_contexts->type_from[i]);
		}
		for (int i = 0; i < APM_SAMPLE_SIZE; i++) {
			bpf_printk("sample[%d] = %02x", i, span_contexts->sample[i]);
		}
		for (int i = 0; i < APM_HOST_ID_SIZE; i++) {
			bpf_printk("host_id[%d] = %02x", i, span_contexts->host_id[i]);
		}
		for (int i = 0; i < APM_APP_ID_SIZE; i++) {
			bpf_printk("app_id[%d] = %02x", i, span_contexts->app_id[i]);
		}
		for (int i = 0; i < APM_INSTANCE_ID_SIZE; i++) {
			bpf_printk("instance_id[%d] = %02x", i, span_contexts->instance_id[i]);
		}
		for (int i = 0; i < APM_TRACE_ID_SIZE; i++) {
			bpf_printk("trace_id[%d] = %02x", i, span_contexts->trace_id[i]);
		}
		for (int i = 0; i < APM_ASSUMED_APP_ID_SIZE; i++) {
			bpf_printk("assumed_app_id[%d] = %02x", i, span_contexts->assumed_app_id[i]);
		}
		for (int i = 0; i < APM_SPAN_ID_SIZE; i++) {
			bpf_printk("span_id[%d] = %02x", i, span_contexts->span_id[i]);
		}*/
		apm_sc = span_contexts;
	}
//	bpf_printk("-------end");

	return apm_sc;
}

static __inline __attribute__((__always_inline__))
__u64 clear_parent_span_context() {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	bpf_map_delete_elem(&apm_parent_span_context_map, &trace_key);
	return 0;
}

static __inline __attribute__((__always_inline__))
__u64 clear_current_span_context() {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	bpf_map_delete_elem(&apm_current_span_context_map, &trace_key);
	return 0;
}

static __inline __attribute__((__always_inline__))
__u64 cw_clear_trace(__u32 tgid, __u32 pid, __u32 fd) {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	struct fd_trace_key_t fd_trace_key = get_fd_trace_key(tgid, fd);
	// 清除trace信息
	struct apm_trace_info_t *trace_info = bpf_map_lookup_elem(&trace_info_heap, &trace_key);
	if (trace_info) {
		bpf_map_delete_elem(&thread_trace_info_heap, &trace_info->thread_trace_key);
	}
	// 清除trace信息
	bpf_map_delete_elem(&trace_info_heap, &trace_key);
	bpf_map_delete_elem(&fd_trace_info_heap, &fd_trace_key);
	return 0;
}

// Client当前sc
static __inline __attribute__((__always_inline__))
void cw_save_current_tracking_span(struct apm_span_context *sc) {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	bpf_map_update_elem(&apm_current_span_context_map, &trace_key, sc, BPF_ANY);
}


static __inline __attribute__((__always_inline__))
struct apm_span_context *cw_get_current_tracking_span() {
	struct apm_trace_key_t trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	struct apm_span_context *apm_sc = {0};
	struct apm_span_context *span_contexts = bpf_map_lookup_elem(&apm_current_span_context_map, &trace_key);
//	bpf_printk("-------");

	if (span_contexts) {
		/*for (int i = 0; i < APM_TYPE_FROM_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-type_from[%d] = %02x", i, span_contexts->type_from[i]);
		}
		for (int i = 0; i < APM_SAMPLE_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-sample[%d] = %02x", i, span_contexts->sample[i]);
		}
		for (int i = 0; i < APM_HOST_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-host_id[%d] = %02x", i, span_contexts->host_id[i]);
		}
		for (int i = 0; i < APM_APP_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-app_id[%d] = %02x", i, span_contexts->app_id[i]);
		}
		for (int i = 0; i < APM_INSTANCE_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-instance_id[%d] = %02x", i, span_contexts->instance_id[i]);
		}
		for (int i = 0; i < APM_TRACE_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-trace_id[%d] = %02x", i, span_contexts->trace_id[i]);
		}
		for (int i = 0; i < APM_ASSUMED_APP_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-assumed_app_id[%d] = %02x", i, span_contexts->assumed_app_id[i]);
		}
		for (int i = 0; i < APM_SPAN_ID_SIZE; i++) {
			bpf_printk("cw_get_current_tracking_span-span_id[%d] = %02x", i, span_contexts->span_id[i]);
		}*/
		apm_sc = span_contexts;
	}
//	bpf_printk("-------end");

	return apm_sc;
}

static __inline __attribute__((__always_inline__))
__u32 cw_add_event_count(__u64 trace_id) {
	__u32 *event_count = bpf_map_lookup_elem(&trace_event_count_heap, &trace_id);
	if (event_count != NULL) {
		*event_count = *event_count + 1;
		return *event_count;
	}
	return 0;
}

static __inline __attribute__((__always_inline__))
__u32 cw_get_event_count(__u64 trace_id) {
	__u32 *event_count = bpf_map_lookup_elem(&trace_event_count_heap, &trace_id);
	if (event_count != NULL){
		return *event_count;
	}
	return 0;
}

static __inline __attribute__((__always_inline__))
struct apm_trace_info_t cw_save_trace_info(__u64 id, __u32 pid, __u64 fd) {
	// apm trace
//	struct apm_trace_key_t trace_key = {0};
//	trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);

	// fd trace
//	struct fd_trace_key_t fd_trace_key = {0};
//	fd_trace_key = get_fd_trace_key(pid, fd);

	// trace info
	struct apm_trace_info_t trace_info = {0};
	__u64 uid_base = bpf_ktime_get_ns();
	trace_info.trace_id = bpf_get_current_pid_tgid() + uid_base;
	// apm trace
	trace_info.trace_key = get_apm_trace_key(120 * NS_PER_SEC, true);
	// fd trace
	trace_info.fd_trace_key = get_fd_trace_key(pid, fd);
	// thread trace
	trace_info.thread_trace_key = get_thread_trace_key(id);

	//	入口方法缓存
	bpf_map_update_elem(&trace_info_heap, &trace_info.trace_key, &trace_info, BPF_NOEXIST);
	bpf_map_update_elem(&fd_trace_info_heap, &trace_info.fd_trace_key, &trace_info, BPF_NOEXIST);
	bpf_map_update_elem(&thread_trace_info_heap, &trace_info.thread_trace_key, &trace_info, BPF_NOEXIST);
	// 事件个数
	__u32 event_count = 0;
	bpf_map_update_elem(&trace_event_count_heap, &trace_info.trace_id, &event_count, BPF_NOEXIST);
	return trace_info;
}
