/**
 * @file mvgal_unified_heap.c
 * @brief Unified VRAM Heap Implementation
 *
 * Multi-Vendor GPU Aggregation Layer for Linux
 */

#include "mvgal/mvgal_unified_heap.h"
#include "mvgal/mvgal_p2p_dma.h"
#include "mvgal/mvgal_gpu.h"
#include "mvgal/mvgal_log.h"

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <pthread.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <errno.h>

/* ============================================================================
 * Unified Heap Structure
 * ============================================================================ */

struct mvgal_unified_heap {
    pthread_mutex_t lock;
    bool initialized;

    /* Allocation tracking */
    struct {
        uint64_t *ids;
        void **host_ptrs;
        size_t *sizes;
        uint32_t *primary_gpus;
        uint64_t *gpu_addresses; /* gpu_count * allocation_count */
        bool *populated;
        size_t count;
        size_t capacity;
    } allocations;

    /* Page fault handling */
    mvgal_page_fault_cb_t page_fault_cb;
    void *page_fault_user_data;

    /* Statistics */
    uint64_t total_bytes;
    uint64_t cross_gpu_transfers;
    uint64_t page_faults;

    /* Next allocation ID */
    uint64_t next_alloc_id;
};

#define MVGAL_HEAP_INITIAL_CAPACITY 256

/* ============================================================================
 * PCIe P2P Detection
 * ============================================================================ */

static uint64_t g_p2p_matrix[MVGAL_MAX_GPUS][MVGAL_MAX_GPUS];

bool mvgal_unified_heap_p2p_supported(uint32_t gpu_a, uint32_t gpu_b)
{
    if (gpu_a >= MVGAL_MAX_GPUS || gpu_b >= MVGAL_MAX_GPUS) {
        return false;
    }
    return g_p2p_matrix[gpu_a][gpu_b] != 0;
}

uint64_t mvgal_unified_heap_get_peer_mask(uint32_t gpu_index)
{
    uint64_t mask = 0;
    uint32_t i;

    if (gpu_index >= MVGAL_MAX_GPUS) {
        return 0;
    }

    for (i = 0; i < MVGAL_MAX_GPUS; i++) {
        if (g_p2p_matrix[gpu_index][i]) {
            mask |= (1ULL << i);
        }
    }
    return mask;
}

/* ============================================================================
 * Unified Heap Implementation
 * ============================================================================ */

mvgal_error_t mvgal_unified_heap_init(mvgal_unified_heap_handle_t *heap_out)
{
    mvgal_unified_heap_handle_t heap;
    int i, j;

    heap = calloc(1, sizeof(*heap));
    if (!heap) {
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }

    pthread_mutex_init(&heap->lock, NULL);
    heap->initialized = true;
    heap->next_alloc_id = 1;

    heap->allocations.ids = calloc(MVGAL_HEAP_INITIAL_CAPACITY, sizeof(uint64_t));
    heap->allocations.host_ptrs = calloc(MVGAL_HEAP_INITIAL_CAPACITY, sizeof(void *));
    heap->allocations.sizes = calloc(MVGAL_HEAP_INITIAL_CAPACITY, sizeof(size_t));
    heap->allocations.primary_gpus = calloc(MVGAL_HEAP_INITIAL_CAPACITY, sizeof(uint32_t));
    heap->allocations.gpu_addresses = calloc(MVGAL_HEAP_INITIAL_CAPACITY * MVGAL_MAX_GPUS, sizeof(uint64_t));
    heap->allocations.populated = calloc(MVGAL_HEAP_INITIAL_CAPACITY * MVGAL_MAX_GPUS, sizeof(bool));

    if (!heap->allocations.ids || !heap->allocations.host_ptrs ||
        !heap->allocations.sizes || !heap->allocations.primary_gpus ||
        !heap->allocations.gpu_addresses || !heap->allocations.populated) {
        free(heap->allocations.ids);
        free(heap->allocations.host_ptrs);
        free(heap->allocations.sizes);
        free(heap->allocations.primary_gpus);
        free(heap->allocations.gpu_addresses);
        free(heap->allocations.populated);
        free(heap);
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }

    heap->allocations.capacity = MVGAL_HEAP_INITIAL_CAPACITY;
    heap->allocations.count = 0;

    /* Build P2P matrix from PCIe topology (via mvgal_p2p_dma) */
    {
        uint8_t matrix[MVGAL_MAX_GPUS * MVGAL_MAX_GPUS];
        uint32_t gpu_count = 0;

        memset(g_p2p_matrix, 0, sizeof(g_p2p_matrix));
        if (mvgal_p2p_build_matrix(matrix, sizeof(matrix), &gpu_count) == MVGAL_SUCCESS) {
            for (i = 0; i < (int)gpu_count && i < MVGAL_MAX_GPUS; i++) {
                for (j = 0; j < (int)gpu_count && j < MVGAL_MAX_GPUS; j++) {
                    g_p2p_matrix[i][j] = matrix[i * gpu_count + j];
                }
            }
        } else {
            for (i = 0; i < MVGAL_MAX_GPUS; i++) {
                g_p2p_matrix[i][i] = 1;
            }
        }
    }

    *heap_out = heap;
    return MVGAL_SUCCESS;
}

