/**
 * tcp_socket.c — POSIX-like socket API (production-grade)
 *
 *   Condition-variable based blocking (no polling sleep!)
 *   Byte-stream send with MSS segmentation
 *   Full close sequence for ESTABLISHED / CLOSE_WAIT
 */

#include "tcp_core.h"

/* ── Local fd tracking ──────────────────────────────────────────── */

static unsigned char fd_bitmap[MAX_FD_COUNT] = {0};
static struct localhost *lhost = NULL;

static int alloc_fd(void)
{
    for (int i = DEFAULT_FD_NUM; i < MAX_FD_COUNT; i++)
    {
        if (!(fd_bitmap[i / 8] & (1 << (i % 8))))
        {
            fd_bitmap[i / 8] |= (1 << (i % 8));
            return i;
        }
    }
    return -1;
}

static void free_fd(int fd)
{
    if (fd >= 0 && fd < MAX_FD_COUNT)
        fd_bitmap[fd / 8] &= ~(1 << (fd % 8));
}

void *get_hostinfo_fromfd(int sockfd)
{
    /* Try hash table (TCP) */
    struct ng_tcp_table *table = tcpInstance();
    for (int b = 0; b < TCP_HTABLE_SIZE; b++)
    {
        struct ng_tcp_stream *s;
        for (s = table->htable[b]; s; s = s->ht_next)
        {
            if (s->fd == sockfd)
                return s;
        }
    }
    /* Try UDP localhosts */
    struct localhost *h;
    for (h = lhost; h; h = h->next)
    {
        if (h->fd == sockfd)
            return h;
    }
    return NULL;
}

/* Helper to dequeue and deliver data to user buffer */
static ssize_t deliver_data(struct ng_tcp_stream *s, struct ng_tcp_fragment *f,
                            void *buf, size_t buflen)
{
    if (f->length == 0)
    {
        frag_free(f);
        return 0;
    }

    ssize_t ret;
    if (f->length > (int)buflen)
    {
        /* Partial read — advance offset, no memmove */
        rte_memcpy(buf, f->data + f->read_offset, buflen);
        f->read_offset += buflen;
        f->length -= buflen;
        ret = buflen;
        rte_ring_mp_enqueue(s->rcvbuf, f); /* put back remainder */
    }
    else
    {
        /* Full consumption — return to pool */
        rte_memcpy(buf, f->data + f->read_offset, f->length);
        ret = f->length;
        frag_free(f);
    }
    return ret;
}

/* ══════════════════════════════════════════════════════════════════
 *  nsocket
 * ══════════════════════════════════════════════════════════════════ */

int nsocket(int domain, int type, int protocol)
{
    (void)domain;
    (void)protocol;
    int fd = alloc_fd();
    if (fd < 0)
        return -1;

    if (type == SOCK_STREAM)
    {
        struct ng_tcp_stream *s = rte_malloc("socket_stream",
                                             sizeof(struct ng_tcp_stream), 0);
        if (!s)
        {
            free_fd(fd);
            return -1;
        }
        memset(s, 0, sizeof(*s));

        s->fd = fd;
        s->protocol = IPPROTO_TCP;
        s->status = NG_TCP_STATUS_CLOSED;

        s->rcvbuf = rte_ring_create("tcp_rcv", RING_SIZE,
                                    rte_socket_id(), RING_F_SP_ENQ | RING_F_SC_DEQ);
        s->sndbuf = rte_ring_create("tcp_snd", RING_SIZE,
                                    rte_socket_id(), RING_F_SP_ENQ | RING_F_SC_DEQ);
        s->retransmit_ring = rte_ring_create("tcp_rtx", RING_SIZE,
                                             rte_socket_id(), RING_F_SP_ENQ | RING_F_SC_DEQ);

        pthread_cond_init(&s->cond, NULL);
        pthread_mutex_init(&s->mutex, NULL);
        s->wakeup = 0;

        s->snd_buf = rte_malloc("snd_buf", TCP_SENDBUF_SIZE, 0);
        s->rcv_buf_max = TCP_RECVBUF_SIZE;
        s->rcv_buf_avail = TCP_RECVBUF_SIZE;
        s->rcv_wnd = TCP_INITIAL_WINDOW;
        s->mss = TCP_MSS_DEFAULT;
        s->rcv_mss = TCP_MSS_DEFAULT;
        s->cwnd = TCP_INITIAL_CWND;
        s->ssthresh = 0x7FFFFFFF;
        s->rto = TCP_RTO_INITIAL;

        ng_tcp_fd_set(fd, s);
        /* Hash insert deferred to nconnect/nbind when 4-tuple is known */
    }

    return fd;
}

