// kvs_client_min.c
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>

#define KVS_MAGIC_U32 0x4B565331u // 'KVS1'
#define KVS_HDR_LEN   16
#define KVS_TYPE_REQ  1
#define KVS_TYPE_RESP 2

// OP mapping (match your server enum)
#define OP_SET   0
#define OP_GET   1
#define OP_DEL   2
#define OP_MOD   3
#define OP_EXIST 4

static int write_u8(uint8_t **pp, uint8_t *end, uint8_t v) {
	if (*pp + 1 > end) return -1;
	**pp = v;
	(*pp)++;
	return 0;
}
static int write_u32(uint8_t **pp, uint8_t *end, uint32_t v) {
	if (*pp + 4 > end) return -1;
	uint32_t be = htonl(v);
	memcpy(*pp, &be, 4);
	(*pp) += 4;
	return 0;
}
static int write_bytes(uint8_t **pp, uint8_t *end, const void *buf, size_t n) {
	if (*pp + n > end) return -1;
	memcpy(*pp, buf, n);
	(*pp) += n;
	return 0;
}

static int read_u8(const uint8_t **pp, const uint8_t *end, uint8_t *out) {
	if (*pp + 1 > end) return -1;
	*out = **pp;
	(*pp)++;
	return 0;
}
static int read_u32(const uint8_t **pp, const uint8_t *end, uint32_t *out) {
	if (*pp + 4 > end) return -1;
	uint32_t be;
	memcpy(&be, *pp, 4);
	*out = ntohl(be);
	(*pp) += 4;
	return 0;
}

static ssize_t send_all(int fd, const void *buf, size_t len) {
	const uint8_t *p = (const uint8_t*)buf;
	size_t sent = 0;
	while (sent < len) {
		ssize_t n = send(fd, p + sent, len - sent, 0);
		if (n < 0) {
			if (errno == EINTR) continue;
			return -1;
		}
		if (n == 0) return -1;
		sent += (size_t)n;
	}
	return (ssize_t)sent;
}

static ssize_t recv_all(int fd, void *buf, size_t len) {
	uint8_t *p = (uint8_t*)buf;
	size_t recvd = 0;
	while (recvd < len) {
		ssize_t n = recv(fd, p + recvd, len - recvd, 0);
		if (n < 0) {
			if (errno == EINTR) continue;
			return -1;
		}
		if (n == 0) return -1;
		recvd += (size_t)n;
	}
	return (ssize_t)recvd;
}

// Build one request packet: header + payload(opcount + cmds...)
static int build_req_packet(uint32_t reqId,
                            uint8_t opcount,
                            uint8_t op,
                            int argc, const char *argv[],
                            uint8_t *out, size_t out_cap, size_t *out_len) {
	if (!out || out_cap < KVS_HDR_LEN || !out_len) return -1;

	uint8_t *w = out + KVS_HDR_LEN;
	uint8_t *end = out + out_cap;

	// payload: opcount(4)
	if (write_u32(&w, end, (uint32_t)opcount) < 0) return -1;

	// one cmd (repeat if you want multiple, keep minimal here)
	if (write_u8(&w, end, op) < 0) return -1;
	if (write_u32(&w, end, (uint32_t)argc) < 0) return -1;
	for (int i = 0; i < argc; i++) {
		uint32_t alen = (uint32_t)strlen(argv[i]);
		if (write_u32(&w, end, alen) < 0) return -1;
		if (write_bytes(&w, end, argv[i], alen) < 0) return -1;
	}

	uint32_t payloadLen = (uint32_t)(w - (out + KVS_HDR_LEN));
	uint32_t magic_be = htonl(KVS_MAGIC_U32);
	uint32_t payload_be = htonl(payloadLen);
	uint32_t reqid_be = htonl(reqId);

	// header layout:
	// | magic(4) | type(1) | payloadLen(4) | reqId(4) | flag1(1) | flag2(2) |
	memcpy(out + 0, &magic_be, 4);
	out[4] = KVS_TYPE_REQ;
	memcpy(out + 5, &payload_be, 4);
	memcpy(out + 9, &reqid_be, 4);
	out[13] = 0;           // flag1
	out[14] = 0; out[15] = 0; // flag2

	*out_len = (size_t)(KVS_HDR_LEN + payloadLen);
	return 0;
}

/**
 * 	KVS_STATUS_OK = 0,
	KVS_STATUS_ERROR = 1,
	KVS_STATUS_NO_EXIST = 2,
	KVS_STATUS_EXIST = 3,
	KVS_STATUS_BADREQ = 4
 */
