// Copyright The CW Authors


#include "java_common.h"
static __inline int updataSocket2(struct sock_t *map_data_res, char * payload,int len,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	long res = bpf_probe_write_user((void *) jbytechar_ptr,  &map_data_res->payload, sizeof(map_data_res->payload));
	bpf_printk("sizeof(map_data_res->payload) %d\n", sizeof(map_data_res->payload));
	bpf_printk("sizeof(payload) %d\n", sizeof(payload));
	bpf_printk("&payload 0x%lx\n", &payload);

	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", jbytechar_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", jbytechar_ptr, res);
	}

	unsigned long new_val;
	new_val = len;
	res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", len_from_rbp_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
	}

	bpf_printk("len %d\n", len);

//	bpf_printk("payload %s\n", payload);
//	for (int i = 270; i < len; ++i) {
//		bpf_printk("data[%d]=%c", i,payload[i]);
//		if(payload[i]=='\0'){
//			break;
//		}
//	}
	return 0;
}

static __inline int updataSocket(struct sock_t *map_data_res,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	long res = bpf_probe_write_user((void *) jbytechar_ptr, &map_data_res->payload, sizeof(map_data_res->payload));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", jbytechar_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", jbytechar_ptr, res);
	}

	unsigned long new_val;
	new_val = map_data_res->size;
	res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res == 0) {
		bpf_printk("Successfully wrote value to user address: 0x%lx\n", len_from_rbp_ptr);
	} else {
		bpf_printk("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
	}

	bpf_printk("Successfully %d\n", map_data_res->size);

	bpf_printk("Successfully %s\n", map_data_res->payload);
//	for (int i = 270; i < 290; ++i) {
//		bpf_printk("data[%d]=%c", i,map_data_res->payload[i]);
//		if(map_data_res->payload[i]=='\0'){
//			break;
//		}
//	}
	return 0;
}

static __inline struct apm_span_context* build_sc(struct ebpf_proc_info *proc_info) {
	u32 key = 0;
	struct apm_span_context *cw_sc = bpf_map_lookup_elem(&apm_span_context_heap, &key);

	if (cw_sc == NULL){
		return NULL;
	}

	struct apm_span_context *cw_psc = cw_get_parent_tracking_span();
	if(cw_psc){
		copy_byte_arrays(cw_psc->trace_id, cw_sc->trace_id, APM_TRACE_ID_SIZE);
	}
	// new spanid
	generate_random_bytes(cw_sc->span_id, APM_SPAN_ID_SIZE);

	// set host_id/appid
	u32 k0 = 0;
	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
	if (trace_conf) {
//		for (int i = 0; i < 8; ++i) {
//			cw_bpf_debug("[Client] host_id:%02x", trace_conf->host_id[i]);
//		}
		copy_byte_arrays(trace_conf->host_id, cw_sc->host_id, APM_HOST_ID_SIZE);
		copy_byte_arrays(trace_conf->app_id, cw_sc->app_id, APM_APP_ID_SIZE);
	}

//	__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) {
		copy_byte_arrays(proc_info->instance_id, cw_sc->instance_id, APM_APP_ID_SIZE);
	}

	// set assumed_app_id
	set_assumed_app_id_arrays("111", cw_sc->assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);
	cw_save_current_tracking_span(cw_sc);
	return cw_sc;
//	char header[CW_STREAM_HEADER_LEN];
//	span_context_to_cw_string_stream(cw_span_context, header, '1');
//	bpf_printk("Successfully HEADER  %s\n", header);
//	return header;
}

static __inline struct sock_t* buildHeader(struct sock_t *map_data) {

	int key = 0;
	struct sock_t *map_data_res = bpf_map_lookup_elem(&socket_res, &key);

	if (!map_data_res) {
		bpf_printk("Failed to lookup socket_heap");
		return NULL;
	}

//	__builtin_memset(map_data_res, 0, sizeof(struct sock_t));
//	char header[5] = "cwt\r\n";

//	char header[CW_STREAM_HEADER_LEN] = "cwtrace: 00:00:1015481350055581:5450531005555981:5610250100539899:304775019cd3218a304775019cd3218a:1001025098564810:140acc88cde8773f\r\n";
	struct apm_span_context *cw_span_context = bpf_map_lookup_elem(&apm_span_context_heap, &key);

	if (cw_span_context == NULL){
		return NULL;
	}

	struct apm_span_context *cw_psc = cw_get_parent_tracking_span();
	if(cw_psc){
		copy_byte_arrays(cw_psc->trace_id, cw_span_context->trace_id, APM_TRACE_ID_SIZE);
		// new spanid
		generate_random_bytes(cw_span_context->span_id, APM_SPAN_ID_SIZE);
	}

	// set host_id/appid
	u32 k0 = 0;
	struct trace_conf_t *trace_conf = trace_conf_map__lookup(&k0);
	if (trace_conf) {
		for (int i = 0; i < 8; ++i) {
//			cw_bpf_debug("[Client] host_id:%02x", trace_conf->host_id[i]);
		}
		copy_byte_arrays(trace_conf->host_id, cw_span_context->host_id, APM_HOST_ID_SIZE);
		copy_byte_arrays(trace_conf->app_id, cw_span_context->app_id, APM_APP_ID_SIZE);
	}

	__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) {
		for (int i = 0; i < 8; ++i) {
//			cw_bpf_debug("[Client] instance_id:%02x", proc_info->instance_id[i]);
		}
		copy_byte_arrays(proc_info->instance_id, cw_span_context->instance_id, APM_APP_ID_SIZE);
	}

	// set assumed_app_id
	set_assumed_app_id_arrays("111", cw_span_context->assumed_app_id, APM_ASSUMED_APP_ID_STRING_SIZE);
	cw_save_current_tracking_span(cw_span_context);

	char header[CW_STREAM_HEADER_LEN];
	span_context_to_cw_string_stream(cw_span_context, header, '1');
	bpf_printk("Successfully HEADER  %s\n", header);

