/**
 * @file mvgal_pool.c
 * @brief Network GPU Pool Manager implementation
 *
 * SPDX-License-Identifier: MIT
 */

#include "mvgal_pool.h"
#include "mvgal/mvgal_pool.h"
#include <string.h>
#include <time.h>
#include <stdio.h>

/******************************************************************************
 * Global pool state
 ******************************************************************************/

static mvgal_pool_internal_t g_pool;

/******************************************************************************
 * Time helpers
 ******************************************************************************/

static uint64_t pool_time_ns(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec;
}

/******************************************************************************
 * Remote GPU → local descriptor conversion
 ******************************************************************************/

static void gpu_descriptor_from_remote(const mvgal_net_remote_gpu_t *src,
                                       mvgal_gpu_descriptor_t *dst)
{
    memset(dst, 0, sizeof(*dst));

    dst->id              = src->local_gpu_index;
    dst->vendor          = MVGAL_VENDOR_UNKNOWN;
    dst->vendor_id       = (uint16_t)src->vendor_id;
    dst->device_id       = (uint16_t)src->device_id;
    dst->vram_total      = src->vram_bytes;
    dst->vram_free       = src->vram_free;
    dst->compute_score   = (float)src->compute_units * 10.0f;
    dst->graphics_score  = src->supports_graphics ? 50.0f : 0.0f;
    dst->available       = src->is_online != 0;
    dst->enabled         = true;

    snprintf(dst->name, sizeof(dst->name),
             "Remote GPU %u @ node 0x%llx",
             src->local_gpu_index,
             (unsigned long long)src->peer_node_id);

    dst->pcie_gen        = MVGAL_PCIE_UNKNOWN;

    if (src->api_flags & (1U << 0))
        dst->supported_apis[dst->api_count++] = MVGAL_API_VULKAN;
    if (src->api_flags & (1U << 1))
        dst->supported_apis[dst->api_count++] = MVGAL_API_OPENCL;
    if (src->api_flags & (1U << 2))
        dst->supported_apis[dst->api_count++] = MVGAL_API_CUDA;
}

/******************************************************************************
 * Transport message callback
 ******************************************************************************/

static void pool_msg_callback(uint64_t node_id,
                              const mvgal_net_message_header_t *hdr,
                              const void *payload,
                              void *userdata)
{
    (void)userdata;

    pthread_mutex_lock(&g_pool.lock);

    g_pool.stats.messages_received++;
    g_pool.stats.bytes_received += MVGAL_NET_HEADER_SIZE + hdr->payload_size;

    switch (hdr->msg_type) {

    case MVGAL_NET_MSG_GPU_ANNOUNCE: {
        const mvgal_net_gpu_announce_t *ann = (const mvgal_net_gpu_announce_t *)payload;
        const mvgal_net_gpu_descriptor_t *gpu_descs =
            (const mvgal_net_gpu_descriptor_t *)(ann + 1);
        mvgal_pool_update_gpus(node_id, ann, gpu_descs, ann->num_gpus);
        break;
    }

    case MVGAL_NET_MSG_WORK_COMPLETE: {
        if (hdr->payload_size < sizeof(uint64_t) + sizeof(uint32_t))
            break;
        const uint64_t *wid = (const uint64_t *)payload;
        const uint32_t *res = (const uint32_t *)(wid + 1);
        mvgal_pool_notify_work_complete(node_id, *wid, *res);
        break;
    }

    case MVGAL_NET_MSG_FENCE_SYNC: {
        /* Fence sync from remote peer — forward to local fence manager (stub) */
        break;
    }

    case MVGAL_NET_MSG_MEM_TRANSFER: {
        /* Incoming memory transfer (stub — full impl needs DMA-BUF support) */
        break;
    }

    case MVGAL_NET_MSG_HEARTBEAT: {
        /* Update peer last-seen, reflect HEARTBEAT_ACK back */
        for (uint32_t i = 0; i < g_pool.peer_count; i++) {
            if (g_pool.peers[i].node_id == node_id) {
                g_pool.peers[i].last_seen_ns = pool_time_ns();
                break;
            }
        }
        /* Send HEARTBEAT_ACK */
        mvgal_net_message_header_t ack_hdr = *hdr;
        ack_hdr.msg_type = MVGAL_NET_MSG_HEARTBEAT_ACK;
        mvgal_transport_send(node_id, &ack_hdr, NULL, 0);
        break;
    }

    case MVGAL_NET_MSG_HEARTBEAT_ACK: {
        for (uint32_t i = 0; i < g_pool.peer_count; i++) {
            if (g_pool.peers[i].node_id == node_id) {
                g_pool.peers[i].last_seen_ns = pool_time_ns();
                break;
            }
        }
        break;
    }

    case MVGAL_NET_MSG_PEER_GOODBYE: {
        mvgal_pool_remove_peer(node_id);
        break;
    }

    case MVGAL_NET_MSG_ERROR: {
        g_pool.stats.errors++;
        break;
    }

    default:
        break;
    }

    pthread_mutex_unlock(&g_pool.lock);
}

