#include "test_client.h"
#include <arpa/inet.h>


int connect_tcpserver(const char *ip, unsigned short port) {

	int connfd = socket(AF_INET, SOCK_STREAM, 0);

	struct sockaddr_in server_addr;
	memset(&server_addr, 0, sizeof(struct sockaddr_in));

	server_addr.sin_family = AF_INET;
	server_addr.sin_addr.s_addr = inet_addr(ip);
	server_addr.sin_port = htons(port);

	if (0 !=  connect(connfd, (struct sockaddr*)&server_addr, sizeof(struct sockaddr_in))) {
		perror("connect");
		return -1;
	}
	
	return connfd;
	
}


int send_msg(int connfd, char *msg, int length) {

	int res = send(connfd, msg, length, 0);
	if (res < 0) {
		perror("send");
		exit(1);
	}
	return res;
}

int recv_msg(int connfd, char *msg, int length) {

	int res = recv(connfd, msg, length, 0);
	if (res < 0) {
		perror("recv");
		exit(1);
	}
	return res;

}

void testcase(int connfd, uint8_t op, const char* key, uint32_t key_len, const char* value, 
	uint32_t value_len, rsp_ret_status_e st, const char* rsp_value, uint32_t expect_len, const char* command_name){
	uint8_t buf[CMD_SIZE];
	uint8_t result[CMD_SIZE];
	kvs_response_t rsp;
	int len, recv_len;

	len = getcmd(op, key, key_len, value, value_len, buf);
	send_msg(connfd, buf, len);
	recv_len = recv_msg(connfd, result, CMD_SIZE);
	if (parse_response(result, recv_len, &rsp) > 0) {
		PRESP(command_name, &rsp);
		if(!verify_response(&rsp, op, st, rsp_value, expect_len)) printf("%s\n", command_name);
	}else{
		printf("parser error\n");
	}

	return ;
}


void array_testcase_1w(int connfd) {

	int count = 1000;
	int i = 0;

	struct timeval tv_begin;
	gettimeofday(&tv_begin, NULL);

	for (i = 0;i < count;i ++) {
		testcase(connfd, KVS_CMD_SET, "name", 4, "l\r\0n", 4, KVS_STATUS_OK, NULL, 0, "SET NAME");
		testcase(connfd, KVS_CMD_GET, "name", 4, NULL, 0, KVS_STATUS_OK, "l\r\0n", 4, "GET NAME");
		testcase(connfd, KVS_CMD_MOD, "name", 4, "liu", 3, KVS_STATUS_OK, NULL, 0, "MOD NAME");
		testcase(connfd, KVS_CMD_GET, "name", 4, NULL, 0, KVS_STATUS_OK, "liu", 3, "GET NAME");
		testcase(connfd, KVS_CMD_EXIST, "name", 4, NULL, 0, KVS_STATUS_EXIST, NULL, 0, "EXIST NAME");
		testcase(connfd, KVS_CMD_DEL, "name", 4, NULL, 0, KVS_STATUS_OK, NULL, 0, "DEL NAME");
		testcase(connfd, KVS_CMD_EXIST, "name", 4, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "EXIST NAME");
	
		testcase(connfd, KVS_CMD_MOD, "stu", 3, "liu", 3, KVS_STATUS_NO_EXIST, NULL, 0, "MOD NAME");
		testcase(connfd, KVS_CMD_DEL, "stu", 3, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "DEL SUT");
	}

	struct timeval tv_end;
	gettimeofday(&tv_end, NULL);

	int time_used = TIME_SUB_MS(tv_end, tv_begin); // ms

	printf("array testcase --> time_used: %d, qps: %d\n", time_used, 9000 * 1000 / time_used);

}

