#include "kvstore.h"
#include "kvs_rw_tools.h"
#include <arpa/inet.h>

int kvs_need(const uint8_t *p, const uint8_t *end, size_t n) {
	return (p + n <= end) ? 0 : -1;
}

// 注意u8类型不需要ntoh或者hton
int kvs_read_u8(const uint8_t **pp, const uint8_t *end, uint8_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 1) < 0) return -1;
	*out = *p;
	*pp = p + 1;
	return 0;
}

int kvs_read_u16(const uint8_t **pp, const uint8_t *end, uint16_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 2) < 0) return -1;
	uint16_t v;
	memcpy(&v, p, 2);
	*out = ntohs(v);
	*pp = p + 2;
	return 0;
}

int kvs_read_u32(const uint8_t **pp, const uint8_t *end, uint32_t *out) {
	const uint8_t *p = *pp;
	if (kvs_need(p, end, 4) < 0) return -1;
	uint32_t v;
	memcpy(&v, p, 4);
	*out = ntohl(v);
	*pp = p + 4;
	return 0;
}

int kvs_write_u8(uint8_t **pp, const uint8_t *end, uint8_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 1) < 0) return -1;
	*p = v;
	*pp = p + 1;
	return 0;
}

int kvs_write_u16(uint8_t **pp, const uint8_t *end, uint16_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 2) < 0) return -1;
	uint16_t be = htons(v);
	memcpy(p, &be, 2);
	*pp = p + 2;
	return 0;
}

int kvs_write_u32(uint8_t **pp, const uint8_t *end, uint32_t v) {
	uint8_t *p = *pp;
	if (kvs_need(p, end, 4) < 0) return -1;
	uint32_t be = htonl(v);
	memcpy(p, &be, 4);
	*pp = p + 4;
	return 0;
}

// return: -1 fail, 0 half, >0 consumed
int kvs_parse_request(const uint8_t *msg, int length, kvs_request_t *req_out) {
	if (!msg || length <= 0 || !req_out) return -1;
	memset(req_out, 0, sizeof(*req_out));

	if(length < KVS_HDR_LEN) return 0;

	// header parser
	uint32_t magic_be;
	memcpy(&magic_be, msg, 4);
	uint32_t magic = ntohl(magic_be);
	if (magic != KVS_MAGIC) return -1;

	// 
	uint8_t type = *(msg+4);

	uint32_t payloadlen_be;
	memcpy(&payloadlen_be, msg+5, 4);
	uint32_t payloadlen = ntohl(payloadlen_be);

	uint32_t reqId_be;
	memcpy(&reqId_be, msg + 9, 4);
	uint32_t reqId = ntohl(reqId_be);

	uint8_t flag1 = *(msg + 13);

	uint16_t flag2_be;
	memcpy(&flag2_be, msg + 14, 2);
	uint16_t flag2 = ntohs(flag2_be);

	uint32_t totalLen = (uint32_t)KVS_HDR_LEN + payloadlen;
	if ((uint32_t)length < totalLen) return 0;          // 半包


	// payload parser
	const uint8_t *ptr = msg + KVS_HDR_LEN;
	const uint8_t *end = ptr + payloadlen;

	uint32_t opcount = 0;
	if(kvs_read_u32(&ptr, end, &opcount) < 0) return -1;
	if (opcount > KVS_MAX_OPCOUNT) return -1;

	kvs_op_t *ops = (kvs_op_t *)kvs_malloc(sizeof(kvs_op_t) * opcount);
	if (!ops) return -1;
	memset(ops, 0, sizeof(kvs_op_t) * opcount);

	// operator parser
	for (uint32_t i = 0; i < opcount; i++) {
		uint8_t op = 0;
		uint32_t argc = 0;
		
		if (kvs_read_u8(&ptr, end, &op) < 0) goto FAIL;
		if (kvs_read_u32(&ptr, end, &argc) < 0) goto FAIL;
		if (argc > 16) goto FAIL;

		char **argv = NULL;
		if (argc > 0) {
			argv = (char **)kvs_malloc(sizeof(char*) * argc);
			if (!argv) goto FAIL;
			memset(argv, 0, sizeof(char*) * argc);
		}

		// operator args parser
		for (uint32_t a = 0; a < argc; a++) {
			uint32_t arglen = 0;
			if (kvs_read_u32(&ptr, end, &arglen) < 0) goto FAIL;
			if (kvs_need(ptr, end, arglen) < 0) goto FAIL;

			char *s = (char *)kvs_malloc((size_t)arglen + 1);
			if (!s) goto FAIL;

			memcpy(s, ptr, arglen);
			s[arglen] = '\0';
			argv[a] = s;
			ptr += arglen;
		}

		ops[i].op = op;
		ops[i].argc = argc;
		ops[i].argv = argv;
	}

	req_out->magic = magic;
	req_out->type = type;
	req_out->payloadLen = payloadlen;
	req_out->reqId = reqId;
	req_out->flag1 = flag1;
	req_out->flag2 = flag2;
	req_out->opcount = opcount;
	req_out->ops = ops;
	req_out->consumed = totalLen;

	return 	(int)totalLen;

FAIL:
	if (ops) {
		for (uint32_t i = 0; i < opcount; i++) {
			if (ops[i].argv) {
				for (uint32_t a = 0; a < ops[i].argc; a++) {
					if (ops[i].argv[a]) kvs_free(ops[i].argv[a]);
				}
				kvs_free(ops[i].argv);
			}
		}
		kvs_free(ops);
	}
	return -1;
}

