// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

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

//char __license[] SEC("license") = "Dual MIT/GPL";

#define PATH_MAX_LEN 128
#define MAX_BUCKETS 8
#define METHOD_MAX_LEN 8
#define MAX_CONCURRENT 50
#define REMOTE_ADDR_MAX_LEN 256
#define HOST_MAX_LEN 256
#define PROTO_MAX_LEN 8

struct http_server_span_t {
	BASE_SPAN_PROPERTIES
	u64 status_code;
	char method[METHOD_MAX_LEN];
	char path[PATH_MAX_LEN];
	char remote_addr[REMOTE_ADDR_MAX_LEN];
	char host[HOST_MAX_LEN];
	char proto[PROTO_MAX_LEN];
};

struct uprobe_data_t {
	struct http_server_span_t span;
	// bpf2go doesn't support pointers fields
	// saving the response pointer in the entry probe
	// and using it in the return probe
	u64 resp_ptr;
};

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

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct map_bucket));
	__uint(max_entries, MAX_CONCURRENT);
} golang_mapbucket_storage_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct cw_header_token));
	__uint(max_entries, MAX_CONCURRENT);
} header_range SEC(".maps");

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

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

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

struct {
	__uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
} events SEC(".maps");

struct header_key {
	char str[CW_HEADER_KEY_LENGTH];
};

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(key_size, sizeof(u32));
	__uint(value_size, sizeof(struct header_key));
	__uint(max_entries, 8);
} header_keys_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(u64));
	__uint(value_size, sizeof(char [CW_HEADER_VAL_LENGTH]));
	__uint(max_entries, MAX_CONCURRENT);
} http_server_context_headers SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_LRU_HASH);
	__uint(key_size, sizeof(u64));
	__uint(value_size, sizeof(char [CW_SYS_HEADER_VAL_LENGTH]));
	__uint(max_entries, MAX_CONCURRENT);
} http_server_context_sys_headers SEC(".maps");

// Injected in init
// volatile const u64 ctx_ptr_pos;
// volatile const u64 headers_ptr_pos;
// volatile const u64 buckets_ptr_pos;
//for (int i = 0; i < W3C_KEY_LENGTH; i++)
//{
//if (current_header_key[i] != "traceparent"[i]){
//goto outer_loop;
//}
//if (current_header_key[i] != "Traceparent"[i]){
//goto outer_loop;
//}
//}



MAP_BUCKET_DEFINITION(go_string_ot,  go_slice_ot)

static __always_inline char *
extract_context_from_req_headers_pre_parsed(void *key) {
	char *header_val =
			bpf_map_lookup_elem(&http_server_context_headers, &key);
	if (!header_val) {
		return NULL;
	}

	return header_val;
}


static __always_inline char *
extract_context_from_req_headers(void *key, void *headers_ptr_ptr, __u64 use_swiss) {
//	if (use_swiss == 1) {
		return extract_context_from_req_headers_pre_parsed(key);
//	}
//	return extract_context_from_req_headers_go_map(headers_ptr_ptr);
}


// This instrumentation attaches uprobe to the following function:
// func (sh serverHandler) ServeHTTP(rw ResponseWriter, req *Request)
SEC("uprobe/HandlerFunc_ServeHTTP")
int uprobe_HandlerFunc_ServeHTTP(struct pt_regs *ctx) {
	cw_bpf_debug("[uprobe_HandlerFunc_ServeHTTP]");
	__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)
    {
	    cw_bpf_debug("[uprobe_HandlerFunc_ServeHTTP] no proc info");
        return 0;
    }

	void *req_ctx_ptr = get_Go_context(ctx, 4, proc_info->ctx_ptr_pos, false);
	void *key = get_consistent_key(ctx, req_ctx_ptr);
	void *httpReq_ptr = bpf_map_lookup_elem(&http_server_uprobes, &key);
	if (httpReq_ptr != NULL) {
		cw_bpf_debug("uprobe/HandlerFunc_ServeHTTP already tracked with the current request");
		return 0;
	}
	
	u32 map_id = 0;
	struct uprobe_data_t *uprobe_data = bpf_map_lookup_elem(&http_server_uprobe_storage_map, &map_id);
	if (uprobe_data == NULL) {
		cw_bpf_debug("uprobe/HandlerFunc_ServeHTTP: http_server_span is NULL");
		return 0;
	}

	__builtin_memset(uprobe_data, 0, sizeof(struct uprobe_data_t));

	// Save response writer
	void *resp_impl = get_argument(ctx, 3);
	uprobe_data->resp_ptr = (u64) resp_impl;

	struct http_server_span_t *http_server_span = &uprobe_data->span;
	http_server_span->start_time = bpf_ktime_get_ns();
