/**
 * @file mvgal_p2p_dma.c
 * @brief Peer-to-Peer DMA Transfer Implementation
 *
 * Multi-Vendor GPU Aggregation Layer for Linux
 */

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

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

/* ============================================================================
 * P2P Context
 * ============================================================================ */

struct mvgal_p2p_context {
    bool initialized;
    uint32_t gpu_count;
    uint8_t p2p_matrix[MVGAL_MAX_GPUS][MVGAL_MAX_GPUS];
    uint32_t bandwidth_matrix[MVGAL_MAX_GPUS][MVGAL_MAX_GPUS]; /* MB/s */
};

/* ============================================================================
 * PCIe Topology Detection
 * ============================================================================ */

static bool read_pcie_path(uint32_t gpu_index, char *path, size_t path_len)
{
    char sysfs_path[256];
    int fd;
    ssize_t len;

    snprintf(sysfs_path, sizeof(sysfs_path),
             "/sys/class/drm/card%d/device/uevent", gpu_index);

    fd = open(sysfs_path, O_RDONLY);
    if (fd < 0) {
        /* Try alternative path */
        mvgal_gpu_descriptor_t gpu;
        if (mvgal_gpu_get_descriptor(gpu_index, &gpu) == MVGAL_SUCCESS) {
            snprintf(path, path_len, "%04x:%02x:%02x.%u",
                     gpu.pci_domain, gpu.pci_bus, gpu.pci_device, gpu.pci_function);
            return true;
        }
        return false;
    }

    len = read(fd, path, path_len - 1);
    close(fd);

    if (len < 0) {
        return false;
    }

    path[len] = '\0';
    return true;
}

/**
 * Get the PCI domain for a GPU device
 */
static bool get_pci_domain_for_gpu(uint32_t gpu_index, uint16_t *domain_out)
{
    char uevent_path[256];
    FILE *f;
    char line[512];
    int domain;

    snprintf(uevent_path, sizeof(uevent_path),
             "/sys/class/drm/card%u/device/uevent", gpu_index);

    f = fopen(uevent_path, "r");
    if (!f) {
        return false;
    }

    while (fgets(line, sizeof(line), f)) {
        if (sscanf(line, "PCI_SLOT_NAME=%04x:%*02x:%*02x.%*d", &domain) == 1) {
            break;
        }
    }
    fclose(f);

    if (domain_out) {
        *domain_out = (uint16_t)domain;
    }
    return true;
}

static bool check_same_root_complex(uint32_t gpu_a, uint32_t gpu_b)
{
    uint16_t domain_a, domain_b;

    if (!get_pci_domain_for_gpu(gpu_a, &domain_a))
        return false;
    if (!get_pci_domain_for_gpu(gpu_b, &domain_b))
        return false;

    /* On modern systems with a single root complex, all PCIe devices
     * share the same domain. This is the primary indicator for P2P capability. */
    return domain_a == domain_b;
}

/* ============================================================================
 * P2P Implementation
 * ============================================================================ */

mvgal_error_t mvgal_p2p_init(mvgal_p2p_context_t *ctx_out)
{
    mvgal_p2p_context_t ctx;
    uint32_t i, j, gpu_count;
    int count;

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

    count = mvgal_gpu_get_count();
    gpu_count = count > 0 ? (uint32_t)count : 0;
    ctx->gpu_count = gpu_count;

    /* Initialize P2P matrix */
    for (i = 0; i < MVGAL_MAX_GPUS; i++) {
        for (j = 0; j < MVGAL_MAX_GPUS; j++) {
            if (i == j) {
                ctx->p2p_matrix[i][j] = 1; /* Self-access always works */
                ctx->bandwidth_matrix[i][j] = 100000; /* Very high */
            } else if (i < gpu_count && j < gpu_count) {
                ctx->p2p_matrix[i][j] = check_same_root_complex(i, j) ? 1 : 0;
                /* Estimate bandwidth: PCIe 4.0 x16 = ~32 GB/s bidirectional */
                ctx->bandwidth_matrix[i][j] = 16000; /* 16 GB/s */
            }
        }
    }

    ctx->initialized = true;
    *ctx_out = ctx;

    MVGAL_LOG_INFO("P2P DMA initialized for %u GPUs", gpu_count);
    return MVGAL_SUCCESS;
}