/******************************************************************************
 * Transport peer callback
 ******************************************************************************/

static void pool_peer_callback(uint64_t node_id,
                               mvgal_net_peer_state_t state,
                               void *userdata)
{
    (void)userdata;

    pthread_mutex_lock(&g_pool.lock);

    switch (state) {

    case MVGAL_NET_PEER_ACTIVE: {
        /* Peer fully connected — request GPU announce */
        for (uint32_t i = 0; i < g_pool.peer_count; i++) {
            if (g_pool.peers[i].node_id == node_id) {
                g_pool.peers[i].peer_state = MVGAL_NET_PEER_ACTIVE;
                break;
            }
        }
        g_pool.stats.peers_discovered++;
        break;
    }

    case MVGAL_NET_PEER_DISCONNECTED:
    case MVGAL_NET_PEER_TIMEOUT: {
        mvgal_pool_remove_peer(node_id);
        if (state == MVGAL_NET_PEER_TIMEOUT) {
            g_pool.stats.peers_lost++;
            g_pool.stats.errors++;
        } else {
            g_pool.stats.peers_lost++;
        }
        break;
    }

    default:
        break;
    }

    pthread_mutex_unlock(&g_pool.lock);
}

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

int mvgal_pool_init(const mvgal_net_config_t *config)
{
    if (g_pool.initialized)
        return -1;

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

    /* Apply config or defaults */
    if (config) {
        g_pool.config = *config;
    } else {
        g_pool.config.discovery_port        = MVGAL_NET_DEFAULT_PORT;
        g_pool.config.data_port             = 0;
        g_pool.config.heartbeat_interval_ms = MVGAL_NET_DEFAULT_HEARTBEAT_MS;
        g_pool.config.timeout_ms            = MVGAL_NET_DEFAULT_TIMEOUT_MS;
        g_pool.config.max_peers             = MVGAL_NET_MAX_PEERS;
        g_pool.config.enable_discovery      = 1;
        g_pool.config.enable_tcp            = 1;
        g_pool.config.max_transfer_bytes    = MVGAL_NET_MAX_PAYLOAD;
    }

    if (pthread_mutex_init(&g_pool.lock, NULL) != 0)
        return -1;

    /* Register transport callbacks */
    mvgal_transport_set_msg_callback(pool_msg_callback, NULL);
    mvgal_transport_set_peer_callback(pool_peer_callback, NULL);

    g_pool.initialized = true;
    return 0;
}

void mvgal_pool_shutdown(void)
{
    if (!g_pool.initialized)
        return;

    pthread_mutex_lock(&g_pool.lock);

    /* Send PEER_GOODBYE to all connected peers */
    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_GOODBYE;
    hdr.checksum    = MVGAL_NET_FLAG_NO_CHECKSUM;
    mvgal_transport_broadcast(&hdr, NULL, 0);
    mvgal_transport_disconnect_all();

    g_pool.discovery_running = false;
    g_pool.initialized = false;

    pthread_mutex_unlock(&g_pool.lock);
    pthread_mutex_destroy(&g_pool.lock);

    /* Unregister callbacks */
    mvgal_transport_set_msg_callback(NULL, NULL);
    mvgal_transport_set_peer_callback(NULL, NULL);
}

int32_t mvgal_pool_get_remote_gpu_count(void)
{
    if (!g_pool.initialized)
        return 0;

    pthread_mutex_lock(&g_pool.lock);
    uint32_t count = g_pool.remote_gpu_count;
    pthread_mutex_unlock(&g_pool.lock);
    return (int32_t)count;
}