void rbtree_testcase_1w(int connfd) {

	int count = 1000;
	int i = 0;

	struct timeval tv_begin;
	gettimeofday(&tv_begin, NULL);

	for (i = 0;i < count;i ++) {
		testcase(connfd, KVS_CMD_RSET, "name", 4, "l\r\0n", 4, KVS_STATUS_OK, NULL, 0, "RSET NAME");
		testcase(connfd, KVS_CMD_RGET, "name", 4, NULL, 0, KVS_STATUS_OK, "l\r\0n", 4, "RGET NAME");
		testcase(connfd, KVS_CMD_RMOD, "name", 4, "liu", 3, KVS_STATUS_OK, NULL, 0, "RMOD NAME");
		testcase(connfd, KVS_CMD_RGET, "name", 4, NULL, 0, KVS_STATUS_OK, "liu", 3, "RGET NAME");
		testcase(connfd, KVS_CMD_REXIST, "name", 4, NULL, 0, KVS_STATUS_EXIST, NULL, 0, "REXIST NAME");
		testcase(connfd, KVS_CMD_RDEL, "name", 4, NULL, 0, KVS_STATUS_OK, NULL, 0, "RDEL NAME");
		testcase(connfd, KVS_CMD_REXIST, "name", 4, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "REXIST NAME");
	
		testcase(connfd, KVS_CMD_RMOD, "stu", 3, "liu", 3, KVS_STATUS_NO_EXIST, NULL, 0, "RMOD NAME");
		testcase(connfd, KVS_CMD_RDEL, "stu", 3, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "RDEL SUT");
	}

	struct timeval tv_end;
	gettimeofday(&tv_end, NULL);

	int time_used = TIME_SUB_MS(tv_end, tv_begin); // ms

	printf("array testcase --> time_used: %d, qps: %d\n", time_used, 9000 * 1000 / time_used);

}

void hash_testcase_1w(int connfd) {

	int count = 1000;
	int i = 0;

	struct timeval tv_begin;
	gettimeofday(&tv_begin, NULL);

	for (i = 0;i < count;i ++) {
		testcase(connfd, KVS_CMD_HSET, "name", 4, "l\r\0n", 4, KVS_STATUS_OK, NULL, 0, "HSET NAME");
		testcase(connfd, KVS_CMD_HGET, "name", 4, NULL, 0, KVS_STATUS_OK, "l\r\0n", 4, "HGET NAME");
		testcase(connfd, KVS_CMD_HMOD, "name", 4, "liu", 3, KVS_STATUS_OK, NULL, 0, "HMOD NAME");
		testcase(connfd, KVS_CMD_HGET, "name", 4, NULL, 0, KVS_STATUS_OK, "liu", 3, "HGET NAME");
		testcase(connfd, KVS_CMD_HEXIST, "name", 4, NULL, 0, KVS_STATUS_EXIST, NULL, 0, "HEXIST NAME");
		testcase(connfd, KVS_CMD_HDEL, "name", 4, NULL, 0, KVS_STATUS_OK, NULL, 0, "HDEL NAME");
		testcase(connfd, KVS_CMD_HEXIST, "name", 4, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "HEXIST NAME");
	
		testcase(connfd, KVS_CMD_HMOD, "stu", 3, "liu", 3, KVS_STATUS_NO_EXIST, NULL, 0, "HMOD NAME");
		testcase(connfd, KVS_CMD_HDEL, "stu", 3, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "HDEL SUT");
	}

	struct timeval tv_end;
	gettimeofday(&tv_end, NULL);

	int time_used = TIME_SUB_MS(tv_end, tv_begin); // ms

	printf("array testcase --> time_used: %d, qps: %d\n", time_used, 9*count * 1000 / time_used);

}

void do_batch_test(int fd, int op, const char *key, const char *value, rsp_ret_status_e st, const char *rsp_value){
    kvs_batch_t batch;
    kvs_batch_init(&batch);

	char bkey[256]={0}, bval[256]={0};

    // 组 batch（最多 64 条）
	for(int i = 0;i < 100; ++ i){
		if(value == NULL){
			int klen = sprintf(bkey, "%s%d", key, i);
	    	kvs_batch_add(&batch, op, bkey, klen, NULL, 0);
		}else{
			int klen = sprintf(bkey, "%s%d", key, i);
			int vlen = sprintf(bval, "%s%d", value, i);
			kvs_batch_add(&batch, op, bkey, klen, bval, vlen);
		}
	}

    // 一次性发送
    kvs_batch_send(fd, &batch);

    // 一次性 recv + parse
    uint8_t recvbuf[BATCH_SIZE];
    kvs_response_t rsps[KVS_BATCH_MAX];

    int nrsp = kvs_batch_recv_parse(fd, &batch, rsps, recvbuf, sizeof(recvbuf));

    // 打印/处理
	for (int i = 0; i < nrsp; i++) {
		print_response(bkey, &rsps[i]);
		int vlen;
		if(rsp_value != NULL) vlen = sprintf(bval, "%s%d", rsp_value, i);
		else vlen = 0;
		verify_response(&rsps[i], op, st, bval, vlen);
	}
}