/* ══════════════════════════════════════════════════════════════════
 *  nbind / nlisten / naccept
 * ══════════════════════════════════════════════════════════════════ */

int nbind(int sockfd, const struct sockaddr *addr, socklen_t addrlen)
{
    (void)addrlen;
    struct ng_tcp_stream *s = ng_tcp_fd_get(sockfd);
    if (!s)
        return -1;
    const struct sockaddr_in *in = (const struct sockaddr_in *)addr;
    rte_memcpy(&s->dip, &in->sin_addr.s_addr, sizeof(uint32_t));
    s->dport = in->sin_port;
    rte_memcpy(s->localmac, gSrcMac, RTE_ETHER_ADDR_LEN);
    return 0;
}

int nlisten(int sockfd, int backlog)
{
    (void)backlog;
    struct ng_tcp_stream *s = ng_tcp_fd_get(sockfd);
    if (!s)
        return -1;
    s->status = NG_TCP_STATUS_LISTEN;

    struct ng_tcp_table *table = tcpInstance();
    s->ht_next = table->listeners;
    table->listeners = s;

    printf("[TCP] Listening on port %d\n", ntohs(s->dport));
    return 0;
}

static struct ng_tcp_stream *get_accept_tcb(uint16_t dport)
{
    struct ng_tcp_table *table = tcpInstance();
    int b;
    for (b = 0; b < TCP_HTABLE_SIZE; b++)
    {
        struct ng_tcp_stream *s;
        for (s = table->htable[b]; s; s = s->ht_next)
        {
            if (s->dport == dport && s->fd == -1 &&
                s->status == NG_TCP_STATUS_ESTABLISHED)
                return s;
        }
    }
    return NULL;
}

int naccept(int sockfd, struct sockaddr *addr, socklen_t *addrlen)
{
    struct ng_tcp_stream *listener = ng_tcp_fd_get(sockfd);
    if (!listener || listener->status != NG_TCP_STATUS_LISTEN)
        return -1;

    struct ng_tcp_stream *accepted = NULL;
    while (!accepted)
    {
        accepted = get_accept_tcb(listener->dport);
        if (!accepted)
            usleep(10000);
    }

    int connfd = alloc_fd();
    if (connfd < 0)
        return -1;
    accepted->fd = connfd;
    ng_tcp_fd_set(connfd, accepted);

    struct sockaddr_in *saddr = (struct sockaddr_in *)addr;
    saddr->sin_family = AF_INET;
    saddr->sin_port = accepted->sport;
    rte_memcpy(&saddr->sin_addr.s_addr, &accepted->sip, sizeof(uint32_t));
    if (addrlen)
        *addrlen = sizeof(*saddr);

    printf("[TCP] Accepted connfd=%d\n", connfd);
    return connfd;
}

/* ══════════════════════════════════════════════════════════════════
 *  nconnect — condition variable with timeout
 * ══════════════════════════════════════════════════════════════════ */