mvgal_error_t mvgal_pool_get_remote_gpu(uint32_t index, mvgal_gpu_descriptor_t *gpu)
{
    if (!g_pool.initialized || !gpu)
        return MVGAL_ERROR_INVALID_ARGUMENT;

    pthread_mutex_lock(&g_pool.lock);

    if (index >= g_pool.remote_gpu_count) {
        pthread_mutex_unlock(&g_pool.lock);
        return MVGAL_ERROR_NOT_FOUND;
    }

    /* Walk entries to find the Nth in-use GPU */
    uint32_t found = 0;
    mvgal_error_t ret = MVGAL_ERROR_NOT_FOUND;
    for (uint32_t i = 0; i < MVGAL_NET_MAX_PEERS * MVGAL_NET_MAX_GPU_ANNOUNCE; i++) {
        if (g_pool.remote_gpus[i].in_use) {
            if (found == index) {
                gpu_descriptor_from_remote(&g_pool.remote_gpus[i].gpu, gpu);
                ret = MVGAL_SUCCESS;
                break;
            }
            found++;
        }
    }

    pthread_mutex_unlock(&g_pool.lock);
    return ret;
}

mvgal_error_t mvgal_pool_get_stats(mvgal_net_stats_t *stats)
{
    if (!g_pool.initialized || !stats)
        return MVGAL_ERROR_INVALID_ARGUMENT;

    pthread_mutex_lock(&g_pool.lock);
    *stats = g_pool.stats;

    /* Count active peers */
    stats->active_peers = 0;
    for (uint32_t i = 0; i < g_pool.peer_count; i++) {
        if (g_pool.peers[i].peer_state == MVGAL_NET_PEER_ACTIVE)
            stats->active_peers++;
    }
    stats->active_remote_gpus = g_pool.remote_gpu_count;

    pthread_mutex_unlock(&g_pool.lock);
    return MVGAL_SUCCESS;
}

/******************************************************************************
 * Internal functions
 ******************************************************************************/

int mvgal_pool_add_peer(uint64_t node_id, const char *instance_name,
                        uint32_t num_gpus, uint32_t protocol_version,
                        const uint8_t *addr, uint32_t addr_len)
{
    pthread_mutex_lock(&g_pool.lock);

    /* Check if peer already tracked */
    for (uint32_t i = 0; i < g_pool.peer_count; i++) {
        if (g_pool.peers[i].node_id == node_id) {
            pthread_mutex_unlock(&g_pool.lock);
            return 0; /* Already known */
        }
    }

    if (g_pool.peer_count >= MVGAL_NET_MAX_PEERS) {
        pthread_mutex_unlock(&g_pool.lock);
        return -1;
    }

    mvgal_net_peer_t *peer = &g_pool.peers[g_pool.peer_count++];
    memset(peer, 0, sizeof(*peer));
    peer->node_id          = node_id;
    peer->num_gpus         = num_gpus;
    peer->protocol_version = protocol_version;
    peer->peer_state       = MVGAL_NET_PEER_DISCOVERED;
    peer->last_seen_ns     = pool_time_ns();

    if (instance_name)
        snprintf(peer->instance_name, sizeof(peer->instance_name), "%s", instance_name);

    if (addr && addr_len > 0) {
        size_t copy_len = addr_len < sizeof(peer->addr_storage) ? addr_len : sizeof(peer->addr_storage);
        memcpy(peer->addr_storage, addr, copy_len);
        peer->addr_len = (uint32_t)copy_len;
    }

    pthread_mutex_unlock(&g_pool.lock);
    return 0;
}