void array_testcase_1w_batch(int connfd) {
    kvs_batch_t batch;
	kvs_batch_init(&batch);

	int count = 1000;
	int i = 0;

	struct timeval tv_begin;
	gettimeofday(&tv_begin, NULL);

	for (i = 0;i < count;i ++) {
		batch.cnt = 0;
		batch.len = 0;
		kvs_batch_add(&batch, KVS_CMD_SET, "name", 4, "l\r\0n", 4);
		kvs_batch_add(&batch, KVS_CMD_GET, "name", 4, NULL, 0);
		kvs_batch_add(&batch, KVS_CMD_MOD, "name", 4, "liu", 3);
		kvs_batch_add(&batch, KVS_CMD_GET, "name", 4, NULL, 0);
		kvs_batch_add(&batch, KVS_CMD_EXIST, "name", 4, NULL, 0);
		kvs_batch_add(&batch, KVS_CMD_DEL, "name", 4, NULL, 0);
		kvs_batch_add(&batch, KVS_CMD_EXIST, "name", 4, NULL, 0);
	
		kvs_batch_add(&batch, KVS_CMD_MOD, "stu", 3, "liu", 3);
		kvs_batch_add(&batch, KVS_CMD_DEL, "stu", 3, NULL, 0);

		kvs_batch_send(connfd, &batch);

		uint8_t recvbuf[BATCH_SIZE];
		kvs_response_t rsps[KVS_BATCH_MAX];
		int nrsp = kvs_batch_recv_parse(connfd, &batch, rsps, recvbuf, sizeof(recvbuf));
	}

	struct timeval tv_end;
	gettimeofday(&tv_end, NULL);

	int time_used = TIME_SUB_MS(tv_end, tv_begin); // ms

	printf("array testcase --> time_used: %d, qps: %d\n", time_used, 9000 * 1000 / time_used);

}

