/**
 * main.c — DPDK WebSocket Client
 *
 * Architecture:
 *   ┌──────────┐    ┌───────────────┐    ┌──────────────┐    ┌──────────┐
 *   │  DPDK RX │───▶│ pkt_process() │───▶│ ng_tcp_out() │───▶│ DPDK TX  │
 *   │  (lcore0)│    │   (lcore1)    │    │   (lcore1)   │    │ (lcore0) │
 *   └──────────┘    └───────┬───────┘    └──────▲───────┘    └──────────┘
 *                           │                   │
 *                    ┌──────▼───────┐    ┌──────┴───────┐
 *                    │  tcp_socket  │    │  tcp_socket  │
 *                    │  nsend()     │───▶│  sndbuf ring │
 *                    │  nrecv()  ◀──│─── │  rcvbuf ring │
 *                    └──────┬───────┘    └──────────────┘
 *                           │
 *                    ┌──────▼───────┐
 *                    │  ws_client   │
 *                    │  ws_send()   │
 *                    │  ws_recv()   │
 *                    └──────┬───────┘
 *                           │
 *                    ┌──────▼───────┐
 *                    │  main()      │
 *                    │  app logic   │
 *                    └──────────────┘
 *
 * Usage:
 *   ./dpdk_ws_client -l 0-1 -n 2 -- -a <server_ip> -p <port> -w <ws_path>
 *
 * Example:
 *   ./dpdk_ws_client -l 0-1 -n 2 -- -a 192.168.1.100 -p 8080 -w /ws
 */

#include "tcp_core.h"
#include "ws_client.h"

#include <getopt.h>

/* ── Config ──────────────────────────────────────────────────────── */

static struct
{
    char server_ip[64];
    uint16_t server_port;
    char ws_path[256];
    char subscribe_msg[1024];
    int ping_interval; /* seconds, 0 = no ping */
} g_cfg = {
    .server_ip = "192.168.1.100",
    .server_port = 8080,
    .ws_path = "/ws",
    .subscribe_msg = "",
    .ping_interval = 30,
};

/* ── Command-line parsing ────────────────────────────────────────── */

static void print_usage(const char *prog)
{
    printf("Usage: %s [EAL options] -- [APP options]\n\n", prog);
    printf("APP options:\n");
    printf("  -a, --addr <ip>       Server IP address (default: %s)\n",
           g_cfg.server_ip);
    printf("  -p, --port <port>     Server port (default: %d)\n",
           g_cfg.server_port);
    printf("  -w, --path <path>     WebSocket path (default: %s)\n",
           g_cfg.ws_path);
    printf("  -s, --subscribe <msg> Subscribe message (JSON)\n");
    printf("  -k, --keepalive <s>   Ping interval in seconds (0=off, default: %d)\n",
           g_cfg.ping_interval);
    printf("  -l, --local-ip <ip>   Local IP address\n");
    printf("  -h, --help            Show this help\n");
    printf("\nExample:\n");
    printf("  %s -l 0-1 -n 2 -- -a 192.168.1.100 -p 8080 -s '{\"sub\":\"topic1\"}'\n",
           prog);
}

static int parse_args(int argc, char *argv[])
{
    static struct option long_opts[] = {
        {"addr", required_argument, 0, 'a'},
        {"port", required_argument, 0, 'p'},
        {"path", required_argument, 0, 'w'},
        {"subscribe", required_argument, 0, 's'},
        {"keepalive", required_argument, 0, 'k'},
        {"local-ip", required_argument, 0, 'l'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}};

    int opt;
    while ((opt = getopt_long(argc, argv, "a:p:w:s:k:l:h",
                              long_opts, NULL)) != -1)
    {
        switch (opt)
        {
        case 'a':
            strncpy(g_cfg.server_ip, optarg, sizeof(g_cfg.server_ip) - 1);
            break;
        case 'p':
            g_cfg.server_port = (uint16_t)atoi(optarg);
            break;
        case 'w':
            strncpy(g_cfg.ws_path, optarg, sizeof(g_cfg.ws_path) - 1);
            break;
        case 's':
            strncpy(g_cfg.subscribe_msg, optarg,
                    sizeof(g_cfg.subscribe_msg) - 1);
            break;
        case 'k':
            g_cfg.ping_interval = atoi(optarg);
            break;
        case 'l':
        {
            struct in_addr addr;
            if (inet_pton(AF_INET, optarg, &addr) == 1)
            {
                gLocalIp = addr.s_addr;
            }
            break;
        }
        case 'h':
            print_usage(argv[0]);
            return -1;
        default:
            return -1;
        }
    }

    return 0;
}

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