void mvgal_pool_remove_peer(uint64_t node_id)
{
    pthread_mutex_lock(&g_pool.lock);

    /* Remove peer GPUs */
    for (uint32_t i = 0; i < MVGAL_NET_MAX_PEERS * MVGAL_NET_MAX_GPU_ANNOUNCE; i++) {
        if (g_pool.remote_gpus[i].in_use &&
            g_pool.remote_gpus[i].gpu.peer_node_id == node_id) {
            g_pool.remote_gpus[i].in_use = false;
            g_pool.remote_gpu_count--;
        }
    }

    /* Remove peer entry (compact array) */
    for (uint32_t i = 0; i < g_pool.peer_count; i++) {
        if (g_pool.peers[i].node_id == node_id) {
            uint32_t remaining = g_pool.peer_count - i - 1;
            if (remaining > 0)
                memmove(&g_pool.peers[i], &g_pool.peers[i + 1],
                        remaining * sizeof(mvgal_net_peer_t));
            g_pool.peer_count--;
            break;
        }
    }

    /* Cleanup transport if still connected */
    if (mvgal_transport_is_connected(node_id))
        mvgal_transport_disconnect(node_id);

    pthread_mutex_unlock(&g_pool.lock);
}

int mvgal_pool_update_gpus(uint64_t peer_node_id,
                           const mvgal_net_gpu_announce_t *announce,
                           const mvgal_net_gpu_descriptor_t *gpu_descs,
                           uint32_t num_gpus)
{
    int added = 0;

    pthread_mutex_lock(&g_pool.lock);

    /* First, clear stale entries for this peer */
    for (uint32_t i = 0; i < MVGAL_NET_MAX_PEERS * MVGAL_NET_MAX_GPU_ANNOUNCE; i++) {
        if (g_pool.remote_gpus[i].in_use &&
            g_pool.remote_gpus[i].gpu.peer_node_id == peer_node_id) {
            g_pool.remote_gpus[i].in_use = false;
            g_pool.remote_gpu_count--;
        }
    }

    /* Insert fresh entries */
    uint32_t max_gpus = num_gpus < MVGAL_NET_MAX_GPU_ANNOUNCE ? num_gpus : MVGAL_NET_MAX_GPU_ANNOUNCE;
    for (uint32_t i = 0; i < max_gpus; i++) {
        /* Find free slot */
        for (uint32_t j = 0; j < MVGAL_NET_MAX_PEERS * MVGAL_NET_MAX_GPU_ANNOUNCE; j++) {
            if (!g_pool.remote_gpus[j].in_use) {
                mvgal_remote_gpu_entry_t *entry = &g_pool.remote_gpus[j];
                entry->in_use = true;
                entry->last_seen_ns = pool_time_ns();
                entry->health_check_idx = 0;

                mvgal_net_remote_gpu_t *rg = &entry->gpu;
                memset(rg, 0, sizeof(*rg));
                rg->peer_node_id    = peer_node_id;
                rg->local_gpu_index = gpu_descs[i].gpu_index;
                rg->vendor_id       = gpu_descs[i].vendor_id;
                rg->device_id       = gpu_descs[i].device_id;
                rg->vram_bytes      = gpu_descs[i].vram_bytes;
                rg->compute_units   = gpu_descs[i].compute_units;
                rg->api_flags       = gpu_descs[i].api_flags;
                rg->pcie_gen        = gpu_descs[i].pcie_gen;
                rg->pcie_lanes      = gpu_descs[i].pcie_lanes;
                rg->supports_graphics = gpu_descs[i].supports_graphics;
                rg->supports_compute  = gpu_descs[i].supports_compute;
                rg->is_online         = 1;
                rg->latency_us        = 0; /* Will be estimated by heartbeat RTT */

                /* Estimate free VRAM as ~80% of total initially */
                rg->vram_free = (uint64_t)((float)rg->vram_bytes * 0.8f);

                g_pool.remote_gpu_count++;
                added++;
                break;
            }
        }
    }

    /* Update peer info */
    for (uint32_t i = 0; i < g_pool.peer_count; i++) {
        if (g_pool.peers[i].node_id == peer_node_id) {
            g_pool.peers[i].last_seen_ns = pool_time_ns();
            g_pool.peers[i].load_average = announce->load_average;
            break;
        }
    }

    pthread_mutex_unlock(&g_pool.lock);
    return added;
}

void mvgal_pool_notify_work_complete(uint64_t peer_node_id,
                                     uint64_t workload_id,
                                     uint32_t result)
{
    (void)peer_node_id;
    (void)workload_id;
    (void)result;

    pthread_mutex_lock(&g_pool.lock);
    g_pool.stats.work_completed++;
    pthread_mutex_unlock(&g_pool.lock);
}

mvgal_pool_internal_t *mvgal_pool_get_instance(void)
{
    return &g_pool;
}