int nconnect(int sockfd, const struct sockaddr *addr, socklen_t addrlen)
{
    (void)addrlen;
    struct ng_tcp_stream *s = ng_tcp_fd_get(sockfd);
    if (!s || s->protocol != IPPROTO_TCP)
        return -1;

    const struct sockaddr_in *in = (const struct sockaddr_in *)addr;
    rte_memcpy(&s->dip, &in->sin_addr.s_addr, sizeof(uint32_t));
    s->dport = in->sin_port;
    s->sip = gLocalIp;

    /* Seed rand() with time + pointer for unique port/seq per run */
    {
        unsigned seed = (unsigned)time(NULL) ^ (unsigned)(uintptr_t)s;
        srand(seed);
    }
    s->sport = htons((rand() % 16384) + 49152);
    s->worker_id = ng_compute_worker(s->sip, s->dip, s->sport, s->dport);
    rte_memcpy(s->localmac, gSrcMac, RTE_ETHER_ADDR_LEN);

    /* Switch to stream's worker context for hash table operations */
    int saved_worker = ng_thread_get_worker();
    ng_thread_set_worker(s->worker_id);
    ng_tcp_hash_remove(s);
    ng_tcp_hash_insert(s);
    ng_tcp_stream_activate(s);
    s->status = NG_TCP_STATUS_SYN_SENT;
    ng_thread_set_worker(saved_worker);

    s->iss = rand() % TCP_MAX_SEQ;
    s->snd_una = s->iss;
    s->snd_nxt = s->iss;
    s->snd_max = s->iss;

    printf("[nconnect] Connecting to %s:%d ...\n",
           inet_ntoa(in->sin_addr), ntohs(s->dport));

    struct ng_tcp_fragment *syn = frag_alloc(0);
    if (!syn)
        return -1;
    syn->sport = s->sport;
    syn->dport = s->dport;
    syn->seqnum = s->iss;
    syn->tcp_flags = RTE_TCP_SYN_FLAG;
    syn->windows = (uint16_t)s->rcv_buf_avail;
    syn->hdrlen_off = 0x50;

    printf("[nconnect] Sending SYN (sport=%d, dport=%d, seq=%u)\n",
           ntohs(s->sport), ntohs(s->dport), s->iss);

    rte_ring_mp_enqueue(s->sndbuf, syn);

    /* Busy-poll: service RX/TX through the stream's worker ring */
    struct timespec deadline, now;
    clock_gettime(CLOCK_REALTIME, &deadline);
    deadline.tv_sec += 5;

    struct inout_ring *ring = ringInstanceFor(s->worker_id);
    int spin_count = 0;

    while (s->status != NG_TCP_STATUS_ESTABLISHED)
    {
        if (s->status == NG_TCP_STATUS_CLOSED)
        {
            printf("[nconnect] Connection refused\n");
            return -1;
        }

        /* Service TX */
        struct rte_mbuf *tx[BURST_SIZE];
        unsigned nb_tx = rte_ring_sc_dequeue_burst(ring->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]);
        }

        /* Service RX */
        struct rte_mbuf *rx[BURST_SIZE];
        unsigned nb_rx = rte_eth_rx_burst(gDpdkPortId, 0, rx, BURST_SIZE);
        if (nb_rx > 0)
            rte_ring_sp_enqueue_burst(ring->in, (void **)rx, nb_rx, NULL);

        /* Check pkt_process signaled us (non-blocking status check) */
        if (s->status == NG_TCP_STATUS_ESTABLISHED)
            break;

        /* Check deadline every 1000 spins (~10us at modern CPU) */
        if (++spin_count >= 1000) {
            spin_count = 0;
            clock_gettime(CLOCK_REALTIME, &now);
            if (now.tv_sec > deadline.tv_sec ||
                (now.tv_sec == deadline.tv_sec && now.tv_nsec >= deadline.tv_nsec)) {
                s->status = NG_TCP_STATUS_CLOSED;
                printf("[nconnect] Connection timed out\n");
                return -1;
            }
        }

        /* CPU pause hint: reduces power and lets other hyperthread run */
        rte_pause();
    }

    printf("[nconnect] Connected to %s:%d\n",
           inet_ntoa(in->sin_addr), ntohs(s->dport));
    return 0;
}

/* ══════════════════════════════════════════════════════════════════
 *  nsend — byte-stream with MSS segmentation
 * ══════════════════════════════════════════════════════════════════ */