//	return 1;

	// Propagate context
	void *req_ptr = get_argument(ctx, 4);
//	struct map_bucket * map_bucket_p = get_map_bucket((void *) (req_ptr + proc_info->headers_ptr_pos));

	char *traceparent_header_value = extract_context_from_req_headers(key,
	                                                                  (void *) (req_ptr + proc_info->headers_ptr_pos),
	                                                                  proc_info->use_swiss_map);

//	bpf_printk("111 %s",traceparent_header_value);

//	char * traceparent_header_value = get_header_val_off(map_bucket_p);
	struct apm_span_context *cw_parent_span_context = bpf_map_lookup_elem(&cw_parent_span_context_storage_map, &map_id);
	if (!cw_parent_span_context) {
		return 0;
	}
	__builtin_memset(cw_parent_span_context, 0, sizeof(struct apm_span_context));

	// parentSys parentService
	char *header_sys_val = bpf_map_lookup_elem(&http_server_context_sys_headers, &key);
	if (header_sys_val) {
//		bpf_printk("has sys");
		bpf_probe_read(&cw_parent_span_context->sysvs, sizeof(cw_parent_span_context->sysvs), header_sys_val);
//		bpf_printk("has sys %s", cw_parent_span_context->sysvs);
	}

	if (traceparent_header_value != NULL) {
		bpf_printk("traceparent_header_value != NULL");

		cw_string_to_span_context(traceparent_header_value, cw_parent_span_context);
		// found parent context in http headers
//		http_server_span->psc = *parent_ctx;
		// copy traceID
		copy_byte_arrays(http_server_span->psc.TraceID, http_server_span->sc.TraceID, TRACE_ID_SIZE);
		// 生成随机SpanID
		generate_random_bytes(http_server_span->sc.SpanID, SPAN_ID_SIZE);
		cw_save_parent_tracking_span(cw_parent_span_context);
	} else {
		bpf_printk("traceparent_header_value == NULL");
		http_server_span->sc = generate_span_context();

		struct apm_span_context *cw_parent_span_context = bpf_map_lookup_elem(&cw_parent_span_context_storage_map, &map_id);
		if (!cw_parent_span_context) {
			return 0;
		}

		// struct apm_span_context context = {};
		generate_random_bytes(cw_parent_span_context->trace_id, TRACE_ID_SIZE);
		// generate_random_bytes(context.span_id, SPAN_ID_SIZE);
		// 保存 trace_id 到psc
		cw_save_parent_tracking_span(cw_parent_span_context);
	}

	if (req_ctx_ptr == NULL) {
//		cw_bpf_debug("uprobe/HandlerFunc_ServeHTTP: req_ctx_ptr is NULL");
		return 0;
	}
//	cw_bpf_debug("found parent context in http headers");

//	cw_bpf_debug("parent_span_context-TraceID---");


//    for (int i = 0; i < TRACE_ID_SIZE; i++) {
//        cw_bpf_debug("%02x", uprobe_data->span.sc.TraceID[i]);
//    }

//    cw_bpf_debug("parent_span_context-SpanID---");
//    for (int i = 0; i < SPAN_ID_SIZE; i++) {
//        cw_bpf_debug("%02x", uprobe_data->span.sc.SpanID[i]);
//    }

//	bpf_map_update_elem(&http_server_uprobes, &key, uprobe_data, 0);
//	start_tracking_span(req_ctx_ptr, &http_server_span->sc);

//	void save_tracking_span(struct apm_span_context *sc)

	return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (sh serverHandler) ServeHTTP(rw ResponseWriter, req *Request)
SEC("uprobe/HandlerFunc_ServeHTTP")
int uprobe_HandlerFunc_ServeHTTP_Returns(struct pt_regs *ctx) {
//	return 1;
	cw_bpf_debug("[uprobe_HandlerFunc_ServeHTTP_Returns]");
//	u64 end_time = bpf_ktime_get_ns();
	__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;
    }

	void *req_ctx_ptr = get_Go_context(ctx, 4, proc_info->ctx_ptr_pos, false);
	void *key = get_consistent_key(ctx, req_ctx_ptr);

	struct uprobe_data_t *uprobe_data = bpf_map_lookup_elem(&http_server_uprobes, &key);
	if (uprobe_data == NULL) {
		cw_bpf_debug("uprobe/HandlerFunc_ServeHTTP_Returns: entry_state is NULL");
		bpf_map_delete_elem(&http_server_context_headers, &key);
		bpf_map_delete_elem(&http_server_context_sys_headers, &key);
		return 0;
	}
	bpf_map_delete_elem(&http_server_uprobes, &key);
	bpf_map_delete_elem(&http_server_context_headers, &key);
	bpf_map_delete_elem(&http_server_context_sys_headers, &key);
	struct http_server_span_t *http_server_span = &uprobe_data->span;