void batch_qps(int connfd) {
    const int N = 1000000;
    const int B = 100;                 // do_batch_test() 里写死 50
	static char valA[256];
    static char valB[256];
    static char valC[256];

    static int inited = 0;
    if (!inited) {
        // 填充 255 字节，最后补 '\0'
        memset(valA, 'A', 255); valA[255] = '\0';
        memset(valB, 'B', 255); valB[255] = '\0';
        memset(valC, 'C', 255); valC[255] = '\0';
        inited = 1;
    }

    struct timeval tv_begin, tv_end;
    gettimeofday(&tv_begin, NULL);

    // ---------------- Phase 1: ADD 两条 DEL 一条 (100w) ----------------
    // 每轮：ADD/SET A_i, ADD/SET B_i, DEL A_i
    // 用 batch：每批处理 idx..idx+49
    for (int base = 0; base < N; base += B) {
        // prefix 必须短，避免 do_batch_test 里 bkey[15] 溢出
        // do_batch_test 会生成: prefix + i (0..49)
        // 我们把 base 编进 prefix，保证全局 key 唯一
        char preA[16], preB[16];

        // 例：A123450_0..A123450_49（base=123450）
        // 注意：这里 prefix 长度要尽量 <= 10 左右
        snprintf(preA, sizeof(preA), "A%d_", base/100);
        snprintf(preB, sizeof(preB), "B%d_", base/100);

        do_batch_test(connfd, KVS_CMD_RSET, preA, valA, KVS_STATUS_OK, NULL); // 50次 RSET A
        do_batch_test(connfd, KVS_CMD_RSET, preB, valB, KVS_STATUS_OK, NULL); // 50次 RSET B
        do_batch_test(connfd, KVS_CMD_RDEL, preA, NULL, KVS_STATUS_OK, NULL); // 50次 RDEL A

        if (base % 10000 == 0) printf("P1 base:%d\n", base);
    }
    printf("phase 1 end\n");

    // ---------------- Phase 2: ADD 一条 DEL 两条 (100w) ----------------
    // 每轮：ADD/SET C_i, DEL B_i, DEL C_i
    for (int base = 0; base < N; base += B) {
        char preB[16], preC[16];
        snprintf(preB, sizeof(preB), "B%d_", base/100);
        snprintf(preC, sizeof(preC), "C%d_", base/100);

        do_batch_test(connfd, KVS_CMD_RSET, preC, valC, KVS_STATUS_OK, NULL); // 50次 RSET C
        do_batch_test(connfd, KVS_CMD_RDEL, preB, NULL, KVS_STATUS_OK, NULL); // 50次 RDEL B
        do_batch_test(connfd, KVS_CMD_RDEL, preC, NULL, KVS_STATUS_OK, NULL); // 50次 RDEL C

        if (base % 10000 == 0) printf("P2 base:%d\n", base);
    }
    printf("phase 2 end\n");

    gettimeofday(&tv_end, NULL);
    int time_used = TIME_SUB_MS(tv_end, tv_begin);

    // 真实总 ops 还是 6*N （每轮 6 个操作）
    long long ops = (long long)N * 6;
    long long qps = (time_used > 0) ? (ops * 1000 / time_used) : 0;

    printf("BATCH(do_batch_test) ADD2-DEL1 then ADD1-DEL2 (N=%d) --> time_used=%d ms, ops=%lld, qps=%lld\n",
           N, time_used, ops, qps);
}

void save(int connfd){
	testcase(connfd, KVS_CMD_SAVE, NULL, 0, NULL, 0, KVS_STATUS_OK, NULL, 0, "SAVE");
}