ssize_t nsend(int sockfd, const void *buf, size_t len, int flags)
{
    (void)flags;
    struct ng_tcp_stream *s = ng_tcp_fd_get(sockfd);
    if (!s || s->protocol != IPPROTO_TCP)
        return -1;
    if (s->status != NG_TCP_STATUS_ESTABLISHED &&
        s->status != NG_TCP_STATUS_CLOSE_WAIT)
        return -1;

    ssize_t total = 0;
    const uint8_t *src = (const uint8_t *)buf;

    while (total < (ssize_t)len)
    {
        size_t chunk = len - total;
        if (chunk > s->rcv_mss)
            chunk = s->rcv_mss;

        struct ng_tcp_fragment *f = frag_alloc((int)chunk);
        if (!f)
            break;

        rte_memcpy(f->data, src + total, chunk);
        f->data[chunk] = '\0';

        f->sport = s->sport;
        f->dport = s->dport;
        f->seqnum = s->snd_nxt;
        f->acknum = s->rcv_nxt;
        f->tcp_flags = RTE_TCP_ACK_FLAG | RTE_TCP_PSH_FLAG;
        f->windows = (uint16_t)s->rcv_buf_avail;
        f->hdrlen_off = 0x50;

        rte_ring_mp_enqueue(s->sndbuf, f);
        s->snd_nxt += chunk;
        s->snd_max = s->snd_nxt;
        total += chunk;
    }

    return total;
}

/* ══════════════════════════════════════════════════════════════════
 *  nrecv — condition variable + RX/TX service
 * ══════════════════════════════════════════════════════════════════ */

ssize_t nrecv(int sockfd, void *buf, size_t len, int flags)
{
    (void)flags;
    struct ng_tcp_stream *s = ng_tcp_fd_get(sockfd);
    if (!s || s->protocol != IPPROTO_TCP)
        return -1;

    /* Service RX/TX */
    struct inout_ring *ring = ringInstanceFor(s->worker_id);
    struct rte_mbuf *tx[BURST_SIZE];
    unsigned nb_tx = rte_ring_sc_dequeue_burst(ring->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]);
    }
    struct rte_mbuf *rx[BURST_SIZE];
    unsigned nb_rx = rte_eth_rx_burst(gDpdkPortId, 0, rx, BURST_SIZE);
    if (nb_rx > 0)
        rte_ring_sp_enqueue_burst(ring->in, (void **)rx, nb_rx, NULL);

    /* Try non-blocking dequeue */
    struct ng_tcp_fragment *f = NULL;
    if (rte_ring_mc_dequeue(s->rcvbuf, (void **)&f) >= 0 && f)
    {
        ssize_t r = deliver_data(s, f, buf, len);
        if (r > 0)
        {
            s->rcv_buf_avail += (uint32_t)r;
            if (s->rcv_buf_avail > s->rcv_buf_max)
                s->rcv_buf_avail = s->rcv_buf_max;
        }
        return r;
    }

    /* Busy-poll for data — latency-critical path, same as nconnect */
    while (rte_ring_count(s->rcvbuf) == 0 && s->wakeup == 0)
    {
        if (s->status == NG_TCP_STATUS_CLOSED ||
            s->status == NG_TCP_STATUS_TIME_WAIT)
        {
            return 0;
        }

        /* Service TX */
        struct rte_mbuf *tx_svc[BURST_SIZE];
        unsigned nb_tx_svc = rte_ring_sc_dequeue_burst(ring->out, (void **)tx_svc, BURST_SIZE, NULL);
        if (nb_tx_svc > 0)
        {
            rte_eth_tx_burst(gDpdkPortId, 0, tx_svc, nb_tx_svc);
            for (unsigned i = 0; i < nb_tx_svc; i++)
                rte_pktmbuf_free(tx_svc[i]);
        }

        /* Service RX */
        struct rte_mbuf *rx_svc[BURST_SIZE];
        unsigned nb_rx_svc = rte_eth_rx_burst(gDpdkPortId, 0, rx_svc, BURST_SIZE);
        if (nb_rx_svc > 0) {
            lat_trace_add(1); /* RX data enqueued to ring->in */
            rte_ring_sp_enqueue_burst(ring->in, (void **)rx_svc, nb_rx_svc, NULL);
        }

        /* Check non-blocking */
        if (rte_ring_count(s->rcvbuf) > 0 || s->wakeup != 0) {
            lat_trace_add(4); /* nrecv noticed data */
            break;
        }

        rte_pause();
    }
    s->wakeup = 0;

    if (rte_ring_mc_dequeue(s->rcvbuf, (void **)&f) < 0 || !f)
        return 0;

    ssize_t r = deliver_data(s, f, buf, len);
    if (r > 0)
    {
        s->rcv_buf_avail += (uint32_t)r;
        if (s->rcv_buf_avail > s->rcv_buf_max)
            s->rcv_buf_avail = s->rcv_buf_max;
    }
    return r;
}