//	bpf_probe_read(map_data_res->payload, map_data->header_offset_idx, map_data->payload);
//	bpf_printk("Successfully %s\n", data);
//	return 1;
#pragma unroll
	for (int i = 0; i < MAX_LEN; i++) {
		if (i < map_data->header_offset_idx) {
			map_data_res->payload[i] = map_data->payload[i];
		} else {
			if (i == map_data->header_offset_idx) {
#pragma unroll
				for (int k = 0; k < CW_STREAM_HEADER_LEN; k++) {
					int tmp_len = i + k;
					if (tmp_len < MAX_LEN) {
						map_data_res->payload[tmp_len] = header[k];
					}
				}
			}
			int tmp_len = i + CW_STREAM_HEADER_LEN;
			if (tmp_len < MAX_LEN) {
				map_data_res->payload[tmp_len] = map_data->payload[i];
				if (map_data->payload[i] == '\0')
					break;
//				bpf_printk("map_data->payload: %c->i=%d, offindex %d", map_data_res->payload[i], i, map_data->header_offset_idx);
			}
		}
	}
	map_data_res->size = map_data->size + CW_STREAM_HEADER_LEN;

	return map_data_res;
}

static __inline int insertHeader(struct sock_t *map_data,void * jbytechar_ptr,void * len_from_rbp_ptr) {

	struct sock_t *map_data_res= buildHeader(map_data);
	if (map_data_res == NULL) {
		bpf_printk("Failed to lookup socket_heap");
		return -1;
	}

	/*修改部分*/
	return updataSocket2(map_data_res, map_data_res->payload, map_data_res->size, jbytechar_ptr, len_from_rbp_ptr);
}