void testcase_add2_del1_then_add1_del2_100w(int connfd) {
    const int N = 1000000;

    // 如果你有 KVS_CMD_ADD 就用 ADD；没有就用 SET（但 SET 会覆盖，意义不同）
    // 这里按你说的“会返回 EXIST”，所以我用 KVS_CMD_ADD 来写。
    // 若你实际命令叫 KVS_CMD_SET 且语义是 ADD，请自行替换。
    const char *valA = "va";
    const char *valB = "vb";
    const char *valC = "vc";

    char keyA[64], keyB[64], keyC[64];

    struct timeval tv_begin, tv_end;
    gettimeofday(&tv_begin, NULL);

    // ---------------- Phase 1: ADD 两条 DEL 一条 (100w) ----------------
    // 每轮：ADD A_i, ADD B_i, DEL A_i -> 留下 B_i
    for (int i = 0; i < N; i++) {
        int klenA = snprintf(keyA, sizeof(keyA), "A_%d", i);
        int klenB = snprintf(keyB, sizeof(keyB), "B_%d", i);

        testcase(connfd, KVS_CMD_RSET, keyA, klenA, valA, 2, KVS_STATUS_OK, NULL, 0, "P1 ADD A_i");
        testcase(connfd, KVS_CMD_RSET, keyB, klenB, valB, 2, KVS_STATUS_OK, NULL, 0, "P1 ADD B_i");
        testcase(connfd, KVS_CMD_RDEL, keyA, klenA, NULL, 0, KVS_STATUS_OK, NULL, 0, "P1 DEL A_i");
	
		if(i%10000 == 0) printf("i:%d\n", i);
	}

	printf("phase 1 end\n");

    // ---------------- Phase 2: ADD 一条 DEL 两条 (100w) ----------------
    // 每轮：ADD C_i, DEL B_i, DEL C_i -> 每轮净删一个 B_i
    for (int i = 0; i < N; i++) {
        int klenC = snprintf(keyC, sizeof(keyC), "C_%d", i);
        int klenB = snprintf(keyB, sizeof(keyB), "B_%d", i);

        testcase(connfd, KVS_CMD_RSET, keyC, klenC, valC, 2, KVS_STATUS_OK, NULL, 0, "P2 ADD C_i");
        testcase(connfd, KVS_CMD_RDEL, keyB, klenB, NULL, 0, KVS_STATUS_OK, NULL, 0, "P2 DEL B_i");
        testcase(connfd, KVS_CMD_RDEL, keyC, klenC, NULL, 0, KVS_STATUS_OK, NULL, 0, "P2 DEL C_i");
		
		if(i%10000 == 0) printf("i:%d\n", i);
	}

	printf("phase 2 end\n");

    // for (int j = 0; j < 5; j++) {
    //     int idx = (j == 0) ? 0 : (j == 1) ? (N/2) : (N-1);

    //     int klenA = snprintf(keyA, sizeof(keyA), "A_%d", idx);
    //     int klenB = snprintf(keyB, sizeof(keyB), "B_%d", idx);
    //     int klenC = snprintf(keyC, sizeof(keyC), "C_%d", idx);

    //     testcase(connfd, KVS_CMD_EXIST, keyA, klenA, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "FINAL A not exist");
    //     testcase(connfd, KVS_CMD_EXIST, keyB, klenB, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "FINAL B not exist");
    //     testcase(connfd, KVS_CMD_EXIST, keyC, klenC, NULL, 0, KVS_STATUS_NO_EXIST, NULL, 0, "FINAL C not exist");
    // }

    gettimeofday(&tv_end, NULL);
    int time_used = TIME_SUB_MS(tv_end, tv_begin);

    // 统计：Phase1 每轮3 ops，Phase2 每轮3 ops
    long long ops = (long long)N * 3 + (long long)N * 3;
    long long qps = (time_used > 0) ? (ops * 1000 / time_used) : 0;

    printf("ADD2-DEL1 then ADD1-DEL2 (N=%d) --> time_used=%d ms, ops=%lld, qps=%lld\n",
           N, time_used, ops, qps);
}