/* ══════════════════════════════════════════════════════════════════
 *  nclose — full close sequence
 * ══════════════════════════════════════════════════════════════════ */

int nclose(int fd)
{
    struct ng_tcp_stream *s = ng_tcp_fd_get(fd);
    if (!s)
        return -1;

    free_fd(fd);
    ng_tcp_fd_clear(fd);

    if (s->protocol != IPPROTO_TCP)
        return 0;

    switch (s->status)
    {
    case NG_TCP_STATUS_ESTABLISHED:
        ng_tcp_send_fin(s);
        s->status = NG_TCP_STATUS_FIN_WAIT_1;
        printf("[nclose] Sent FIN → FIN_WAIT_1\n");
        {
            struct timespec deadline, now;
            clock_gettime(CLOCK_REALTIME, &deadline);
            deadline.tv_sec += 5;
            int spin = 0;
            struct inout_ring *ring = ringInstanceFor(s->worker_id);
            while (s->status != NG_TCP_STATUS_CLOSED &&
                   s->status != NG_TCP_STATUS_TIME_WAIT)
            {
                /* Service TX */
                struct rte_mbuf *tx[BURST_SIZE];
                unsigned nb_tx = rte_ring_sc_dequeue_burst(ring->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]);
                }
                /* Service RX */
                struct rte_mbuf *rx[BURST_SIZE];
                unsigned nb_rx = rte_eth_rx_burst(gDpdkPortId, 0, rx, BURST_SIZE);
                if (nb_rx > 0)
                    rte_ring_sp_enqueue_burst(ring->in, (void **)rx, nb_rx, NULL);
                /* Deadline check */
                if (++spin >= 1000) {
                    spin = 0;
                    clock_gettime(CLOCK_REALTIME, &now);
                    if (now.tv_sec > deadline.tv_sec) break;
                }
                rte_pause();
            }
        }
        break;

    case NG_TCP_STATUS_CLOSE_WAIT:
        ng_tcp_send_fin(s);
        s->status = NG_TCP_STATUS_LAST_ACK;
        printf("[nclose] Sent FIN → LAST_ACK\n");
        {
            struct timespec deadline, now;
            clock_gettime(CLOCK_REALTIME, &deadline);
            deadline.tv_sec += 5;
            int spin = 0;
            struct inout_ring *ring = ringInstanceFor(s->worker_id);
            while (s->status != NG_TCP_STATUS_CLOSED)
            {
                /* Service TX */
                struct rte_mbuf *tx[BURST_SIZE];
                unsigned nb_tx = rte_ring_sc_dequeue_burst(ring->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]);
                }
                /* Service RX */
                struct rte_mbuf *rx[BURST_SIZE];
                unsigned nb_rx = rte_eth_rx_burst(gDpdkPortId, 0, rx, BURST_SIZE);
                if (nb_rx > 0)
                    rte_ring_sp_enqueue_burst(ring->in, (void **)rx, nb_rx, NULL);
                /* Deadline check */
                if (++spin >= 1000) {
                    spin = 0;
                    clock_gettime(CLOCK_REALTIME, &now);
                    if (now.tv_sec > deadline.tv_sec) break;
                }
                rte_pause();
            }
        }
        break;

    case NG_TCP_STATUS_SYN_RCVD:
    case NG_TCP_STATUS_SYN_SENT:
        ng_tcp_send_rst(s, s->snd_nxt, 0);
        break;

    default:
        break;
    }

    return 0;
}