SEC("uprobe/Java_java_net_SocketOutputStream_socketWrite0")
int uprobe_Java_java_net_SocketOutputStream_socketWrite0(struct pt_regs *ctx) {
//	char subs2[HEADER_LEN +1];
//	int a =0;
//	struct apm_span_context *cw_sc2 = build_sc();
////	struct apm_span_context *cw_sc2 = bpf_map_lookup_elem(&apm_span_context_heap, &a);
//
//	if (cw_sc2 == NULL){
//		return 0;
//	}
///*111*/
////	char subs[HEADER_LEN +1];
//	span_context_to_cw_string_stream(cw_sc2, subs2, '1');
//	bpf_printk("%s",subs2);
//
//	struct stream_header_t *stream_header1 = bpf_map_lookup_elem(&stream_header_heap, &a);
//	if (!stream_header1) {
//		bpf_printk("Failed to lookup socket_heap\n");
//		return 1;
//	}
//	bpf_probe_read_user(stream_header1->header, sizeof(stream_header1->header), (void *) (subs2));
//	bpf_printk("stream_header1->header------%s",stream_header1->header);
//
//	return 0;



	// 捕获第六个参数 data_count
	int data_count = PT_REGS_PARM6(ctx);
	if (data_count < MIN_LEN) {
		return 0;
	}


//	void *len_from_rsp_ptr = (void *) PT_REGS_RSP(ctx) - 65640 +8 ;
	void *len_from_rbp_ptr = (void *) PT_LEN_REGS(ctx) ;
//	void *jbytearray_from_rbp_p_p_p = (void *) (PT_REGS_RBP(ctx)) - 0x10050;

//	long res;
	void *len_ptr = 0;
	cw_bpf_debug("address: len_from_rbp_ptr<0x%lx>\n", len_from_rbp_ptr);
//	bpf_printk("len_from_rsp_ptr address: 0x%lx\n", len_from_rsp_ptr);

	bpf_probe_read(&len_ptr, sizeof(len_ptr), (void *) len_from_rbp_ptr);
	cw_bpf_debug("[len_ptr] before addr<0x%lx>, %d \n", len_from_rbp_ptr, len_ptr);

//	bpf_probe_read(&len_ptr, sizeof(len_ptr), (void *) len_from_rbp_ptr);
//	cw_bpf_debug("[len_ptr]after  addr<0x%lx>, %d \n", len_from_rbp_ptr, len_ptr);

	if ( (long) len_ptr != data_count ){
		cw_bpf_debug("[len_ptr] check error.");
		return 0;
	}

	// 获取jbytearray
	__u32 tgid = (__u32)(bpf_get_current_pid_tgid() >> 32);
	struct ebpf_proc_info *proc_info = bpf_map_lookup_elem(&proc_info_map, &tgid);
	if (!proc_info) {
		return 0;
	}

	cw_bpf_debug("[Java] code_type:%d\n", proc_info->code_type);
	void *jbytearray_ptr = NULL;
	// 1003 是 javaAOT 的
	if (proc_info->code_type == CodeTypeJavaAot) {
		jbytearray_ptr = (void *) (PT_REGS_SP(ctx)) + 0x8 + 0x10060+0x20;
	} else{
		jbytearray_ptr = (void *) PT_REGS_PARM4(ctx);
	}
	
	cw_bpf_debug("jbytechar_ptr_from_rcx <0x%lx>", jbytearray_ptr);

	unsigned long jbytechar_head_ptr;
	// 读取一级指针
	long ret = bpf_probe_read_user(&jbytechar_head_ptr, sizeof(unsigned long), (void *) jbytearray_ptr);
	cw_bpf_debug("[jbytechar_head_ptr] <0x%lx>", jbytechar_head_ptr);

	if (ret != 0) {
		cw_bpf_debug("Failed to read first level ptr: %d\n", ret);
		return 0;
	}

	// 检查一级指针是否有效
	if (!jbytechar_head_ptr) {
		cw_bpf_debug("First level pointer is null.\n");
		return 0;
	}

	// 定义 Map 键值和 Map 数据结构
	int key = 0;
	struct sock_t *map_data = bpf_map_lookup_elem(&socket_heap, &key);
	if (!map_data) {
		cw_bpf_debug("Failed to lookup socket_heap\n");
		return 1;
	}
//	__builtin_memset(map_data, 0, sizeof(struct sock_t));
	// 读取用户空间数据到 map_data->payload
	void *jbytechar_ptr = (void *) (jbytechar_head_ptr + 16);
	cw_bpf_debug("[jbytechar_ptr] <0x%lx>", jbytechar_ptr);
	long err = bpf_probe_read_user_str(map_data->payload, sizeof(map_data->payload), jbytechar_ptr);

	if (err < 0) {
		cw_bpf_debug("bpf_probe_read_user failed with return code: %d\n", err);
		return 0;
	}
	map_data->size = data_count;
	map_data->payload_char_p = (void *) jbytechar_ptr;
	map_data->payload_len_p = (void *) len_from_rbp_ptr;

	// http协议过滤
	if (!is_http_request2(map_data->payload, map_data->size)) {
		return -1;
	}

	// 查找 Header 开始位置
#pragma clang loop unroll(full)
	for (int i = 0; i < 200 - 5; i++) {
		if (map_data->payload[i] == '1' &&
		    map_data->payload[i + 1] == '.' &&
		    map_data->payload[i + 2] == '1' &&
		    map_data->payload[i + 3] == '\r' &&
		    map_data->payload[i + 4] == '\n') {
			map_data->header_offset_idx = i + 5;
			break;
		}
	}

	if (map_data->header_offset_idx == 0) {
		return -1;
	}


/*设置初始值*/
	char *full_payload= bpf_map_lookup_elem(&large_array_map, &key);
	if (!full_payload)
		return -1;
	// 读取完整http
	bpf_probe_read_user_str(full_payload, JAVA_MAX_BUFFER_SIZE, jbytechar_ptr);
	cw_bpf_debug("origin full_payload %s\n", full_payload);


	// 尾部payload
	u32 key2 = 1;
	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &key2);
	if (!tail_payload) {
		cw_bpf_debug("header ->end %s\n", tail_payload);
		return -1;
	}

	u32 copy_size = L7_IOVEC_BUF_SIZE;
	long end_str_len = bpf_probe_read_user_str(tail_payload, JAVA_MAX_BUFFER_SIZE,
	                                           (void *) (jbytechar_ptr + map_data->header_offset_idx));
	cw_bpf_debug("end_str_len %d\n", end_str_len);
	map_data->end_str_len = end_str_len;
	if (end_str_len > copy_size) {
		u32 chunk = (u32) end_str_len / copy_size;
		if (chunk > MAX_CHUNK) {
			cw_bpf_debug("Exceeding the max_chunk:%d need_chunk:%d\n", MAX_CHUNK, chunk);
			return -1;
		}
	}
	u32 insert_pos = map_data->header_offset_idx;