void send_spec_chars(int connfd){
	const char *v_ws = "li an\tok\nend";  /* 内容含空格、\t、\n */
	int v_ws_len = 12;                   /* 手算：l i ' ' a n \t o k \n e n d = 12 */

	testcase(connfd, KVS_CMD_RSET,
			"ws", 2,
			v_ws, v_ws_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET WHITESPACE");

	testcase(connfd, KVS_CMD_RGET,
			"ws", 2,
			NULL, 0,
			KVS_STATUS_OK,
			v_ws, v_ws_len,
			"RGET WHITESPACE");

	/* 2) 引号与反斜杠：测试是否被错误转义 */
	const char *v_quote = "he\"llo\\world'!"; 
	/* 字节数：h e " l l o \ w o r l d ' ! = 15 */
	int v_quote_len = 15;

	testcase(connfd, KVS_CMD_RSET,
			"quote", 5,
			v_quote, v_quote_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET QUOTE BACKSLASH");

	testcase(connfd, KVS_CMD_RGET,
			"quote", 5,
			NULL, 0,
			KVS_STATUS_OK,
			v_quote, v_quote_len,
			"RGET QUOTE BACKSLASH");


	/* 3) 分隔符字符：冒号/逗号/分号/竖线 */
	const char *v_sep = "a:b,c;d|e";
	int v_sep_len = 9;   /* a : b , c ; d | e = 9 */

	testcase(connfd, KVS_CMD_RSET,
			"sep", 3,
			v_sep, v_sep_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET SEPARATORS");

	testcase(connfd, KVS_CMD_RGET,
			"sep", 3,
			NULL, 0,
			KVS_STATUS_OK,
			v_sep, v_sep_len,
			"RGET SEPARATORS");

	/* 4) CRLF：\r\n 组合（最容易把一条请求拆错/响应拼错） */
	const char *v_crlf = "line1\r\nline2";
	int v_crlf_len = 12; /* line1(5) + \r(1) + \n(1) + line2(5) = 12 */

	testcase(connfd, KVS_CMD_RSET,
			"crlf", 4,
			v_crlf, v_crlf_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET CRLF");

	testcase(connfd, KVS_CMD_RGET,
			"crlf", 4,
			NULL, 0,
			KVS_STATUS_OK,
			v_crlf, v_crlf_len,
			"RGET CRLF");

	/* 5) 二进制数据：包含 \0，必须按长度处理，不能用 strlen */
	char v_bin[] = { 'A', 0x00, 'B', 'C', 0x00, 'D' };
	int v_bin_len = 6;

	testcase(connfd, KVS_CMD_RSET,
			"bin", 3,
			v_bin, v_bin_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET BINARY WITH NUL");

	testcase(connfd, KVS_CMD_RGET,
			"bin", 3,
			NULL, 0,
			KVS_STATUS_OK,
			v_bin, v_bin_len,
			"RGET BINARY WITH NUL");

	/* 6) UTF-8：中文 + emoji（测试多字节） */
	const char *v_utf8 = "中文🙂"; 
	/* "中"(3字节) + "文"(3字节) + "🙂"(4字节) = 10字节 */
	int v_utf8_len = 10;

	testcase(connfd, KVS_CMD_RSET,
			"utf8", 4,
			v_utf8, v_utf8_len,
			KVS_STATUS_OK,
			NULL, 0,
			"RSET UTF8");

	testcase(connfd, KVS_CMD_RGET,
			"utf8", 4,
			NULL, 0,
			KVS_STATUS_OK,
			v_utf8, v_utf8_len,
			"RGET UTF8");

}

int main(int argc, char *argv[]) {
	if (argc != 4) {
		printf("arg error\n");
		return -1;
	}

	char *ip = argv[1];
	int port = atoi(argv[2]);
	int mode = atoi(argv[3]);

	int connfd = connect_tcpserver(ip, port);

	if(mode == 0){
		array_testcase_1w(connfd);
	}else if(mode == 1){
		array_testcase_1w_batch(connfd);
	}else if(mode == 2){
		rbtree_testcase_1w(connfd);
	}else if(mode == 3){
		hash_testcase_1w(connfd);
	}else if(mode == 4){
		testcase_add2_del1_then_add1_del2_100w(connfd);
	}else if(mode == 10){
		do_batch_test(connfd, KVS_CMD_SET, "array_set", "array_val", KVS_STATUS_OK, NULL);
	}else if(mode == 11){
		do_batch_test(connfd, KVS_CMD_GET, "array_set", NULL, KVS_STATUS_OK, "array_val");
	}else if(mode == 12){
		do_batch_test(connfd, KVS_CMD_EXIST, "array_set", NULL, KVS_STATUS_EXIST, NULL);
	}else if(mode == 13){
		do_batch_test(connfd, KVS_CMD_DEL, "array_set", NULL, KVS_STATUS_OK, NULL);
	}else if(mode == 20){
		do_batch_test(connfd, KVS_CMD_RSET, "rbtree_set", "rbtree_val", KVS_STATUS_OK, NULL);
	}else if(mode == 21){
		do_batch_test(connfd, KVS_CMD_RGET, "rbtree_set", NULL, KVS_STATUS_OK, "rbtree_val");
	}else if(mode == 22){
		do_batch_test(connfd, KVS_CMD_REXIST, "rbtree_set", NULL, KVS_STATUS_OK, NULL);
	}else if(mode == 30){
		do_batch_test(connfd, KVS_CMD_HSET, "hash_set", "hash_val", KVS_STATUS_OK, NULL);
	}else if(mode == 31){
		do_batch_test(connfd, KVS_CMD_HGET, "hash_set", NULL, KVS_STATUS_OK, "hash_val");
	}else if(mode == 32){
		do_batch_test(connfd, KVS_CMD_HEXIST, "hash_set", NULL, KVS_STATUS_OK, NULL);
	}else if(mode == -1){
		save(connfd);
	}else if(mode == 5){
		batch_qps(connfd);
	}else if(mode == 6){
		send_spec_chars(connfd);
	}
	
	return 0;
}