/**
 * test_ws_protocol.c — WebSocket 协议层单元测试
 *
 * 使用标准 Linux socket 替代 DPDK 数据面，验证:
 *   1. WebSocket 握手 (HTTP Upgrade → 101)
 *   2. WS 帧编码/解码 (masked send / unmasked recv)
 *   3. 订阅/推送 消息交互
 *
 * 编译: gcc -o test_ws_protocol test_ws_protocol.c -lpthread
 * 运行: ./test_ws_protocol [server_ip] [port]
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <time.h>
#include <arpa/inet.h>
#include <netinet/tcp.h>
#include <sys/socket.h>

/* ═══════════════════════════════════════════════════════
 *  WebSocket 常量 (RFC 6455)
 * ═══════════════════════════════════════════════════════ */

#define WS_OPCODE_TEXT          0x1
#define WS_OPCODE_BINARY        0x2
#define WS_OPCODE_CLOSE         0x8
#define WS_OPCODE_PING          0x9
#define WS_OPCODE_PONG          0xA

#define WS_MASK_KEY_SIZE        4
#define WS_MAX_FRAME_SIZE       (64 * 1024)

/* ── Base64 ─────────────────────────────────────────── */

static const char b64_table[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

static int base64_encode(const unsigned char *src, int len, char *dst) {
    int i, j = 0;
    for (i = 0; i < len; i += 3) {
        unsigned int val  = (unsigned char)src[i] << 16;
        unsigned int val2 = (i + 1 < len) ? ((unsigned char)src[i + 1] << 8) : 0;
        unsigned int val3 = (i + 2 < len) ? (unsigned char)src[i + 2] : 0;
        val |= val2 | val3;
        dst[j++] = b64_table[(val >> 18) & 0x3F];
        dst[j++] = b64_table[(val >> 12) & 0x3F];
        dst[j++] = (i + 1 < len) ? b64_table[(val >> 6) & 0x3F] : '=';
        dst[j++] = (i + 2 < len) ? b64_table[val & 0x3F] : '=';
    }
    dst[j] = '\0';
    return j;
}

/* ── WebSocket 握手 ─────────────────────────────────── */

static int ws_handshake(int fd, const char *host, int port, const char *path) {
    /* Generate random key */
    uint8_t key_bytes[16];
    char ws_key[32];
    srand(time(NULL));
    for (int i = 0; i < 16; i++) key_bytes[i] = rand() % 256;
    base64_encode(key_bytes, 16, ws_key);

    /* Build HTTP Upgrade request */
    char req[2048];
    int len = snprintf(req, sizeof(req),
        "GET %s HTTP/1.1\r\n"
        "Host: %s:%d\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: %s\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n",
        path, host, port, ws_key);

    printf("[WS] Sending handshake:\n%s", req);

    if (send(fd, req, len, 0) != len) {
        printf("[WS] ERROR: send handshake failed\n");
        return -1;
    }

    /* Receive response */
    char resp[4096];
    memset(resp, 0, sizeof(resp));
    int total = 0;
    while (total < (int)sizeof(resp) - 1) {
        int n = recv(fd, resp + total, sizeof(resp) - 1 - total, 0);
        if (n <= 0) { printf("[WS] recv failed: %d\n", n); return -1; }
        total += n;
        resp[total] = '\0';
        if (strstr(resp, "\r\n\r\n")) break;
    }

    printf("[WS] Response:\n%s\n", resp);

    if (!strstr(resp, "101")) {
        printf("[WS] ERROR: not 101\n");
        return -1;
    }
    printf("[WS] Handshake OK!\n");
    return 0;
}

/* ── WebSocket 帧编码 (client → server, MASKED) ─────── */

static int ws_send_frame(int fd, int opcode, const uint8_t *payload, int len) {
    uint8_t frame[WS_MAX_FRAME_SIZE];
    int pos = 0;

    frame[pos++] = 0x80 | (opcode & 0x0F); /* FIN=1 */

    uint8_t mask[4];
    for (int i = 0; i < 4; i++) mask[i] = rand() % 256;

    if (len < 126) {
        frame[pos++] = 0x80 | (uint8_t)len;
    } else if (len < 65536) {
        frame[pos++] = 0x80 | 126;
        frame[pos++] = (len >> 8) & 0xFF;
        frame[pos++] = len & 0xFF;
    } else {
        frame[pos++] = 0x80 | 127;
        for (int i = 7; i >= 0; i--)
            frame[pos++] = (len >> (i * 8)) & 0xFF;
    }

    memcpy(&frame[pos], mask, 4);
    pos += 4;

    for (int i = 0; i < len; i++)
        frame[pos + i] = payload[i] ^ mask[i % 4];
    pos += len;

    int sent = send(fd, frame, pos, 0);
    return sent == pos ? sent : -1;
}

static int ws_send_text(int fd, const char *text) {
    return ws_send_frame(fd, WS_OPCODE_TEXT, (const uint8_t *)text, strlen(text));
}

/* ── WebSocket 帧解码 (server → client, UNMASKED) ──── */

static int ws_recv_frame(int fd, int *opcode, uint8_t *buf, int *buflen) {
    uint8_t header[10];
    int total = 0;
    while (total < 2) {
        int n = recv(fd, header + total, 2 - total, 0);
        if (n <= 0) return n;
        total += n;
    }

    uint8_t byte0 = header[0], byte1 = header[1];
    *opcode = byte0 & 0x0F;
    int masked = (byte1 >> 7) & 0x01;
    uint64_t payload_len = byte1 & 0x7F;

    if (payload_len == 126) {
        while (total < 4) { int n = recv(fd, header + total, 4 - total, 0); if (n <= 0) return n; total += n; }
        payload_len = ((uint64_t)header[2] << 8) | header[3];
    } else if (payload_len == 127) {
        while (total < 10) { int n = recv(fd, header + total, 10 - total, 0); if (n <= 0) return n; total += n; }
        payload_len = 0;
        for (int i = 2; i < 10; i++) payload_len = (payload_len << 8) | header[i];
    }

    uint8_t mask_key[4] = {0};
    if (masked) {
        uint8_t mk[4]; int mk_total = 0;
        while (mk_total < 4) { int n = recv(fd, mk + mk_total, 4 - mk_total, 0); if (n <= 0) return n; mk_total += n; }
        memcpy(mask_key, mk, 4);
    }

    if (*opcode == WS_OPCODE_CLOSE) { printf("[WS] CLOSE frame\n"); return 0; }

    if (*opcode == WS_OPCODE_PING) {
        uint8_t pd[128]; int pl = payload_len < 128 ? (int)payload_len : 128;
        int rt = 0; while (rt < pl) { int n = recv(fd, pd + rt, pl - rt, 0); if (n <= 0) return n; rt += n; }
        if (masked) for (int i = 0; i < pl; i++) pd[i] ^= mask_key[i % 4];
        ws_send_frame(fd, WS_OPCODE_PONG, pd, pl);
        printf("[WS] PING → PONG\n");
        *buflen = 0; return 0;
    }

    if (*opcode == WS_OPCODE_PONG) {
        uint8_t pd[128]; int pl = payload_len < 128 ? (int)payload_len : 128;
        int rt = 0; while (rt < pl) { int n = recv(fd, pd + rt, pl - rt, 0); if (n <= 0) return n; rt += n; }
        printf("[WS] PONG\n");
        *buflen = 0; return 0;
    }

    if (payload_len > (uint64_t)(*buflen)) {
        printf("[WS] payload too large: %lu > %d\n", payload_len, *buflen);
        return -1;
    }

    int rt = 0;
    while (rt < (int)payload_len) {
        int n = recv(fd, buf + rt, (int)payload_len - rt, 0);
        if (n <= 0) return n;
        rt += n;
    }

    if (masked) for (int i = 0; i < (int)payload_len; i++) buf[i] ^= mask_key[i % 4];

    *buflen = (int)payload_len;
    return *buflen;
}

/* ═══════════════════════════════════════════════════════
 *  Main test
 * ═══════════════════════════════════════════════════════ */

static int test_case = 1;

static void test_send_recv(int fd) {
    printf("\n=== Test %d: echo ===\n", test_case++);
    ws_send_text(fd, "hello dpdk websocket!");
    uint8_t buf[4096]; int opcode, len = sizeof(buf);
    int n = ws_recv_frame(fd, &opcode, buf, &len);
    if (n > 0) {
        buf[len] = '\0';
        printf("[OK] echo reply: %s\n", buf);
    } else {
        printf("[FAIL] no response\n");
    }
}

static void test_subscribe(int fd) {
    printf("\n=== Test %d: subscribe ===\n", test_case++);
    ws_send_text(fd, "{\"action\":\"subscribe\",\"topic\":\"dpdk_test\"}");

    /* First response: subscription confirmation */
    uint8_t buf[4096]; int opcode, len = sizeof(buf);
    int n = ws_recv_frame(fd, &opcode, buf, &len);
    if (n > 0) {
        buf[len] = '\0';
        printf("[OK] subscribe response: %s\n", buf);
    }

    /* Second response: push data with welcome message */
    len = sizeof(buf);
    n = ws_recv_frame(fd, &opcode, buf, &len);
    if (n > 0) {
        buf[len] = '\0';
        printf("[OK] push data: %s\n", buf);
    }
}

static void test_periodic_push(int fd, int wait_sec) {
    printf("\n=== Test %d: periodic push (wait %ds) ===\n", test_case++, wait_sec);
    sleep(wait_sec);
    /* Check for push data (non-blocking poll) */
    /* For simplicity, just wait and try one recv with timeout */
    struct timeval tv = { .tv_sec = 1, .tv_usec = 0 };
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    uint8_t buf[4096]; int opcode, len = sizeof(buf);
    int n = ws_recv_frame(fd, &opcode, buf, &len);
    if (n > 0) {
        buf[len] = '\0';
        printf("[OK] periodic push: %s\n", buf);
    } else {
        printf("[INFO] no periodic data received yet\n");
    }
}

int main(int argc, char *argv[]) {
    const char *server_ip = argc > 1 ? argv[1] : "192.168.159.129";
    int port = argc > 2 ? atoi(argv[2]) : 8080;

    printf("╔══════════════════════════════════════════════╗\n");
    printf("║   WebSocket Protocol Unit Test              ║\n");
    printf("╠══════════════════════════════════════════════╣\n");
    printf("║ Server: %s:%d                            \n", server_ip, port);
    printf("╚══════════════════════════════════════════════╝\n\n");

    /* 1. TCP connect */
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) { perror("socket"); return 1; }

    /* Set TCP_NODELAY for lower latency */
    int opt = 1;
    setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port   = htons(port),
    };
    inet_pton(AF_INET, server_ip, &addr.sin_addr);

    printf("[TCP] Connecting to %s:%d...\n", server_ip, port);
    if (connect(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("connect");
        close(fd);
        return 1;
    }
    printf("[TCP] Connected!\n");

    /* 2. WebSocket handshake */
    if (ws_handshake(fd, server_ip, port, "/ws") < 0) {
        printf("[FAIL] WebSocket handshake failed\n");
        close(fd);
        return 1;
    }

    /* 3. Run tests */
    test_send_recv(fd);
    test_subscribe(fd);
    test_periodic_push(fd, 5);

    printf("\n╔══════════════════════════════════════════════╗\n");
    printf("║   All tests passed!                         ║\n");
    printf("╚══════════════════════════════════════════════╝\n");

    close(fd);
    return 0;
}