//	if (chunk > MAX_CHUNK) {
//		cw_bpf_debug("Exceeding the max_chunk:%d need_chunk:%d\n", MAX_CHUNK, chunk);
//		return -1;
//	}

	// 挪动数据以腾出空间插入子串
	// 插入header后的结束位置
//	int tail_offset_idx = map_data->header_offset_idx + CW_STREAM_HEADER_LEN -1;

//	char subs[HEADER_LEN +1] = "cwtrace: 00:00:1015481350055581:5450531005555981:5610250100539899:304775019cd3218a304775019cd3218a:1001025098564810:140acc88cde8773f\r\n\0";

//	char subs[HEADER_LEN +1];
//	struct apm_span_context *cw_sc = build_sc();

/*111*/
//	char subs[HEADER_LEN +1];
//	span_context_to_cw_string_stream(cw_sc, subs, '1');






	char tmp_header[CW_STREAM_HEADER_LEN];
	struct apm_span_context *cw_sc = build_sc(proc_info);
	if (cw_sc == NULL){
		return 0;
	}
	span_context_to_cw_string_stream(cw_sc, tmp_header, '1');
//	cw_bpf_debug("header:[%s],sizeof %d ",map_data->header_stream,sizeof(map_data->header_stream));


	bpf_probe_read(map_data->header_stream, sizeof(map_data->header_stream), (void *) (tmp_header));
//	cw_bpf_debug("stream_header1->header------%s",map_data->header_stream);

//	if (insert_pos >= 0 && insert_pos <= JAVA_MAX_BUFFER_SIZE - sizeof(map_data->header_stream)) {
//		__builtin_memcpy(full_payload + insert_pos, map_data->header_stream, sizeof(map_data->header_stream));
//	}


	int tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1;


	// 修改长度
	long new_val;
	new_val = map_data->size + HEADER_LEN;

	long res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
	if (res < 0) {
		cw_bpf_debug("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);
//		cw_bpf_debug("Failed to write len to user address: %p\n", len_from_rbp_ptr);
		return -1;
	}


	// 单次
	if (end_str_len <= copy_size) {
		// 写头
		bpf_probe_write_user((void *) jbytechar_ptr + insert_pos, &map_data->header_stream, sizeof(map_data->header_stream));
		// 写尾
		bpf_probe_write_user((void *) jbytechar_ptr + tail_offset_idx, tail_payload, L7_IOVEC_BUF_SIZE);
	}

	// 多批
	if (end_str_len > copy_size) {
//		bpf_tail_call(ctx, &jmp_table1, 1);
		bpf_tail_call(ctx, &NAME(progs_jmp_up_map), PROG_DATA_JAVA_UPDATE_HEADER_UP_IDX);
	}

	return 0;
}

// bpf_prog_up_update_header
PROGUP(java_update_header)(struct pt_regs *ctx) {
	u32 k0 = 0;
	u32 copy_size = L7_IOVEC_BUF_SIZE;

	u32 k1 = 1;
	char *tail_payload = bpf_map_lookup_elem(&large_array_map, &k1);
	if (!tail_payload) {
		cw_bpf_debug("header ->end %s\n", tail_payload);
		return -1;
	}

	struct sock_t *map_data = bpf_map_lookup_elem(&socket_heap, &k0);
	if (!map_data) {
		cw_bpf_debug("Failed to lookup socket_heap\n");
		return -1;
	}

	u32 insert_pos = map_data->header_offset_idx;
	u32 tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1;
	u32 end_str_len = map_data->end_str_len;
	void * jbytechar_ptr = map_data->payload_char_p;

	if (!jbytechar_ptr) {
		return -1;
	}

	// 多批
	u32 chunk = end_str_len / copy_size;
	bpf_probe_write_user((void *) jbytechar_ptr + insert_pos, &map_data->header_stream, sizeof(map_data->header_stream));
	for (u32 i = 0; i < MAX_CHUNK; i++) {
		if (i <= chunk) {
			tail_offset_idx = insert_pos + sizeof(map_data->header_stream) - 1 + i * copy_size;
			bpf_probe_write_user((void *) jbytechar_ptr + tail_offset_idx, tail_payload + i * copy_size, L7_IOVEC_BUF_SIZE);
		}
	}
//
//	long new_val;
//	new_val = map_data->size + HEADER_LEN;
//
//	long res = bpf_probe_write_user((void *) len_from_rbp_ptr, &new_val, sizeof(new_val));
//	if (res < 0) cw_bpf_debug("Failed to write value to user address: %p, error: %ld\n", len_from_rbp_ptr, res);

	return 0;
}
