#include "mvgal_transport.h"

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

/* ── Constants ───────────────────────────────────────────────────────── */

#define MAX_EVENTS               256
#define PARTIAL_READ_BUF         4096
#define ACCEPT_BACKLOG           128
#define INVALID_FD               (-1)
#define MAX_NODE_ID_STR          32   /* enough for uint64_t hex */

/* ── Per-connection state ────────────────────────────────────────────── */

typedef struct connection {
    int            fd;
    uint64_t       node_id;
    uint8_t       *recv_buf;
    size_t         recv_buf_size;
    size_t         recv_offset;
    mvgal_net_peer_state_t state;
    uint64_t       last_heartbeat_ms;
    bool           marked_for_close;
    struct connection *next;       /* hash chain */
} connection_t;

/* ── Transport state ─────────────────────────────────────────────────── */

static struct {
    mvgal_transport_config_t  config;
    mvgal_transport_stats_t   stats;

    int    server_fd;
    int    epoll_fd;
    bool   running;
    bool   initialized;

    /* Connection table: simple hash by node_id (mod CONN_HASH_SIZE) */
#define CONN_HASH_SIZE  128
    connection_t *conns[CONN_HASH_SIZE];
    uint32_t      conn_count;

    /* Callbacks */
    mvgal_transport_msg_callback  msg_cb;
    void                         *msg_cb_userdata;
    mvgal_transport_peer_callback peer_cb;
    void                         *peer_cb_userdata;

    /* Heartbeat thread */
    pthread_t  heartbeat_thread;
    pthread_t  listener_thread;
    pthread_mutex_t lock;
} g_transport = {0};

/* ── Helpers ─────────────────────────────────────────────────────────── */

static uint32_t conn_hash(uint64_t node_id) {
    return (uint32_t)(node_id ^ (node_id >> 32)) % CONN_HASH_SIZE;
}

static int set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    return (flags == -1) ? -1 : fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

static int set_socket_opts(int fd) {
    int yes = 1;
    if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes)) < 0)
        return -1;
    struct timeval tv = { .tv_sec = 0, .tv_usec = 500000 };
    if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0)
        return -1;
    return 0;
}

static bool bind_addr_is_ipv4(const char *addr)
{
    struct in_addr tmp;

    return addr != NULL && inet_pton(AF_INET, addr, &tmp) == 1;
}

/* ── Connection table operations (caller must hold lock) ─────────────── */

static connection_t *conn_find_unlocked(uint64_t node_id) {
    connection_t *c = g_transport.conns[conn_hash(node_id)];
    while (c) {
        if (c->node_id == node_id) return c;
        c = c->next;
    }
    return NULL;
}

static int conn_insert_unlocked(connection_t *c) {
    if (g_transport.conn_count >= g_transport.config.max_peers) {
        g_transport.stats.connection_errors++;
        return -1;
    }
    uint32_t h = conn_hash(c->node_id);
    c->next = g_transport.conns[h];
    g_transport.conns[h] = c;
    g_transport.conn_count++;
    g_transport.stats.total_connections++;
    g_transport.stats.active_connections++;
    return 0;
}

static void conn_remove_unlocked(uint64_t node_id) {
    uint32_t h = conn_hash(node_id);
    connection_t **pp = &g_transport.conns[h];
    while (*pp) {
        if ((*pp)->node_id == node_id) {
            connection_t *victim = *pp;
            *pp = victim->next;

            if (victim->fd != INVALID_FD) {
                epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_DEL, victim->fd, NULL);
                close(victim->fd);
            }
            free(victim->recv_buf);
            free(victim);

            g_transport.conn_count--;
            if (g_transport.conn_count > 0)
                g_transport.stats.active_connections--;
            else
                g_transport.stats.active_connections = 0;
            return;
        }
        pp = &(*pp)->next;
    }
}

/* ── Create connection object ────────────────────────────────────────── */