void kvs_free_request(kvs_request_t *req) {
	if (!req) return;
	if (req->ops) {
		for (uint32_t i = 0; i < req->opcount; i++) {
			if (req->ops[i].argv) {
				for (uint32_t a = 0; a < req->ops[i].argc; a++) {
					if (req->ops[i].argv[a]) kvs_free(req->ops[i].argv[a]);
				}
				kvs_free(req->ops[i].argv);
			}
		}
		kvs_free(req->ops);
	}
	memset(req, 0, sizeof(*req));
}

int kvs_execute_request(const kvs_request_t *req, kvs_response_t *results /* size opcount */) {
	if (!req || !results) return -1;

	for (uint32_t i = 0; i < req->opcount; i++) {
		uint8_t status = KVS_STATUS_ERROR;
		const char *data = NULL;
		uint32_t dlen = 0;

		kvs_exec_one_op(req->ops[i].op, req->ops[i].argc, req->ops[i].argv,
		                &status, &data, &dlen);

		results[i].status = status;
		results[i].datalen = dlen;
		results[i].data_ptr = data;
	}
	return 0;
}


 int kvs_build_response(const kvs_request_t *req,
                              const kvs_response_t *results,
                              uint8_t *response, int response_cap) {
	if (!req || !results || !response || response_cap < KVS_HDR_LEN) return -1;

	uint8_t *out = response;
	uint8_t *w = out + KVS_HDR_LEN;
	const uint8_t *w_end = out + response_cap;

	// payload: opcount
	if (kvs_write_u32(&w, w_end, req->opcount) < 0) return -1;

	for (uint32_t i = 0; i < req->opcount; i++) {
		if (kvs_write_u8(&w, w_end, results[i].status) < 0) return -1;
		if (kvs_write_u32(&w, w_end, results[i].datalen) < 0) return -1;

		if (results[i].datalen > 0) {
			if (kvs_need(w, w_end, results[i].datalen) < 0) return -1;
			memcpy(w, results[i].data_ptr, results[i].datalen);
			w += results[i].datalen;
		}
	}

	uint32_t respPayloadLen = (uint32_t)(w - (out + KVS_HDR_LEN));
	uint32_t total = KVS_HDR_LEN + respPayloadLen;

	// header
	{
		uint32_t be = htonl(KVS_MAGIC);
		memcpy(out + 0, &be, 4);
	}
	out[4] = (uint8_t)KVS_TYPE_RESP;
	{
		uint32_t be = htonl(respPayloadLen);
		memcpy(out + 5, &be, 4);
	}
	{
		uint32_t be = htonl(req->reqId);
		memcpy(out + 9, &be, 4);
	}
	out[13] = req->flag1;
	{
		uint16_t be = htons(req->flag2);
		memcpy(out + 14, &be, 2);
	}

	return (int)total;
}