/**
 * @file mvgal_discovery.c
 * @brief UDP broadcast peer discovery
 *
 * SPDX-License-Identifier: MIT
 */

#include "mvgal_pool.h"
#include <string.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>

/******************************************************************************
 * Discovery state
 ******************************************************************************/

static struct {
    int             sock_fd;
    uint16_t        discovery_port;
    pthread_t       thread;
    volatile bool   running;
} g_discovery;

/******************************************************************************
 * Build a PEER_DISCOVERY message payload
 ******************************************************************************/

static void build_discovery_payload(mvgal_net_peer_discovery_t *payload)
{
    memset(payload, 0, sizeof(*payload));

    /* Generate a simple node ID from hostname hash */
    char hostname[256] = {0};
    if (gethostname(hostname, sizeof(hostname)) != 0)
        snprintf(hostname, sizeof(hostname), "unknown");

    uint64_t hash = 5381;
    for (const char *p = hostname; *p; p++)
        hash = ((hash << 5) + hash) + (uint64_t)(unsigned char)*p;

    payload->node_id         = hash;
    payload->protocol_version = MVGAL_NET_VERSION;

    /* instance_name = "hostname:port" capped at 64 bytes total */
    char name_part[58] = {0};
    strncpy(name_part, hostname, sizeof(name_part) - 1);
    snprintf(payload->instance_name, sizeof(payload->instance_name),
             "%s:%u", name_part, g_discovery.discovery_port);

    payload->num_gpus = 0; /* Filled in by caller if needed */
}

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

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

    struct sockaddr_in bind_addr;
    memset(&bind_addr, 0, sizeof(bind_addr));
    bind_addr.sin_family      = AF_INET;
    bind_addr.sin_port        = htons(g_discovery.discovery_port);
    bind_addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(g_discovery.sock_fd, (struct sockaddr *)&bind_addr,
             sizeof(bind_addr)) < 0) {
        fprintf(stderr, "mvgal_discovery: bind failed: %s\n", strerror(errno));
        return NULL;
    }

    /* Set receive timeout so we can check g_discovery.running periodically */
    struct timeval tv;
    tv.tv_sec  = 0;
    tv.tv_usec = 500000; /* 500ms */
    setsockopt(g_discovery.sock_fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    uint8_t buf[MVGAL_NET_MAX_PAYLOAD];

    while (g_discovery.running) {
        struct sockaddr_in from;
        socklen_t from_len = sizeof(from);
        ssize_t n = recvfrom(g_discovery.sock_fd, buf, sizeof(buf), 0,
                             (struct sockaddr *)&from, &from_len);

        if (n < 0) {
            if (errno == EINTR || errno == EAGAIN)
                continue; /* Timeout — loop and check running flag */
            break;
        }

        /* Need at least a full header + discovery payload */
        if ((size_t)n < sizeof(mvgal_net_message_header_t) + sizeof(mvgal_net_peer_discovery_t))
            continue;

        const mvgal_net_message_header_t *hdr = (const mvgal_net_message_header_t *)buf;

        if (hdr->magic != MVGAL_NET_MAGIC || hdr->version != MVGAL_NET_VERSION)
            continue;

        if (hdr->msg_type != MVGAL_NET_MSG_PEER_DISCOVERY)
            continue;

        const mvgal_net_peer_discovery_t *disc =
            (const mvgal_net_peer_discovery_t *)(hdr + 1);

        /* Skip ourselves (match on node ID) */
        mvgal_net_peer_discovery_t self;
        build_discovery_payload(&self);
        if (disc->node_id == self.node_id)
            continue;

        /* Convert address for transport connect */
        char addr_str[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &from.sin_addr, addr_str, sizeof(addr_str));

        mvgal_pool_internal_t *pool = mvgal_pool_get_instance();

        /* Register peer in pool */
        mvgal_pool_add_peer(disc->node_id, disc->instance_name,
                            disc->num_gpus, disc->protocol_version,
                            (const uint8_t *)&from, from_len);

        /* Initiate TCP data connection */
        uint16_t data_port = pool->config.data_port;
        if (data_port == 0)
            data_port = from.sin_port; /* Use same port as discovery */

        mvgal_transport_connect(addr_str, data_port, disc->node_id);
    }

    return NULL;
}

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

int mvgal_discovery_start(const mvgal_net_config_t *config)
{
    if (g_discovery.running)
        return -1;

    if (!config || !config->enable_discovery)
        return 0; /* Discovery disabled — not an error */

    g_discovery.discovery_port = config->discovery_port;

    g_discovery.sock_fd = socket(AF_INET, SOCK_DGRAM | SOCK_CLOEXEC, 0);
    if (g_discovery.sock_fd < 0)
        return -1;

    /* Allow broadcast */
    int broadcast_enable = 1;
    setsockopt(g_discovery.sock_fd, SOL_SOCKET, SO_BROADCAST,
               &broadcast_enable, sizeof(broadcast_enable));

    /* Reuse address so multiple instances can coexist on same port */
    int reuse = 1;
    setsockopt(g_discovery.sock_fd, SOL_SOCKET, SO_REUSEADDR,
               &reuse, sizeof(reuse));

    g_discovery.running = true;

    if (pthread_create(&g_discovery.thread, NULL, discovery_listener, NULL) != 0) {
        g_discovery.running = false;
        close(g_discovery.sock_fd);
        g_discovery.sock_fd = -1;
        return -1;
    }

    return 0;
}

int mvgal_discovery_stop(void)
{
    if (!g_discovery.running)
        return -1;

    g_discovery.running = false;

    pthread_join(g_discovery.thread, NULL);

    if (g_discovery.sock_fd >= 0) {
        close(g_discovery.sock_fd);
        g_discovery.sock_fd = -1;
    }

    return 0;
}

int mvgal_discovery_announce(void)
{
    if (!g_discovery.running)
        return -1;

    mvgal_net_message_header_t hdr;
    memset(&hdr, 0, sizeof(hdr));
    hdr.magic        = MVGAL_NET_MAGIC;
    hdr.version      = MVGAL_NET_VERSION;
    hdr.msg_type     = MVGAL_NET_MSG_PEER_DISCOVERY;
    hdr.flags        = MVGAL_NET_FLAG_NO_CHECKSUM;

    mvgal_net_peer_discovery_t payload;
    build_discovery_payload(&payload);

    hdr.payload_size = sizeof(payload);

    /* Broadcast on discovery port */
    struct sockaddr_in bc_addr;
    memset(&bc_addr, 0, sizeof(bc_addr));
    bc_addr.sin_family      = AF_INET;
    bc_addr.sin_port        = htons(g_discovery.discovery_port);
    bc_addr.sin_addr.s_addr = INADDR_BROADCAST;

    uint8_t buf[MVGAL_NET_HEADER_SIZE + sizeof(payload)];
    memcpy(buf, &hdr, sizeof(hdr));
    memcpy(buf + sizeof(hdr), &payload, sizeof(payload));

    ssize_t sent = sendto(g_discovery.sock_fd, buf, sizeof(buf), 0,
                          (struct sockaddr *)&bc_addr, sizeof(bc_addr));

    if (sent < 0)
        return -1;

    return 0;
}

bool mvgal_discovery_is_running(void)
{
    return g_discovery.running;
}