static connection_t *conn_create(int fd, uint64_t node_id) {
    connection_t *c = calloc(1, sizeof(connection_t));
    if (!c) return NULL;

    c->fd      = fd;
    c->node_id = node_id;
    c->state   = MVGAL_NET_PEER_CONNECTING;
    c->recv_buf_size = g_transport.config.recv_buffer_size;
    c->recv_buf = malloc(c->recv_buf_size);
    if (!c->recv_buf) {
        free(c);
        return NULL;
    }
    return c;
}

/* ── Read exact bytes from socket ────────────────────────────────────── */

/* ── Write all bytes to socket ───────────────────────────────────────── */

static int write_all(int fd, const void *buf, size_t len) {
    const uint8_t *p = (const uint8_t *)buf;
    while (len > 0) {
        ssize_t n = write(fd, p, len);
        if (n <= 0) return -1;
        p    += n;
        len  -= (size_t)n;
    }
    return 0;
}

/* ── Process one incoming message for a connection ───────────────────── */

static void process_message(connection_t *c, mvgal_net_message_header_t *hdr,
                             void *payload)
{
    g_transport.stats.messages_received++;
    g_transport.stats.bytes_received += sizeof(*hdr) + hdr->payload_size;

    /* Handle heartbeat reply */
    if (hdr->msg_type == MVGAL_NET_MSG_HEARTBEAT_ACK) {
        /* Record heartbeat timestamp — updated via recv timestamps elsewhere */
        return;
    }

    /* Deliver to callback */
    if (g_transport.msg_cb) {
        g_transport.msg_cb(c->node_id, hdr, payload, g_transport.msg_cb_userdata);
    }
}

/* ── Drain available data from a connection ──────────────────────────── */

static int drain_connection(connection_t *c) {
    if (!c || c->fd == INVALID_FD) return -1;

    /* Read as much as available into recv_buf */
    for (;;) {
        size_t space = c->recv_buf_size - c->recv_offset;
        if (space < PARTIAL_READ_BUF) {
            /* Grow buffer */
            size_t new_size = c->recv_buf_size * 2;
            uint8_t *new_buf = realloc(c->recv_buf, new_size);
            if (!new_buf) break;
            c->recv_buf      = new_buf;
            c->recv_buf_size = new_size;
            space = c->recv_buf_size - c->recv_offset;
        }

        ssize_t n = read(c->fd, c->recv_buf + c->recv_offset, space);
        if (n > 0) {
            c->recv_offset += (size_t)n;
            continue;
        }
        if (n == 0) {
            /* Peer disconnected */
            return -1;
        }
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            break; /* no more data */
        }
        return -1;
    }

    /* Parse complete messages from buffer */
    size_t consumed = 0;
    while (consumed + sizeof(mvgal_net_message_header_t) <= c->recv_offset) {
        mvgal_net_message_header_t *hdr =
            (mvgal_net_message_header_t *)(c->recv_buf + consumed);
        size_t msg_total = sizeof(*hdr) + hdr->payload_size;

        if (consumed + msg_total > c->recv_offset) break; /* incomplete */

        void *payload = (msg_total > sizeof(*hdr))
                        ? c->recv_buf + consumed + sizeof(*hdr)
                        : NULL;

        process_message(c, hdr, payload);
        consumed += msg_total;
    }

    /* Shift remaining data to front */
    if (consumed > 0 && consumed < c->recv_offset) {
        memmove(c->recv_buf, c->recv_buf + consumed,
                c->recv_offset - consumed);
    }
    c->recv_offset -= consumed;
    return 0;
}

/* ── Listener thread ─────────────────────────────────────────────────── */

