/**
 * @file mvgal_remote.c
 * @brief Remote work dispatch, memory transfers, and fence sync
 *
 * SPDX-License-Identifier: MIT
 */

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

/******************************************************************************
 * Serialize a workload descriptor into the wire format
 ******************************************************************************/

static void serialize_workload(const mvgal_workload_descriptor_t *desc,
                               mvgal_net_work_submit_t *wire)
{
    memset(wire, 0, sizeof(*wire));

    /* Use the workload ID + a hash as the global workload identifier */
    wire->workload_id  = (uint64_t)desc->id;
    wire->source_gpu   = MVGAL_NET_MAX_GPU_ANNOUNCE; /* Local sentinel index */
    wire->target_gpu   = desc->assigned_gpu;
    wire->workload_type = (uint32_t)desc->type;

    /* Estimate size from memory requirement */
    wire->size_bytes   = (uint64_t)desc->memory_required;

    /* Map priority: mvgal uses 0-100, wire uses 0-15 */
    wire->priority     = (desc->priority > 0)
                             ? (uint32_t)(desc->priority / 7) + 1
                             : 0;
    if (wire->priority > 15)
        wire->priority = 15;

    /* Map deadline to timeout */
    wire->timeout_ms   = (desc->deadline > 0 && desc->deadline > desc->submission_time)
                             ? (uint32_t)((desc->deadline - desc->submission_time) / 1000000ULL)
                             : 0;
}

/******************************************************************************
 * Serialize a memory transfer descriptor
 ******************************************************************************/

static void serialize_mem_transfer(uint64_t transfer_id,
                                   uint32_t direction,
                                   uint32_t source_gpu,
                                   uint32_t target_gpu,
                                   uint64_t offset,
                                   uint64_t size_bytes,
                                   mvgal_net_mem_transfer_t *wire)
{
    memset(wire, 0, sizeof(*wire));
    wire->transfer_id = transfer_id;
    wire->direction   = direction;
    wire->source_gpu  = source_gpu;
    wire->target_gpu  = target_gpu;
    wire->offset      = offset;
    wire->size_bytes  = size_bytes;
    wire->dmabuf_fd   = (uint32_t)-1; /* DMA-BUF not yet supported over network */
}

/******************************************************************************
 * Build a network message header
 ******************************************************************************/

static void build_header(mvgal_net_message_header_t *hdr,
                         uint32_t msg_type,
                         size_t payload_size)
{
    memset(hdr, 0, sizeof(*hdr));
    hdr->magic        = MVGAL_NET_MAGIC;
    hdr->version      = MVGAL_NET_VERSION;
    hdr->msg_type     = msg_type;
    hdr->payload_size = (uint32_t)payload_size;
    hdr->flags        = MVGAL_NET_FLAG_NO_CHECKSUM;
}

/******************************************************************************
 * Public API — Remote work submission
 ******************************************************************************/

int mvgal_remote_submit_workload(const mvgal_workload_descriptor_t *desc,
                                 uint64_t target_peer_node_id)
{
    if (!desc || !mvgal_transport_is_connected(target_peer_node_id))
        return -1;

    mvgal_net_work_submit_t wire;
    serialize_workload(desc, &wire);

    mvgal_net_message_header_t hdr;
    build_header(&hdr, MVGAL_NET_MSG_WORK_SUBMIT, sizeof(wire));

    int ret = mvgal_transport_send(target_peer_node_id, &hdr, &wire, sizeof(wire));

    if (ret == 0) {
        mvgal_pool_internal_t *pool = mvgal_pool_get_instance();
        pthread_mutex_lock(&pool->lock);
        pool->stats.work_submitted++;
        pool->stats.messages_sent++;
        pool->stats.bytes_sent += MVGAL_NET_HEADER_SIZE + sizeof(wire);
        pthread_mutex_unlock(&pool->lock);
    }

    return ret;
}

/******************************************************************************
 * Public API — Remote memory transfer
 ******************************************************************************/

int mvgal_remote_transfer_memory(uint64_t target_peer_node_id,
                                 uint32_t direction,
                                 uint32_t source_gpu,
                                 uint32_t target_gpu,
                                 uint64_t offset,
                                 uint64_t size_bytes)
{
    if (!mvgal_transport_is_connected(target_peer_node_id))
        return -1;

    static uint64_t next_transfer_id = 1;

    mvgal_net_mem_transfer_t wire;
    serialize_mem_transfer(next_transfer_id++, direction,
                           source_gpu, target_gpu, offset, size_bytes, &wire);

    mvgal_net_message_header_t hdr;
    build_header(&hdr, MVGAL_NET_MSG_MEM_TRANSFER, sizeof(wire));

    int ret = mvgal_transport_send(target_peer_node_id, &hdr, &wire, sizeof(wire));

    if (ret == 0) {
        mvgal_pool_internal_t *pool = mvgal_pool_get_instance();
        pthread_mutex_lock(&pool->lock);
        pool->stats.transfers_initiated++;
        pool->stats.messages_sent++;
        pool->stats.bytes_sent += MVGAL_NET_HEADER_SIZE + sizeof(wire);
        pthread_mutex_unlock(&pool->lock);
    }

    return ret;
}

/******************************************************************************
 * Public API — Remote fence sync
 ******************************************************************************/

int mvgal_remote_fence_sync(uint64_t fence_id,
                            uint32_t signal_value,
                            uint32_t wait_value,
                            uint32_t flags,
                            uint64_t *target_peer_ids,
                            uint32_t peer_count)
{
    if (!target_peer_ids || peer_count == 0)
        return -1;

    mvgal_net_fence_sync_t wire;
    memset(&wire, 0, sizeof(wire));
    wire.fence_id     = fence_id;
    wire.signal_value = signal_value;
    wire.wait_value   = wait_value;
    wire.flags        = flags;

    mvgal_net_message_header_t hdr;
    build_header(&hdr, MVGAL_NET_MSG_FENCE_SYNC, sizeof(wire));

    int success_count = 0;

    for (uint32_t i = 0; i < peer_count; i++) {
        if (mvgal_transport_is_connected(target_peer_ids[i])) {
            int ret = mvgal_transport_send(target_peer_ids[i], &hdr, &wire, sizeof(wire));
            if (ret == 0) {
                success_count++;
                mvgal_pool_internal_t *pool = mvgal_pool_get_instance();
                pthread_mutex_lock(&pool->lock);
                pool->stats.messages_sent++;
                pool->stats.bytes_sent += MVGAL_NET_HEADER_SIZE + sizeof(wire);
                pthread_mutex_unlock(&pool->lock);
            }
        }
    }

    return (success_count > 0) ? 0 : -1;
}

/******************************************************************************
 * Public API — Get online peers for a fence broadcast
 ******************************************************************************/

uint32_t mvgal_remote_get_active_peers(uint64_t *peer_ids, uint32_t max_count)
{
    if (!peer_ids || max_count == 0)
        return 0;

    mvgal_pool_internal_t *pool = mvgal_pool_get_instance();
    uint32_t count = 0;

    pthread_mutex_lock(&pool->lock);

    for (uint32_t i = 0; i < pool->peer_count && count < max_count; i++) {
        if (pool->peers[i].peer_state == MVGAL_NET_PEER_ACTIVE) {
            peer_ids[count++] = pool->peers[i].node_id;
        }
    }

    pthread_mutex_unlock(&pool->lock);
    return count;
}