int main(int argc, char *argv[])
{
    setbuf(stdout, NULL); /* disable buffering for real-time output */

    /* ── DPDK EAL init ── */
    int ret = rte_eal_init(argc, argv);
    if (ret < 0)
    {
        rte_exit(EXIT_FAILURE, "EAL init failed\n");
    }
    argc -= ret;
    argv += ret;

    /* ── Parse app args ── */
    if (parse_args(argc, argv) < 0)
    {
        rte_exit(EXIT_FAILURE, "Invalid arguments\n");
    }

    printf("╔══════════════════════════════════════════════╗\n");
    printf("║   DPDK WebSocket Client                     ║\n");
    printf("╠══════════════════════════════════════════════╣\n");
    printf("║ Server:   %s:%d%s\n",
           g_cfg.server_ip, g_cfg.server_port, g_cfg.ws_path);
    struct in_addr lip;
    lip.s_addr = gLocalIp;
    printf("║ Local IP: %s\n", inet_ntoa(lip));
    printf("║ Ping:     %ds\n", g_cfg.ping_interval);
    printf("╚══════════════════════════════════════════════╝\n\n");

    /* ── mbuf pool ── */
    struct rte_mempool *mbuf_pool = rte_pktmbuf_pool_create(
        "mbuf_pool", NUM_MBUFS, 0, 0,
        RTE_MBUF_DEFAULT_BUF_SIZE, rte_socket_id());
    if (mbuf_pool == NULL)
    {
        rte_exit(EXIT_FAILURE, "mbuf pool create failed\n");
    }

    /* ── Init fragment mempools (replaces rte_malloc in hot path) ── */
    ng_frag_pool_init();

    /* ── Init port ── */
    ng_init_port(mbuf_pool);
    rte_eth_macaddr_get(gDpdkPortId, (struct rte_ether_addr *)gSrcMac);

    /* ── Init per-worker ring buffers ── */
    int num_workers = ng_get_num_workers();
    for (int w = 1; w <= num_workers; w++)
    {
        struct inout_ring *r = ringInstanceFor(w);
        char name_in[32], name_out[32];
        snprintf(name_in,  sizeof(name_in),  "in_ring_w%d", w);
        snprintf(name_out, sizeof(name_out), "out_ring_w%d", w);
        r->in  = rte_ring_create(name_in,  RING_SIZE, rte_socket_id(),
                                 RING_F_SP_ENQ | RING_F_SC_DEQ);
        r->out = rte_ring_create(name_out, RING_SIZE, rte_socket_id(),
                                 RING_F_SP_ENQ | RING_F_SC_DEQ);
    }

    /* ── Launch pkt_process workers on lcores 1..N ── */
    static struct worker_args wargs[MAX_WORKERS + 1];
    unsigned lcore_id = rte_lcore_id();
    for (int w = 1; w <= num_workers; w++)
    {
        lcore_id = rte_get_next_lcore(lcore_id, 1, 0);
        wargs[w].mbuf_pool = mbuf_pool;
        wargs[w].worker_id = w;
        printf("[MAIN] Launching worker %d on lcore %u\n", w, lcore_id);
        rte_eal_remote_launch(pkt_process, &wargs[w], lcore_id);
    }

    /* Give workers time to start */
    rte_delay_ms(200);

    /* ══════════════════════════════════════════════════════════
     *  Application: WebSocket Client
     * ══════════════════════════════════════════════════════════ */

    printf("\n[APP] Starting WebSocket client...\n\n");

    /* 1. Create TCP socket */
    int sockfd = nsocket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0)
    {
        rte_exit(EXIT_FAILURE, "nsocket failed\n");
    }
    printf("[APP] Created socket fd=%d\n", sockfd);

    /* 2. Connect + WebSocket handshake */
    ret = ws_connect(sockfd, g_cfg.server_ip, g_cfg.server_port,
                     g_cfg.ws_path, NULL);
    if (ret < 0)
    {
        printf("[APP] ws_connect failed\n");
        nclose(sockfd);
        rte_exit(EXIT_FAILURE, "WebSocket connection failed\n");
    }

    /* 3. Send subscription message */
    if (strlen(g_cfg.subscribe_msg) > 0)
    {
        printf("[APP] Sending subscription: %s\n", g_cfg.subscribe_msg);
        ws_send_text(sockfd, g_cfg.subscribe_msg);
    }

    /* ══════════════════════════════════════════════════════════
     *  Main loop: RX/TX + WS receive
     * ══════════════════════════════════════════════════════════ */

    uint64_t last_ping = rte_get_timer_cycles();
    uint64_t ping_cycles = g_cfg.ping_interval * rte_get_timer_hz();

    printf("[APP] Entering main loop...\n\n");

    int running = 1;
    while (running)
    {
        /* ── DPDK RX (distribute by flow hash to workers) ── */
        {
            struct rte_mbuf *rx[BURST_SIZE];
            unsigned num_recvd = rte_eth_rx_burst(gDpdkPortId, 0, rx, BURST_SIZE);
            for (unsigned i = 0; i < num_recvd; i++)
            {
                struct rte_ether_hdr *ehdr = rte_pktmbuf_mtod(rx[i], struct rte_ether_hdr *);
                int wid = 1; /* default to worker 1 for non-IP */
                if (ehdr->ether_type == rte_cpu_to_be_16(RTE_ETHER_TYPE_IPV4))
                {
                    struct rte_ipv4_hdr *iph = rte_pktmbuf_mtod_offset(rx[i],
                        struct rte_ipv4_hdr *, sizeof(struct rte_ether_hdr));
                    if (iph->next_proto_id == IPPROTO_TCP)
                    {
                        struct rte_tcp_hdr *th = (struct rte_tcp_hdr *)((uint8_t *)iph +
                            ((iph->version_ihl & 0x0F) << 2));
                        wid = ng_compute_worker(iph->src_addr, iph->dst_addr,
                                                th->src_port, th->dst_port);
                    }
                }
                rte_ring_sp_enqueue(ringInstanceFor(wid)->in, rx[i]);
            }
        }

        /* ── DPDK TX (collect from all workers) ── */
        for (int w = 1; w <= num_workers; w++)
        {
            struct rte_mbuf *tx[BURST_SIZE];
            unsigned nb_tx = rte_ring_sc_dequeue_burst(ringInstanceFor(w)->out,
                                                       (void **)tx, BURST_SIZE, NULL);
            if (nb_tx > 0)
            {
                rte_eth_tx_burst(gDpdkPortId, 0, tx, nb_tx);
                for (unsigned i = 0; i < nb_tx; i++)
                    rte_pktmbuf_free(tx[i]);
            }
        }

        /* ── WS receive (non-blocking poll) ── */
        /* Peek: is data available? nrecv blocks, so we check via short timeout */
        /* For simplicity, we do a blocking receive; in production use
         * non-blocking mode or a separate lcore for app logic. */
        uint8_t ws_buf[WS_MAX_FRAME_SIZE];
        int ws_opcode = 0;
        int ws_len = sizeof(ws_buf);

        /* Non-blocking recv: try, if no data available, skip */
        int ws_ret = ws_recv_frame(sockfd, &ws_opcode, ws_buf, &ws_len);
        if (ws_ret > 0 && ws_len > 0)
        {
            ws_buf[ws_len] = '\0';

            switch (ws_opcode)
            {
            case WS_OPCODE_TEXT:
                printf("[WS RECV TEXT] %s\n", ws_buf);
                running = 0; /* exit after first text msg for latency measurement */
                break;
            case WS_OPCODE_BINARY:
                printf("[WS RECV BINARY] %d bytes\n", ws_len);
                /* Hex dump first 32 bytes */
                printf("  ");
                for (int i = 0; i < (ws_len < 32 ? ws_len : 32); i++)
                {
                    printf("%02x ", ws_buf[i]);
                }
                printf("\n");
                break;
            default:
                printf("[WS RECV] opcode=%d, len=%d\n", ws_opcode, ws_len);
                break;
            }
        }
        else if (ws_ret == 0)
        {
            printf("[WS] Connection closed by server\n");
            running = 0;
        }
        else if (ws_ret < 0)
        {
            printf("[WS] Error receiving\n");
            running = 0;
        }

        /* ── Keep-alive ping ── */
        if (g_cfg.ping_interval > 0)
        {
            uint64_t now = rte_get_timer_cycles();
            if (now - last_ping > ping_cycles)
            {
                ws_send_ping(sockfd);
                last_ping = now;
                printf("[WS] PING sent\n");
            }
        }
    }

    /* ── Cleanup ── */
    printf("\n[APP] Shutting down...\n");

    /* Dump latency trace */
    printf("[TRACE] latency handoff points (rdtsc deltas, TSC=%.1f MHz):\n",
           rte_get_tsc_hz() / 1e6);
    int max = g_lat_trace_idx < LAT_TRACE_SIZE ? g_lat_trace_idx : LAT_TRACE_SIZE;
    for (int i = 1; i < max; i++) {
        int a = (i - 1) % LAT_TRACE_SIZE;
        int b = i % LAT_TRACE_SIZE;
        uint64_t delta = g_lat_trace[b].ts - g_lat_trace[a].ts;
        double us = (double)delta / (rte_get_tsc_hz() / 1000000.0);
        const char *evt[] = {"?","rx_enq","pkt_deq","rcvbuf","nrecv_seen"};
        int e1 = g_lat_trace[a].event, e2 = g_lat_trace[b].event;
        if (e1 < 1 || e1 > 4) e1 = 0;
        if (e2 < 1 || e2 > 4) e2 = 0;
        printf("  [%d] %s -> %s: %.1fus\n", i, evt[e1], evt[e2], us);
    }

    nclose(sockfd);

    /* Service TX until FIN and pending ACKs are sent (all workers) */
    for (int flush = 0; flush < 500; flush++) {
        int any = 0;
        for (int w = 1; w <= num_workers; w++) {
            struct rte_mbuf *txf[BURST_SIZE];
            unsigned nb = rte_ring_sc_dequeue_burst(ringInstanceFor(w)->out, (void **)txf, BURST_SIZE, NULL);
            if (nb > 0) {
                rte_eth_tx_burst(gDpdkPortId, 0, txf, nb);
                for (unsigned i = 0; i < nb; i++) rte_pktmbuf_free(txf[i]);
                any = 1;
            }
        }
        /* RX for any remaining packets */
        struct rte_mbuf *rxf[BURST_SIZE];
        unsigned nr = rte_eth_rx_burst(gDpdkPortId, 0, rxf, BURST_SIZE);
        if (nr > 0) {
            for (unsigned i = 0; i < nr; i++)
                rte_ring_sp_enqueue(ringInstanceFor(1)->in, rxf[i]);
            any = 1;
        }
        if (!any) rte_delay_us(200);
    }

    rte_eal_mp_wait_lcore();

    return 0;
}