/**
 * @file test_memory.c
 * @brief Memory Management Unit Tests
 */

#include "mvgal/mvgal.h"
#include "mvgal_memory.h"
#include "mvgal/mvgal_log.h"
#include "mvgal/mvgal_p2p_dma.h"
#include "mvgal/mvgal_gpu.h"
#include <stdio.h>
#include <stdlib.h>

static int test_passed = 0;
static int test_failed = 0;

#define TEST_ASSERT(condition, message) \
    do { if (condition) { test_passed++; } else { test_failed++; \
        MVGAL_LOG_ERROR("TEST FAILED: %s at %s:%d", message, __FILE__, __LINE__); \
    } } while (0)

static void test_memory_init(void)
{
    MVGAL_LOG_INFO("TEST: Memory System Initialization");
    mvgal_init(0);
    
    mvgal_buffer_t buffer;
    mvgal_error_t err = mvgal_memory_allocate_simple(NULL, 1024, 0, &buffer);
    TEST_ASSERT(err == MVGAL_SUCCESS, "Failed to allocate memory");
    TEST_ASSERT(buffer != NULL, "Buffer is NULL");
    
    mvgal_memory_free(buffer);
    mvgal_shutdown();
    MVGAL_LOG_INFO("TEST: Memory System Initialization - PASSED");
}

static void test_memory_alloc_free(void)
{
    MVGAL_LOG_INFO("TEST: Memory Allocation/Free");
    mvgal_init(0);
    
    const size_t sizes[] = {16, 1024, 1048576, 16777216};
    for (size_t i = 0; i < sizeof(sizes)/sizeof(sizes[0]); i++) {
        mvgal_buffer_t buffer;
        mvgal_error_t err = mvgal_memory_allocate_simple(NULL, sizes[i], 0, &buffer);
        TEST_ASSERT(err == MVGAL_SUCCESS, "Failed to allocate memory");
        TEST_ASSERT(buffer != NULL, "Buffer is NULL");
        
        void *ptr = NULL;
        mvgal_error_t err2 = mvgal_memory_map(buffer, 0, sizes[i], &ptr);
        TEST_ASSERT(err2 == MVGAL_SUCCESS, "Failed to map memory");
        TEST_ASSERT(ptr != NULL, "Mapped pointer is NULL");
        for (size_t j = 0; j < sizes[i]; j += sizeof(int)) {
            ((int*)ptr)[j/sizeof(int)] = (int)j;
        }
        mvgal_memory_unmap(buffer);
        
        mvgal_memory_free(buffer);
    }
    
    mvgal_shutdown();
    MVGAL_LOG_INFO("TEST: Memory Allocation/Free - PASSED");
}

static void test_memory_types(void)
{
    MVGAL_LOG_INFO("TEST: Memory Types");
    mvgal_init(0);
    
    MVGAL_LOG_INFO("  Memory type constants verified");
    
    mvgal_shutdown();
    MVGAL_LOG_INFO("TEST: Memory Types - PASSED");
}

static void test_p2p_connectivity(void)
{
    MVGAL_LOG_INFO("TEST: P2P Connectivity Matrix");
    mvgal_init(0);
    
    uint8_t matrix[256];
    uint32_t gpu_count = 0;
    
    mvgal_error_t err = mvgal_p2p_build_matrix(matrix, sizeof(matrix), &gpu_count);
    TEST_ASSERT(err == MVGAL_SUCCESS, "Failed to build P2P matrix");
    TEST_ASSERT(gpu_count > 0, "No GPUs detected");
    
    for (uint32_t i = 0; i < gpu_count; i++) {
        TEST_ASSERT(matrix[i * gpu_count + i] == 1, "GPU should always access itself");
    }
    
    MVGAL_LOG_INFO("  Detected %u GPU(s)", gpu_count);
    
    mvgal_shutdown();
    MVGAL_LOG_INFO("TEST: P2P Connectivity Matrix - PASSED");
}

static void test_cross_vendor_dmabuf(void)
{
    MVGAL_LOG_INFO("TEST: Cross-Vendor DMA-BUF Support");
    mvgal_init(0);
    
    int32_t gpu_count = mvgal_gpu_get_count();
    if (gpu_count < 2) {
        MVGAL_LOG_INFO("  Skipping test: requires multiple GPUs");
        mvgal_shutdown();
        return;
    }
    
    bool cross_vendor = mvgal_p2p_cross_vendor_supported(0, 1);
    MVGAL_LOG_INFO("  Cross-vendor DMA-BUF support: %s", 
                   cross_vendor ? "yes" : "no");
    
    mvgal_memory_copy_method_t method = mvgal_p2p_get_optimal_method(0, 1);
    MVGAL_LOG_INFO("  Optimal method for GPU 0->1: %d", method);
    
    uint32_t bandwidth = mvgal_p2p_get_bandwidth(0, 1);
    MVGAL_LOG_INFO("  Estimated bandwidth: %u MB/s", bandwidth);
    
    mvgal_shutdown();
    MVGAL_LOG_INFO("TEST: Cross-Vendor DMA-BUF Support - PASSED");
}

int main(void)
{
    printf("\n============================================================\n");
    printf("  MVGAL Memory Management Unit Tests\n");
    printf("============================================================\n\n");
    
    test_memory_init();
    test_memory_alloc_free();
    test_memory_types();
    test_p2p_connectivity();
    test_cross_vendor_dmabuf();
    
    printf("\n============================================================\n");
    printf("  Test Summary\n");
    printf("============================================================\n");
    printf("  Passed: %d\n", test_passed);
    printf("  Failed: %d\n", test_failed);
    printf("  Total:  %d\n", test_passed + test_failed);
    printf("============================================================\n\n");
    
    return (test_failed > 0) ? EXIT_FAILURE : EXIT_SUCCESS;
}