static void *listener_thread_fn(void *arg) {
    (void)arg;

    while (g_transport.running) {
        struct sockaddr_in6 cli_addr;
        socklen_t cli_len = sizeof(cli_addr);
        int cli_fd = accept(g_transport.server_fd,
                            (struct sockaddr *)&cli_addr, &cli_len);
        if (cli_fd < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                usleep(10000); /* 10 ms */
                continue;
            }
            if (!g_transport.running) break;
            continue;
        }

        set_nonblocking(cli_fd);
        set_socket_opts(cli_fd);

        /* We don't know node_id yet — will be learned via discovery */
        /* For now create a connection with node_id=0 — updated on first msg */
        pthread_mutex_lock(&g_transport.lock);
        connection_t *c = conn_create(cli_fd, 0);
        if (!c) {
            close(cli_fd);
            pthread_mutex_unlock(&g_transport.lock);
            continue;
        }
        c->state = MVGAL_NET_PEER_CONNECTED;

        if (conn_insert_unlocked(c) != 0) {
            free(c->recv_buf);
            free(c);
            close(cli_fd);
            pthread_mutex_unlock(&g_transport.lock);
            continue;
        }

        struct epoll_event ev = { .events = EPOLLIN | EPOLLRDHUP | EPOLLET,
                                   .data.ptr = c };
        epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_ADD, cli_fd, &ev);
        pthread_mutex_unlock(&g_transport.lock);

        if (g_transport.peer_cb) {
            g_transport.peer_cb(0, MVGAL_NET_PEER_CONNECTED,
                                g_transport.peer_cb_userdata);
        }
    }
    return NULL;
}

/* ── Heartbeat thread ────────────────────────────────────────────────── */

static void *heartbeat_thread_fn(void *arg) {
    (void)arg;
    uint32_t interval_ms = g_transport.config.heartbeat_interval_ms;

    mvgal_net_message_header_t hdr = {0};
    hdr.magic        = MVGAL_NET_MAGIC;
    hdr.version      = MVGAL_NET_VERSION;
    hdr.msg_type     = MVGAL_NET_MSG_HEARTBEAT;
    hdr.payload_size = 0;
    hdr.checksum     = 0;
    hdr.flags        = 0;

    while (g_transport.running) {
        usleep(interval_ms * 1000);

        pthread_mutex_lock(&g_transport.lock);
        if (!g_transport.running) {
            pthread_mutex_unlock(&g_transport.lock);
            break;
        }

        /* Send heartbeat to all connected peers */
        for (uint32_t i = 0; i < CONN_HASH_SIZE; i++) {
            connection_t *c = g_transport.conns[i];
            while (c) {
                if (c->fd != INVALID_FD) {
                    hdr.sequence_num++;
                    if (write_all(c->fd, &hdr, sizeof(hdr)) < 0) {
                        c->marked_for_close = true;
                        g_transport.stats.heartbeat_timeouts++;
                    } else {
                        g_transport.stats.messages_sent++;
                        g_transport.stats.bytes_sent += sizeof(hdr);
                    }
                }
                c = c->next;
            }
        }

        /* Clean up marked connections */
        for (uint32_t i = 0; i < CONN_HASH_SIZE; i++) {
            connection_t *c = g_transport.conns[i];
            while (c) {
                if (c->marked_for_close) {
                    uint64_t dead_id = c->node_id;
                    conn_remove_unlocked(dead_id);
                    if (g_transport.peer_cb) {
                        g_transport.peer_cb(dead_id, MVGAL_NET_PEER_TIMEOUT,
                                            g_transport.peer_cb_userdata);
                    }
                    break;
                }
                c = c->next;
            }
        }

        pthread_mutex_unlock(&g_transport.lock);
    }
    return NULL;
}

/* ── Public API ──────────────────────────────────────────────────────── */

int mvgal_transport_init(const mvgal_transport_config_t *config) {
    if (g_transport.initialized) return 0; /* Already initialized */

    memset(&g_transport, 0, sizeof(g_transport));

    if (config) {
        g_transport.config = *config;
    } else {
        g_transport.config = (mvgal_transport_config_t)MVGAL_TRANSPORT_DEFAULT_CONFIG;
    }

    /* Default values */
    if (g_transport.config.heartbeat_interval_ms == 0)
        g_transport.config.heartbeat_interval_ms = 5000;
    if (g_transport.config.timeout_ms == 0)
        g_transport.config.timeout_ms = 30000;
    if (g_transport.config.max_peers == 0)
        g_transport.config.max_peers = 64;
    if (g_transport.config.recv_buffer_size == 0)
        g_transport.config.recv_buffer_size = 65536;

    g_transport.server_fd = INVALID_FD;
    g_transport.epoll_fd  = epoll_create1(0);
    if (g_transport.epoll_fd < 0) return -1;

    pthread_mutex_init(&g_transport.lock, NULL);
    g_transport.initialized = true;
    return 0;
}