void mvgal_p2p_fini(mvgal_p2p_context_t ctx)
{
    if (ctx) {
        free(ctx);
    }
}

mvgal_error_t mvgal_p2p_query_caps(uint32_t gpu_src, uint32_t gpu_dst, mvgal_p2p_caps_t *caps)
{
    if (!caps) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    memset(caps, 0, sizeof(*caps));

    /* Check if within bounds */
    if (gpu_src >= MVGAL_MAX_GPUS || gpu_dst >= MVGAL_MAX_GPUS) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    caps->supported = check_same_root_complex(gpu_src, gpu_dst);
    caps->same_root_complex = caps->supported;
    caps->nvlink_connected = false; /* Would need NVLink detection */
    caps->max_bandwidth_mb_s = 16000; /* PCIe 4.0 x16 estimate */
    caps->latency_us = 10; /* Estimate */

    return MVGAL_SUCCESS;
}

mvgal_error_t mvgal_p2p_transfer(mvgal_p2p_context_t ctx,
                                  void *src_addr,
                                  void *dst_addr,
                                  size_t size,
                                  uint32_t src_gpu,
                                  uint32_t dst_gpu,
                                  mvgal_p2p_direction_t direction)
{
    if (!ctx || !ctx->initialized) {
        return MVGAL_ERROR_NOT_INITIALIZED;
    }

    if (src_addr == NULL || dst_addr == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    /* GPU-to-GPU P2P */
    if (src_gpu != 0xFFFFFFFF && dst_gpu != 0xFFFFFFFF) {
        if (!check_same_root_complex(src_gpu, dst_gpu)) {
            return MVGAL_ERROR_UNSUPPORTED;
        }
    }

    /* Perform the transfer - in production, this would use:
     * - NVIDIA: cudaMemcpyPeer() or DMA-BUF
     * - AMD: memcpy with SDI mapping
     * - Intel: pagemap + DMA
     * For now, use host staging */
    memcpy(dst_addr, src_addr, size);

    return MVGAL_SUCCESS;
}

mvgal_error_t mvgal_p2p_transfer_async(mvgal_p2p_context_t ctx,
                                        void *src_addr,
                                        void *dst_addr,
                                        size_t size,
                                        uint32_t src_gpu,
                                        uint32_t dst_gpu,
                                        mvgal_fence_t *fence)
{
    if (!ctx || !ctx->initialized) {
        return MVGAL_ERROR_NOT_INITIALIZED;
    }

    /* In production: queue transfer and return fence */
    /* For now: synchronous transfer */
    return mvgal_p2p_transfer(ctx, src_addr, dst_addr, size, src_gpu, dst_gpu,
                              MVGAL_P2P_DMA_GPU_TO_GPU);
}

/**
 * @brief Check if two GPUs support cross-vendor DMA-BUF sharing
 *
 * For cross-vendor transfers, we need DMA-BUF support on both GPUs.
 * This checks for driver capability to import/export DMABUF.
 */
bool mvgal_p2p_cross_vendor_supported(uint32_t src_gpu, uint32_t dst_gpu)
{
    mvgal_gpu_descriptor_t src_desc, dst_desc;
    bool src_supports_dmabuf = false;
    bool dst_supports_dmabuf = false;
    
    /* Query GPU descriptors to check capabilities */
    if (mvgal_gpu_get_descriptor(src_gpu, &src_desc) == MVGAL_SUCCESS) {
        /* All modern Linux GPU drivers support DMA-BUF */
        /* NVIDIA: nvidia-drm.ko with DRM_IOCTL_NV_GPU_FENCE_CREATE */
        /* AMD: amdgpu with DRM_IOCTL_AMDGPU_DMABUF */
        /* Intel: i915 with DRM_IOCTL_I915_GEM_EXPORT_DMABUF */
        src_supports_dmabuf = true; /* Assume all GPUs support */
    }
    
    if (mvgal_gpu_get_descriptor(dst_gpu, &dst_desc) == MVGAL_SUCCESS) {
        dst_supports_dmabuf = true;
    }
    
    /* Cross-vendor DMA-BUF is supported if both GPUs can do it */
    return src_supports_dmabuf && dst_supports_dmabuf;
}

/**
 * @brief Get optimal transfer method for cross-vendor memory copy
 *
 * Returns the preferred method based on:
 * - Same vendor same root: P2P native
 * - Cross vendor same root: DMA-BUF
 * - Cross vendor different root: Host staging
 * - With NVLink: NVLink transfer
 */
mvgal_memory_copy_method_t mvgal_p2p_get_optimal_method(
    uint32_t src_gpu,
    uint32_t dst_gpu
) {
    mvgal_gpu_descriptor_t src_desc, dst_desc;
    
    if (src_gpu == dst_gpu) {
        return MVGAL_MEMORY_COPY_GPU_LOCAL;
    }
    
    /* Check same root complex */
    if (check_same_root_complex(src_gpu, dst_gpu)) {
        /* Check if same vendor */
        if (mvgal_gpu_get_descriptor(src_gpu, &src_desc) == MVGAL_SUCCESS &&
            mvgal_gpu_get_descriptor(dst_gpu, &dst_desc) == MVGAL_SUCCESS) {
            
            if (src_desc.vendor == dst_desc.vendor) {
                /* Same vendor - use native P2P */
                return MVGAL_MEMORY_COPY_P2P;
            } else {
                /* Cross vendor - use DMA-BUF */
                return MVGAL_MEMORY_COPY_DMA_BUF;
            }
        }
    }
    
    /* Different root complexes - host staging */
    return MVGAL_MEMORY_COPY_CPU;
}

/**
 * @brief Query bandwidth-aware optimal strategy
 *
 * For multi-GPU workloads, suggests whether to:
 * - Replicate buffers (high bandwidth, small buffers)
 * - Migrate buffers (low bandwidth, large buffers)
 * - Keep on host (very low bandwidth)
 */
mvgal_p2p_strategy_t mvgal_p2p_query_optimal_strategy(
    uint32_t src_gpu,
    uint32_t dst_gpu,
    size_t buffer_size
) {
    mvgal_p2p_caps_t caps;
    const size_t HIGH_BANDWIDTH_THRESHOLD = 8000; /* 8 GB/s */
    const size_t SMALL_BUFFER_THRESHOLD = 16 * 1024 * 1024; /* 16 MB */
    
    if (mvgal_p2p_query_caps(src_gpu, dst_gpu, &caps) != MVGAL_SUCCESS) {
        return MVGAL_P2P_STRATEGY_HOST_CACHED;
    }
    
    if (!caps.supported) {
        return MVGAL_P2P_STRATEGY_HOST_CACHED;
    }
    
    if (caps.max_bandwidth_mb_s > HIGH_BANDWIDTH_THRESHOLD &&
        buffer_size < SMALL_BUFFER_THRESHOLD) {
        return MVGAL_P2P_STRATEGY_REPLICATE;
    }
    
    if (caps.max_bandwidth_mb_s > HIGH_BANDWIDTH_THRESHOLD) {
        return MVGAL_P2P_STRATEGY_MIGRATE;
    }
    
    return MVGAL_P2P_STRATEGY_HOST_CACHED;
}

uint32_t mvgal_p2p_get_bandwidth(uint32_t gpu_src, uint32_t gpu_dst)
{
    mvgal_p2p_caps_t caps;

    if (mvgal_p2p_query_caps(gpu_src, gpu_dst, &caps) != MVGAL_SUCCESS) {
        return 0;
    }

    return caps.max_bandwidth_mb_s;
}

bool mvgal_p2p_check_pcie_topology(uint32_t gpu_a, uint32_t gpu_b)
{
    return check_same_root_complex(gpu_a, gpu_b);
}

mvgal_error_t mvgal_p2p_build_matrix(uint8_t *matrix, size_t matrix_size, uint32_t *gpu_count)
{
    int32_t count;
    uint32_t i, j;

    if (!matrix || !gpu_count) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    count = mvgal_gpu_get_count();
    *gpu_count = count > 0 ? (uint32_t)count : 0;

    if (*gpu_count * *gpu_count > matrix_size) {
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }

    for (i = 0; i < *gpu_count; i++) {
        for (j = 0; j < *gpu_count; j++) {
            matrix[i * *gpu_count + j] = (check_same_root_complex(i, j) || i == j) ? 1 : 0;
        }
    }

    return MVGAL_SUCCESS;
}