//	void *resp_ptr = (void *) uprobe_data->resp_ptr;
//	void *req_ptr = NULL;
//	bpf_probe_read(&req_ptr, sizeof(req_ptr), (void *) (resp_ptr + req_ptr_pos));

//	http_server_span->end_time = end_time;
//
//	void *url_ptr = 0;
//	bpf_probe_read(&url_ptr, sizeof(url_ptr), (void *) (req_ptr + url_ptr_pos));
//	// Collect fields from response
//	read_go_string(req_ptr, method_ptr_pos, http_server_span->method, sizeof(http_server_span->method),
//	               "method from request");
//	read_go_string(url_ptr, path_ptr_pos, http_server_span->path, sizeof(http_server_span->path),
//	               "path from Request.URL");
//	read_go_string(req_ptr, remote_addr_pos, http_server_span->remote_addr, sizeof(http_server_span->remote_addr),
//	               "remote addr from Request.RemoteAddr");
//	read_go_string(req_ptr, host_pos, http_server_span->host, sizeof(http_server_span->host), "host from Request.Host");
//	read_go_string(req_ptr, proto_pos, http_server_span->proto, sizeof(http_server_span->proto),
//	               "proto from Request.Proto");
//
//	// status code
//	bpf_probe_read(&http_server_span->status_code, sizeof(http_server_span->status_code),
//	               (void *) (resp_ptr + status_code_pos));
//
//	bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, http_server_span, sizeof(*http_server_span));
	stop_tracking_span(&http_server_span->sc, &http_server_span->psc);


	clear_parent_span_context();
	cw_bpf_debug("HTTP_END");
	return 0;
}

// This instrumentation attaches uprobe to the following function:
// func (r *Reader) readContinuedLineSlice(lim int64, validateFirstLine func([]byte) error) ([]byte, error) {
SEC("uprobe/textproto_Reader_readContinuedLineSlice")
int uprobe_textproto_Reader_readContinuedLineSlice_Returns(struct pt_regs *ctx) {
	void *key = (void *) GOROUTINE(ctx);
	u64 len = (u64) GO_PARAM2(ctx);
	u8 *buf = (u8 *) GO_PARAM1(ctx);
	if (len == (CW_HEADER_KEY_LENGTH + CW_HEADER_VAL_LENGTH + 2)) {
		long cw_header_native = 0x3A65636172747763LL; // 小端序下的 "cwtrace:"
		__u64 key64 = 0;
		bpf_probe_read_user(&key64, sizeof(key64), buf);
		if (key64 == cw_header_native) {
			char header_val[CW_HEADER_VAL_LENGTH] = {};
			bpf_probe_read_user(header_val, sizeof(header_val), buf + CW_HEADER_KEY_LENGTH + 2);
			cw_bpf_debug("Swiss Map:%s", header_val);
			bpf_map_update_elem(&http_server_context_headers, &key, &header_val, BPF_ANY);
		}
	}
	// cwother: 08:service:03:tag
	// 解析 cwother header (变长，需要检查是否匹配)
	if (len >= (CW_SYS_HEADER_KEY_LENGTH + 2 + 8)) {  // 至少 "cwother: NN:name:MM " 的长度
		long cw_sys_header_native = 0x3A726568746F7763LL; // 小端序下的 "cwother:"
		__u64 key64 = 0;
		bpf_probe_read_user(&key64, sizeof(key64), buf);
		if (key64 == cw_sys_header_native) {
			char header_val[CW_SYS_HEADER_VAL_LENGTH] = {};
			bpf_probe_read_user(header_val, sizeof(header_val), buf + CW_SYS_HEADER_KEY_LENGTH + 2);
			cw_bpf_debug("Swiss Map:%s key %llu", header_val,key);
			bpf_map_update_elem(&http_server_context_sys_headers, &key, &header_val, BPF_ANY);
		}
	}
	return 0;
}