void mvgal_transport_shutdown(void) {
    if (!g_transport.initialized) return;

    g_transport.running = false;

    /* Stop listener */
    if (g_transport.server_fd != INVALID_FD) {
        close(g_transport.server_fd);
        g_transport.server_fd = INVALID_FD;
    }

    /* Wait for threads */
    if (g_transport.listener_thread)
        pthread_join(g_transport.listener_thread, NULL);
    if (g_transport.heartbeat_thread)
        pthread_join(g_transport.heartbeat_thread, NULL);

    /* Disconnect all */
    pthread_mutex_lock(&g_transport.lock);
    for (uint32_t i = 0; i < CONN_HASH_SIZE; i++) {
        connection_t *c = g_transport.conns[i];
        while (c) {
            connection_t *next = c->next;
            if (c->fd != INVALID_FD) {
                epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_DEL, c->fd, NULL);
                close(c->fd);
            }
            free(c->recv_buf);
            free(c);
            c = next;
        }
        g_transport.conns[i] = NULL;
    }
    g_transport.conn_count = 0;
    g_transport.stats.active_connections = 0;
    pthread_mutex_unlock(&g_transport.lock);

    if (g_transport.epoll_fd != INVALID_FD) {
        close(g_transport.epoll_fd);
    }

    pthread_mutex_destroy(&g_transport.lock);
    memset(&g_transport, 0, sizeof(g_transport));
}