void mvgal_unified_heap_fini(mvgal_unified_heap_handle_t heap)
{
    size_t i, j;

    if (!heap) {
        return;
    }

    pthread_mutex_lock(&heap->lock);

    /* Free all allocations */
    for (i = 0; i < heap->allocations.count; i++) {
        if (heap->allocations.host_ptrs[i]) {
            munmap(heap->allocations.host_ptrs[i], heap->allocations.sizes[i]);
        }
    }

    free(heap->allocations.ids);
    free(heap->allocations.host_ptrs);
    free(heap->allocations.sizes);
    free(heap->allocations.primary_gpus);
    free(heap->allocations.gpu_addresses);
    free(heap->allocations.populated);
    pthread_mutex_unlock(&heap->lock);
    pthread_mutex_destroy(&heap->lock);
    free(heap);
}

uint64_t mvgal_unified_heap_allocate(mvgal_unified_heap_handle_t heap,
                                      size_t size,
                                      mvgal_placement_hint_t placement,
                                      mvgal_access_hint_t access_hint,
                                      void **ptr_out)
{
    uint64_t alloc_id;
    void *host_ptr = NULL;
    uint32_t gpu_count;
    uint32_t i, slot;

    if (!heap || !ptr_out) {
        return 0;
    }

    pthread_mutex_lock(&heap->lock);

    /* Expand capacity if needed */
    if (heap->allocations.count >= heap->allocations.capacity) {
        size_t new_cap = heap->allocations.capacity * 2;
        uint64_t *new_ids = realloc(heap->allocations.ids, new_cap * sizeof(uint64_t));
        void **new_ptrs = realloc(heap->allocations.host_ptrs, new_cap * sizeof(void *));
        size_t *new_sizes = realloc(heap->allocations.sizes, new_cap * sizeof(size_t));
        uint32_t *new_gpus = realloc(heap->allocations.primary_gpus, new_cap * sizeof(uint32_t));
        uint64_t *new_addrs = realloc(heap->allocations.gpu_addresses,
                                      new_cap * MVGAL_MAX_GPUS * sizeof(uint64_t));
        bool *new_pop = realloc(heap->allocations.populated,
                               new_cap * MVGAL_MAX_GPUS * sizeof(bool));

        if (!new_ids || !new_ptrs || !new_sizes || !new_gpus || !new_addrs || !new_pop) {
            pthread_mutex_unlock(&heap->lock);
            return 0;
        }

        heap->allocations.ids = new_ids;
        heap->allocations.host_ptrs = new_ptrs;
        heap->allocations.sizes = new_sizes;
        heap->allocations.primary_gpus = new_gpus;
        heap->allocations.gpu_addresses = new_addrs;
        heap->allocations.populated = new_pop;
        heap->allocations.capacity = new_cap;
    }

    slot = heap->allocations.count;
    alloc_id = heap->next_alloc_id++;

    /* Allocate host memory (backed for DMA-BUF export) */
    host_ptr = mmap(NULL, size, PROT_READ | PROT_WRITE,
                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (host_ptr == MAP_FAILED) {
        pthread_mutex_unlock(&heap->lock);
        return 0;
    }

    /* Get GPU count */
    gpu_count = mvgal_gpu_get_count();

    /* Initialize allocation record */
    heap->allocations.ids[slot] = alloc_id;
    heap->allocations.host_ptrs[slot] = host_ptr;
    heap->allocations.sizes[slot] = size;
    heap->allocations.primary_gpus[slot] = 0; /* Default to first GPU */

    for (i = 0; i < gpu_count; i++) {
        heap->allocations.gpu_addresses[slot * MVGAL_MAX_GPUS + i] = 0;
        heap->allocations.populated[slot * MVGAL_MAX_GPUS + i] = false;
    }

    heap->allocations.count++;
    heap->total_bytes += size;

    pthread_mutex_unlock(&heap->lock);

    *ptr_out = host_ptr;
    return alloc_id;
}

void mvgal_unified_heap_free(mvgal_unified_heap_handle_t heap, uint64_t alloc_id)
{
    size_t i;

    if (!heap) {
        return;
    }

    pthread_mutex_lock(&heap->lock);

    for (i = 0; i < heap->allocations.count; i++) {
        if (heap->allocations.ids[i] == alloc_id) {
            munmap(heap->allocations.host_ptrs[i], heap->allocations.sizes[i]);
            heap->allocations.ids[i] = 0;
            heap->allocations.host_ptrs[i] = NULL;
            heap->allocations.sizes[i] = 0;
            heap->allocations.primary_gpus[i] = 0;
            memset(&heap->allocations.gpu_addresses[i * MVGAL_MAX_GPUS], 0,
                   MVGAL_MAX_GPUS * sizeof(uint64_t));
            memset(&heap->allocations.populated[i * MVGAL_MAX_GPUS], 0,
                   MVGAL_MAX_GPUS * sizeof(bool));

            /* Compact array */
            if (i < heap->allocations.count - 1) {
                memmove(&heap->allocations.ids[i],
                        &heap->allocations.ids[i + 1],
                        (heap->allocations.count - i - 1) * sizeof(uint64_t));
                memmove(&heap->allocations.host_ptrs[i],
                        &heap->allocations.host_ptrs[i + 1],
                        (heap->allocations.count - i - 1) * sizeof(void *));
                memmove(&heap->allocations.sizes[i],
                        &heap->allocations.sizes[i + 1],
                        (heap->allocations.count - i - 1) * sizeof(size_t));
            }
            heap->allocations.count--;
            break;
        }
    }

    pthread_mutex_unlock(&heap->lock);
}

uint64_t mvgal_unified_heap_get_gpu_address(mvgal_unified_heap_handle_t heap,
                                             uint64_t alloc_id,
                                             uint32_t gpu_index)
{
    size_t i;

    if (!heap) {
        return 0;
    }

    pthread_mutex_lock(&heap->lock);

    for (i = 0; i < heap->allocations.count; i++) {
        if (heap->allocations.ids[i] == alloc_id) {
            uint64_t addr = heap->allocations.gpu_addresses[i * MVGAL_MAX_GPUS + gpu_index];
            pthread_mutex_unlock(&heap->lock);
            return addr;
        }
    }

    pthread_mutex_unlock(&heap->lock);
    return 0;
}

mvgal_error_t mvgal_unified_heap_migrate(mvgal_unified_heap_handle_t heap,
                                         uint64_t alloc_id,
                                         uint32_t target_gpu)
{
    size_t i, src_gpu;

    if (!heap) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    pthread_mutex_lock(&heap->lock);

    for (i = 0; i < heap->allocations.count; i++) {
        if (heap->allocations.ids[i] == alloc_id) {
            src_gpu = heap->allocations.primary_gpus[i];

            if (!mvgal_unified_heap_p2p_supported(src_gpu, target_gpu)) {
                pthread_mutex_unlock(&heap->lock);
                return MVGAL_ERROR_UNSUPPORTED;
            }

            /* Host-backed migration: copy bytes then update placement */
            {
                void *host_ptr = heap->allocations.host_ptrs[i];
                size_t size = heap->allocations.sizes[i];
                mvgal_p2p_context_t p2p = NULL;

                if (host_ptr != NULL && size > 0 &&
                    mvgal_p2p_init(&p2p) == MVGAL_SUCCESS) {
                    (void)mvgal_p2p_transfer(
                        p2p, host_ptr, host_ptr, size,
                        src_gpu, target_gpu, MVGAL_P2P_DMA_GPU_TO_GPU);
                    mvgal_p2p_fini(p2p);
                }
            }

            heap->allocations.primary_gpus[i] = target_gpu;
            heap->allocations.populated[i * MVGAL_MAX_GPUS + target_gpu] = true;
            heap->allocations.populated[i * MVGAL_MAX_GPUS + src_gpu] = false;
            heap->cross_gpu_transfers++;

            pthread_mutex_unlock(&heap->lock);
            return MVGAL_SUCCESS;
        }
    }

    pthread_mutex_unlock(&heap->lock);
    return MVGAL_ERROR_NOT_FOUND;
}

mvgal_error_t mvgal_unified_heap_replicate(mvgal_unified_heap_handle_t heap,
                                           uint64_t alloc_id,
                                           const uint32_t *gpu_indices,
                                           uint32_t gpu_count)
{
    size_t i, j;

    if (!heap || !gpu_indices) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    pthread_mutex_lock(&heap->lock);

    for (i = 0; i < heap->allocations.count; i++) {
        if (heap->allocations.ids[i] == alloc_id) {
            for (j = 0; j < gpu_count; j++) {
                if (gpu_indices[j] < MVGAL_MAX_GPUS) {
                    heap->allocations.populated[i * MVGAL_MAX_GPUS + gpu_indices[j]] = true;
                }
            }
            heap->cross_gpu_transfers++;
            pthread_mutex_unlock(&heap->lock);
            return MVGAL_SUCCESS;
        }
    }

    pthread_mutex_unlock(&heap->lock);
    return MVGAL_ERROR_NOT_FOUND;
}

mvgal_error_t mvgal_unified_heap_set_page_fault_handler(
    mvgal_unified_heap_handle_t heap,
    mvgal_page_fault_cb_t callback,
    void *user_data)
{
    if (!heap) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    pthread_mutex_lock(&heap->lock);
    heap->page_fault_cb = callback;
    heap->page_fault_user_data = user_data;
    pthread_mutex_unlock(&heap->lock);

    return MVGAL_SUCCESS;
}

mvgal_error_t mvgal_unified_heap_get_stats(mvgal_unified_heap_handle_t heap,
                                         uint64_t *total_bytes,
                                         uint64_t *cross_gpu_transfers,
                                         uint64_t *page_faults)
{
    if (!heap) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    pthread_mutex_lock(&heap->lock);
    if (total_bytes)
        *total_bytes = heap->total_bytes;
    if (cross_gpu_transfers)
        *cross_gpu_transfers = heap->cross_gpu_transfers;
    if (page_faults)
        *page_faults = heap->page_faults;
    pthread_mutex_unlock(&heap->lock);

    return MVGAL_SUCCESS;
}