const char *rsp_status[5] = {"OK", "ERROR", "NO_EXIST", "EXIST", "ERROR"};
static int recv_and_print_resp(int fd) {
	uint8_t hdr[KVS_HDR_LEN];

	if (recv_all(fd, hdr, KVS_HDR_LEN) < 0) {
		perror("recv header");
		return -1;
	}

	uint32_t magic_be, payload_be, reqid_be;
	memcpy(&magic_be, hdr + 0, 4);
	memcpy(&payload_be, hdr + 5, 4);
	memcpy(&reqid_be, hdr + 9, 4);

	uint32_t magic = ntohl(magic_be);
	uint8_t type = hdr[4];
	uint32_t payloadLen = ntohl(payload_be);
	uint32_t reqId = ntohl(reqid_be);

	if (magic != KVS_MAGIC_U32 || type != KVS_TYPE_RESP) {
		fprintf(stderr, "bad response header: magic=0x%x type=%u\n", magic, type);
		return -1;
	}

	uint8_t *payload = (uint8_t*)malloc(payloadLen);
	if (!payload) return -1;

	if (payloadLen > 0) {
		if (recv_all(fd, payload, payloadLen) < 0) {
			perror("recv payload");
			free(payload);
			return -1;
		}
	}

	// parse response payload:
	// | opcount(4) | repeat Cmd |
	// Cmd: | status(1) | datalen(4) | data |
	const uint8_t *p = payload;
	const uint8_t *end = payload + payloadLen;

	uint32_t opcount = 0;
	if (read_u32(&p, end, &opcount) < 0) {
		fprintf(stderr, "resp parse failed\n");
		free(payload);
		return -1;
	}

	printf("RESP reqId=%u opcount=%u\n", reqId, opcount);

	for (uint32_t i = 0; i < opcount; i++) {
		uint8_t status = 0;
		uint32_t dlen = 0;
		if (read_u8(&p, end, &status) < 0) goto bad;
		if (read_u32(&p, end, &dlen) < 0) goto bad;
		if (p + dlen > end) goto bad;

		printf("  op[%u]: status=%s datalen=%u", i, rsp_status[status], dlen);
		if (dlen > 0) {
			// printf(" data=\"%.*s\"", (int)dlen, (const char*)p);
            printf(" data=\"");
            for (int i = 0; i < dlen; i++) {
                printf("%02X", (unsigned char)p[i]);
                if (i + 1 < dlen)
                    printf(" ");
            }
            printf("\"");
		}
		printf("\n");
		p += dlen;
	}

	free(payload);
	return 0;

bad:
	fprintf(stderr, "resp parse failed (truncated)\n");
	free(payload);
	return -1;
}

uint8_t buf[4096];
size_t pkt_len = 0;
uint32_t reqId = 1;
int fd;

int testcase(int op, const char *args[], int argnum){
    if (build_req_packet(reqId++, 1, op, argnum, args, buf, sizeof(buf), &pkt_len) < 0) {
        fprintf(stderr, "build failed\n");
        close(fd);
        return 1;
    }
    if (send_all(fd, buf, pkt_len) < 0) { perror("send SET"); close(fd); return 1; }
    if (recv_and_print_resp(fd) < 0) { close(fd); return 1; } 
    return 0;
}

#define BUF_CAP  65536
int main(int argc, char **argv) {
	if (argc != 3) {
		fprintf(stderr, "Usage: %s <ip> <port>\n", argv[0]);
		return 1;
	}
	const char *ip = argv[1];
	int port = atoi(argv[2]);

	fd = socket(AF_INET, SOCK_STREAM, 0);
	if (fd < 0) { perror("socket"); return 1; }

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = htons((uint16_t)port);
	if (inet_pton(AF_INET, ip, &addr.sin_addr) != 1) {
		fprintf(stderr, "bad ip\n");
		close(fd);
		return 1;
	}
	if (connect(fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
		perror("connect");
		close(fd);
		return 1;
    }

    size_t pkt_len = 0;
    uint32_t reqId = 1;

	// ---- SET key value ----
	{
		const char *args[] = {"foo", "bar"};
        printf("SET ");
        for (const unsigned char *p = (const unsigned char *)args[1]; *p; p++) {
            printf("%02X ", *p);
        }
        printf("\n");
        testcase(OP_SET, args, 2);
	}

	// ---- GET key ----
	{
		const char *args[] = {"foo"};
		testcase(OP_GET, args, 1);
	}

    // ---- SET key ----
	{
		const char *args[] = {"text", "\r\n\0aaaa\r\n\0"};
        printf("SET ");
        for (const unsigned char *p = (const unsigned char *)args[1]; *p; p++) {
            printf("%02X ", *p);
        }
        printf("\n");
		testcase(OP_SET, args, 2);
	}

    // ---- GET key ----
	{
		const char *args[] = {"text"};
		testcase(OP_GET, args, 1);
	}

	close(fd);
	return 0;
}