int mvgal_transport_listen(void) {
    if (!g_transport.initialized) return -1;
    if (g_transport.server_fd != INVALID_FD) return 0; /* Already listening */

    bool use_ipv4 = bind_addr_is_ipv4(g_transport.config.bind_addr);

    /* Create TCP socket matching the requested bind address family. */
    g_transport.server_fd = socket(use_ipv4 ? AF_INET : AF_INET6,
                                   SOCK_STREAM | SOCK_NONBLOCK, 0);
    if (g_transport.server_fd < 0) return -1;

    int yes = 1;
    setsockopt(g_transport.server_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

    if (use_ipv4) {
        struct sockaddr_in addr = {0};
        addr.sin_family = AF_INET;
        addr.sin_port   = htons(g_transport.config.tcp_port);
        if (inet_pton(AF_INET, g_transport.config.bind_addr, &addr.sin_addr) != 1) {
            close(g_transport.server_fd);
            g_transport.server_fd = INVALID_FD;
            return -1;
        }

        if (bind(g_transport.server_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
            close(g_transport.server_fd);
            g_transport.server_fd = INVALID_FD;
            return -1;
        }
    } else {
        int no = 0;
        setsockopt(g_transport.server_fd, IPPROTO_IPV6, IPV6_V6ONLY, &no, sizeof(no));

        struct sockaddr_in6 addr = {0};
        addr.sin6_family = AF_INET6;
        addr.sin6_port   = htons(g_transport.config.tcp_port);
        if (g_transport.config.bind_addr) {
            if (inet_pton(AF_INET6, g_transport.config.bind_addr, &addr.sin6_addr) != 1) {
                close(g_transport.server_fd);
                g_transport.server_fd = INVALID_FD;
                return -1;
            }
        } else {
            addr.sin6_addr = in6addr_any;
        }

        if (bind(g_transport.server_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
            close(g_transport.server_fd);
            g_transport.server_fd = INVALID_FD;
            return -1;
        }
    }

    if (listen(g_transport.server_fd, ACCEPT_BACKLOG) < 0) {
        close(g_transport.server_fd);
        g_transport.server_fd = INVALID_FD;
        return -1;
    }

    /* Get actual port if ephemeral */
    if (g_transport.config.tcp_port == 0) {
        if (use_ipv4) {
            struct sockaddr_in bound;
            socklen_t bound_len = sizeof(bound);
            if (getsockname(g_transport.server_fd, (struct sockaddr *)&bound, &bound_len) == 0)
                g_transport.config.tcp_port = ntohs(bound.sin_port);
        } else {
            struct sockaddr_in6 bound;
            socklen_t bound_len = sizeof(bound);
            if (getsockname(g_transport.server_fd, (struct sockaddr *)&bound, &bound_len) == 0)
                g_transport.config.tcp_port = ntohs(bound.sin6_port);
        }
    }

    g_transport.running = true;

    if (pthread_create(&g_transport.listener_thread, NULL,
                       listener_thread_fn, NULL) != 0)
    {
        close(g_transport.server_fd);
        g_transport.server_fd = INVALID_FD;
        g_transport.running = false;
        return -1;
    }

    if (g_transport.config.enable_heartbeat) {
        if (pthread_create(&g_transport.heartbeat_thread, NULL,
                           heartbeat_thread_fn, NULL) != 0)
        {
            g_transport.running = false;
            close(g_transport.server_fd);
            g_transport.server_fd = INVALID_FD;
            return -1;
        }
    }

    return 0;
}

int mvgal_transport_stop(void) {
    g_transport.running = false;
    if (g_transport.server_fd != INVALID_FD) {
        close(g_transport.server_fd);
        g_transport.server_fd = INVALID_FD;
    }
    return 0;
}

int mvgal_transport_connect(const char *addr, uint16_t port, uint64_t node_id) {
    if (!g_transport.initialized || !addr) return -1;

    pthread_mutex_lock(&g_transport.lock);

    if (conn_find_unlocked(node_id)) {
        pthread_mutex_unlock(&g_transport.lock);
        return -1; /* already connected */
    }

    int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
    if (fd < 0) {
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }

    set_socket_opts(fd);

    struct sockaddr_in srv = {0};
    srv.sin_family = AF_INET;
    srv.sin_port   = htons(port);
    if (inet_pton(AF_INET, addr, &srv.sin_addr) <= 0) {
        close(fd);
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }

    int rc = connect(fd, (struct sockaddr *)&srv, sizeof(srv));
    if (rc < 0 && errno != EINPROGRESS) {
        close(fd);
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }

    connection_t *c = conn_create(fd, node_id);
    if (!c) {
        close(fd);
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }
    c->state = MVGAL_NET_PEER_CONNECTING;

    if (conn_insert_unlocked(c) != 0) {
        free(c->recv_buf);
        free(c);
        close(fd);
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }

    struct epoll_event ev = { .events = EPOLLIN | EPOLLOUT | EPOLLRDHUP | EPOLLET,
                               .data.ptr = c };
    epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_ADD, fd, &ev);
    pthread_mutex_unlock(&g_transport.lock);
    return 0;
}

int mvgal_transport_disconnect(uint64_t node_id) {
    pthread_mutex_lock(&g_transport.lock);
    connection_t *c = conn_find_unlocked(node_id);
    if (!c) {
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }
    uint64_t dead_id = c->node_id;
    conn_remove_unlocked(node_id);
    pthread_mutex_unlock(&g_transport.lock);

    if (g_transport.peer_cb) {
        g_transport.peer_cb(dead_id, MVGAL_NET_PEER_DISCONNECTED,
                            g_transport.peer_cb_userdata);
    }
    return 0;
}

int mvgal_transport_disconnect_all(void) {
    pthread_mutex_lock(&g_transport.lock);
    for (uint32_t i = 0; i < CONN_HASH_SIZE; i++) {
        connection_t *c = g_transport.conns[i];
        while (c) {
            connection_t *next = c->next;
            if (c->fd != INVALID_FD) {
                epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_DEL, c->fd, NULL);
                close(c->fd);
            }
            free(c->recv_buf);
            free(c);
            c = next;
        }
        g_transport.conns[i] = NULL;
    }
    g_transport.conn_count = 0;
    g_transport.stats.active_connections = 0;
    pthread_mutex_unlock(&g_transport.lock);
    return 0;
}

int mvgal_transport_send(uint64_t node_id,
                          const mvgal_net_message_header_t *hdr,
                          const void *payload,
                          size_t payload_size)
{
    if (!g_transport.initialized || !hdr) return -1;

    pthread_mutex_lock(&g_transport.lock);
    connection_t *c = conn_find_unlocked(node_id);
    if (!c || c->fd == INVALID_FD) {
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }

    if (write_all(c->fd, hdr, sizeof(*hdr)) < 0) {
        c->marked_for_close = true;
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }
    if (payload && payload_size > 0) {
        if (write_all(c->fd, payload, payload_size) < 0) {
            c->marked_for_close = true;
            pthread_mutex_unlock(&g_transport.lock);
            return -1;
        }
    }

    g_transport.stats.messages_sent++;
    g_transport.stats.bytes_sent += sizeof(*hdr) + payload_size;
    pthread_mutex_unlock(&g_transport.lock);
    return 0;
}

int mvgal_transport_broadcast(const mvgal_net_message_header_t *hdr,
                               const void *payload,
                               size_t payload_size)
{
    if (!g_transport.initialized || !hdr) return -1;
    int failures = 0;

    pthread_mutex_lock(&g_transport.lock);
    for (uint32_t i = 0; i < CONN_HASH_SIZE; i++) {
        connection_t *c = g_transport.conns[i];
        while (c) {
            if (c->fd != INVALID_FD && !c->marked_for_close) {
                if (write_all(c->fd, hdr, sizeof(*hdr)) < 0 ||
                    (payload && payload_size > 0 &&
                     write_all(c->fd, payload, payload_size) < 0))
                {
                    c->marked_for_close = true;
                    failures++;
                } else {
                    g_transport.stats.messages_sent++;
                    g_transport.stats.bytes_sent += sizeof(*hdr) + payload_size;
                }
            }
            c = c->next;
        }
    }
    pthread_mutex_unlock(&g_transport.lock);
    return failures;
}

int mvgal_transport_poll(uint32_t timeout_ms) {
    if (!g_transport.initialized || g_transport.epoll_fd == INVALID_FD)
        return -1;

    struct epoll_event events[MAX_EVENTS];
    int nfds = epoll_wait(g_transport.epoll_fd, events, MAX_EVENTS,
                          (int)timeout_ms);
    if (nfds <= 0) return nfds;

    for (int i = 0; i < nfds; i++) {
        connection_t *c = (connection_t *)events[i].data.ptr;
        if (!c) continue;

        if (events[i].events & (EPOLLERR | EPOLLHUP | EPOLLRDHUP)) {
            pthread_mutex_lock(&g_transport.lock);
            uint64_t dead_id = c->node_id;
            conn_remove_unlocked(c->node_id);
            pthread_mutex_unlock(&g_transport.lock);

            if (g_transport.peer_cb) {
                g_transport.peer_cb(dead_id, MVGAL_NET_PEER_DISCONNECTED,
                                    g_transport.peer_cb_userdata);
            }
            continue;
        }

        if (events[i].events & EPOLLOUT) {
            /* Connection established only if SO_ERROR reports success. */
            int so_error = 0;
            socklen_t so_error_len = sizeof(so_error);
            if (getsockopt(c->fd, SOL_SOCKET, SO_ERROR,
                           &so_error, &so_error_len) < 0 || so_error != 0) {
                pthread_mutex_lock(&g_transport.lock);
                uint64_t dead_id = c->node_id;
                conn_remove_unlocked(c->node_id);
                g_transport.stats.connection_errors++;
                pthread_mutex_unlock(&g_transport.lock);

                if (g_transport.peer_cb) {
                    g_transport.peer_cb(dead_id, MVGAL_NET_PEER_DISCONNECTED,
                                        g_transport.peer_cb_userdata);
                }
                continue;
            }

            pthread_mutex_lock(&g_transport.lock);
            c->state = MVGAL_NET_PEER_CONNECTED;
            /* Stop polling for OUT unless we need to write */
            struct epoll_event ev = { .events = EPOLLIN | EPOLLRDHUP | EPOLLET,
                                       .data.ptr = c };
            epoll_ctl(g_transport.epoll_fd, EPOLL_CTL_MOD, c->fd, &ev);
            pthread_mutex_unlock(&g_transport.lock);

            if (g_transport.peer_cb) {
                g_transport.peer_cb(c->node_id, MVGAL_NET_PEER_CONNECTED,
                                    g_transport.peer_cb_userdata);
            }
        }

        if (events[i].events & EPOLLIN) {
            pthread_mutex_lock(&g_transport.lock);
            int rc = drain_connection(c);
            pthread_mutex_unlock(&g_transport.lock);
            if (rc < 0) {
                pthread_mutex_lock(&g_transport.lock);
                uint64_t dead_id = c->node_id;
                conn_remove_unlocked(c->node_id);
                pthread_mutex_unlock(&g_transport.lock);

                if (g_transport.peer_cb) {
                    g_transport.peer_cb(dead_id, MVGAL_NET_PEER_DISCONNECTED,
                                        g_transport.peer_cb_userdata);
                }
            }
        }
    }
    return nfds;
}

bool mvgal_transport_is_connected(uint64_t node_id) {
    pthread_mutex_lock(&g_transport.lock);
    connection_t *c = conn_find_unlocked(node_id);
    bool ok = (c && c->fd != INVALID_FD && !c->marked_for_close);
    pthread_mutex_unlock(&g_transport.lock);
    return ok;
}

int mvgal_transport_get_peer(uint64_t node_id, mvgal_net_peer_t *peer) {
    if (!peer) return -1;

    pthread_mutex_lock(&g_transport.lock);
    connection_t *c = conn_find_unlocked(node_id);
    if (!c) {
        pthread_mutex_unlock(&g_transport.lock);
        return -1;
    }
    peer->node_id   = c->node_id;
    peer->peer_state = c->state;
    peer->last_seen_ns = 0; /* would need timestamp tracking */
    pthread_mutex_unlock(&g_transport.lock);
    return 0;
}

int mvgal_transport_get_peers(mvgal_net_peer_t *peers, uint32_t *count) {
    if (!peers || !count) return -1;
    uint32_t max = *count;
    uint32_t idx = 0;

    pthread_mutex_lock(&g_transport.lock);
    for (uint32_t i = 0; i < CONN_HASH_SIZE && idx < max; i++) {
        connection_t *c = g_transport.conns[i];
        while (c && idx < max) {
            peers[idx].node_id   = c->node_id;
            peers[idx].peer_state = c->state;
            peers[idx].last_seen_ns = 0;
            idx++;
            c = c->next;
        }
    }
    pthread_mutex_unlock(&g_transport.lock);

    *count = idx;
    return 0;
}

void mvgal_transport_get_stats(mvgal_transport_stats_t *stats) {
    if (!stats) return;
    pthread_mutex_lock(&g_transport.lock);
    *stats = g_transport.stats;
    pthread_mutex_unlock(&g_transport.lock);
}

void mvgal_transport_set_msg_callback(mvgal_transport_msg_callback cb, void *userdata) {
    pthread_mutex_lock(&g_transport.lock);
    g_transport.msg_cb          = cb;
    g_transport.msg_cb_userdata = userdata;
    pthread_mutex_unlock(&g_transport.lock);
}

void mvgal_transport_set_peer_callback(mvgal_transport_peer_callback cb, void *userdata) {
    pthread_mutex_lock(&g_transport.lock);
    g_transport.peer_cb          = cb;
    g_transport.peer_cb_userdata = userdata;
    pthread_mutex_unlock(&g_transport.lock);
}
