#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/kthread.h>
#include <linux/delay.h>
#include <linux/mm.h>
#include <linux/dma-mapping.h>
#include <linux/vmalloc.h>
#include <linux/spinlock.h>
#include <linux/ktime.h>
#include <linux/atomic.h>

#include "gsp_emulation.h"
#include "gsp_rpc_client.h"
#include "memory.h"

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Open-IOV");
MODULE_DESCRIPTION("GSP Emulation Unit Tests");

#define TEST_ALLOC_SIZE (4096)
#define TEST_MAP_SIZE   (4096)
#define TEST_PIN_PAGES  (1)

static struct gsp_emu_context *test_ctx = NULL;
static struct gsp_client *test_client = NULL;
static void *test_shared_mem = NULL;
static dma_addr_t test_shared_dma = 0;
static struct page **test_pages = NULL;
static int test_pages_count = 0;

static int test_seq_id = 1;

static int get_next_seq_id(void)
{
    return test_seq_id++;
}

static void init_test_context(void)
{
    test_shared_mem = dma_alloc_coherent(NULL, 4096, &test_shared_dma, GFP_KERNEL);
    if (!test_shared_mem) {
        pr_err("GSP_EMU_TEST: Failed to allocate shared memory\n");
        return;
    }
    memset(test_shared_mem, 0, 4096);

    test_pages_count = TEST_ALLOC_SIZE / PAGE_SIZE;
    test_pages = kmalloc_array(test_pages_count, sizeof(struct page *), GFP_KERNEL);
    if (!test_pages) {
        pr_err("GSP_EMU_TEST: Failed to allocate page array\n");
        return;
    }
    for (int i = 0; i < test_pages_count; i++) {
        test_pages[i] = alloc_page(GFP_KERNEL | __GFP_ZERO);
        if (!test_pages[i]) {
            pr_err("GSP_EMU_TEST: Failed to allocate page %d\n", i);
            test_pages_count = i;
            break;
        }
    }

    test_client = gsp_client_init(NULL);
    if (!test_client) {
        pr_err("GSP_EMU_TEST: Failed to init gsp_client\n");
        return;
    }

    test_ctx = kzalloc(sizeof(struct gsp_emu_context), GFP_KERNEL);
    if (!test_ctx) {
        pr_err("GSP_EMU_TEST: Failed to allocate context\n");
        return;
    }

    test_ctx->vgpu = NULL;
    test_ctx->bar0_vaddr = test_shared_mem;
    test_ctx->doorbell_reg = test_shared_mem + 0x110C00;
    test_ctx->shared_mem = test_shared_mem + 0x1000;
    test_ctx->rpc_client = test_client;
    spin_lock_init(&test_ctx->lock);
    test_ctx->running = true;
    atomic_set(&test_ctx->processed_rpcs, 0);
    atomic_set(&test_ctx->failed_rpcs, 0);
    test_ctx->total_latency_ns = 0;

    pr_info("GSP_EMU_TEST: Test context initialized\n");
}

static void cleanup_test_context(void)
{
    if (test_ctx) {
        test_ctx->running = false;
        kfree(test_ctx);
        test_ctx = NULL;
    }
    if (test_client) {
        gsp_client_fini(test_client);
        test_client = NULL;
    }
    if (test_pages) {
        for (int i = 0; i < test_pages_count; i++) {
            if (test_pages[i])
                __free_page(test_pages[i]);
        }
        kfree(test_pages);
        test_pages = NULL;
    }
    if (test_shared_mem) {
        dma_free_coherent(NULL, 4096, test_shared_mem, test_shared_dma);
        test_shared_mem = NULL;
    }
    pr_info("GSP_EMU_TEST: Test context cleaned up\n");
}

static int test_alloc_free(void)
{
    struct gsp_rpc_alloc_request req = {0};
    struct gsp_rpc_alloc_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing ALLOC/FREE...\n");

    req.header.type = GSP_RPC_ALLOC;
    req.header.seq_id = get_next_seq_id();
    req.header.flags = 0;
    req.size = TEST_ALLOC_SIZE;
    req.alignment = PAGE_SIZE;
    req.flags = GSP_ALLOC_FLAG_VRAM;

    ret = gsp_rpc_alloc_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: ALLOC failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: ALLOC returned status=%d\n", resp.status);
        return -1;
    }
    if (resp.handle == 0) {
        pr_err("GSP_EMU_TEST: ALLOC returned zero handle\n");
        return -1;
    }

    pr_info("GSP_EMU_TEST: ALLOC succeeded, handle=0x%llx\n", resp.handle);

    struct gsp_rpc_free_request free_req = {0};
    struct gsp_rpc_free_response free_resp = {0};

    free_req.header.type = GSP_RPC_FREE;
    free_req.header.seq_id = get_next_seq_id();
    free_req.handle = resp.handle;

    ret = gsp_rpc_free_emulation(test_client, &free_req, &free_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: FREE failed with ret=%d\n", ret);
        return -1;
    }
    if (free_resp.status != 0) {
        pr_err("GSP_EMU_TEST: FREE returned status=%d\n", free_resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: FREE succeeded\n");
    return 0;
}

static int test_map_unmap(void)
{
    struct gsp_rpc_alloc_request alloc_req = {0};
    struct gsp_rpc_alloc_response alloc_resp = {0};
    struct gsp_rpc_map_request map_req = {0};
    struct gsp_rpc_map_response map_resp = {0};
    struct gsp_rpc_unmap_request unmap_req = {0};
    struct gsp_rpc_unmap_response unmap_resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing MAP/UNMAP...\n");

    alloc_req.header.type = GSP_RPC_ALLOC;
    alloc_req.header.seq_id = get_next_seq_id();
    alloc_req.size = TEST_MAP_SIZE;
    alloc_req.alignment = PAGE_SIZE;
    alloc_req.flags = GSP_ALLOC_FLAG_VRAM;

    ret = gsp_rpc_alloc_emulation(test_client, &alloc_req, &alloc_resp);
    if (ret != 0 || alloc_resp.status != 0) {
        pr_err("GSP_EMU_TEST: Pre-alloc for MAP failed\n");
        return -1;
    }

    map_req.header.type = GSP_RPC_MAP;
    map_req.header.seq_id = get_next_seq_id();
    map_req.handle = alloc_resp.handle;
    map_req.offset = 0;
    map_req.size = TEST_MAP_SIZE;
    map_req.gpu_va = 0x100000000ULL;
    map_req.flags = GSP_MAP_FLAG_READ | GSP_MAP_FLAG_WRITE;

    ret = gsp_rpc_map_emulation(test_client, &map_req, &map_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MAP failed with ret=%d\n", ret);
        return -1;
    }
    if (map_resp.status != 0) {
        pr_err("GSP_EMU_TEST: MAP returned status=%d\n", map_resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: MAP succeeded, gpu_va=0x%llx\n", map_resp.gpu_va);

    unmap_req.header.type = GSP_RPC_UNMAP;
    unmap_req.header.seq_id = get_next_seq_id();
    unmap_req.gpu_va = map_resp.gpu_va;
    unmap_req.size = TEST_MAP_SIZE;

    ret = gsp_rpc_unmap_emulation(test_client, &unmap_req, &unmap_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: UNMAP failed with ret=%d\n", ret);
        return -1;
    }
    if (unmap_resp.status != 0) {
        pr_err("GSP_EMU_TEST: UNMAP returned status=%d\n", unmap_resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: UNMAP succeeded\n");

    struct gsp_rpc_free_request free_req = {0};
    struct gsp_rpc_free_response free_resp = {0};
    free_req.header.type = GSP_RPC_FREE;
    free_req.header.seq_id = get_next_seq_id();
    free_req.handle = alloc_resp.handle;
    gsp_rpc_free_emulation(test_client, &free_req, &free_resp);

    return 0;
}

static int test_pin_unpin(void)
{
    struct gsp_rpc_alloc_request alloc_req = {0};
    struct gsp_rpc_alloc_response alloc_resp = {0};
    struct gsp_rpc_map_request map_req = {0};
    struct gsp_rpc_map_response map_resp = {0};
    struct gsp_rpc_pin_request pin_req = {0};
    struct gsp_rpc_pin_response pin_resp = {0};
    struct gsp_rpc_unpin_request unpin_req = {0};
    struct gsp_rpc_unpin_response unpin_resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing PIN/UNPIN...\n");

    alloc_req.header.type = GSP_RPC_ALLOC;
    alloc_req.header.seq_id = get_next_seq_id();
    alloc_req.size = TEST_ALLOC_SIZE;
    alloc_req.alignment = PAGE_SIZE;
    alloc_req.flags = GSP_ALLOC_FLAG_VRAM;

    ret = gsp_rpc_alloc_emulation(test_client, &alloc_req, &alloc_resp);
    if (ret != 0 || alloc_resp.status != 0) {
        pr_err("GSP_EMU_TEST: Pre-alloc for PIN failed\n");
        return -1;
    }

    map_req.header.type = GSP_RPC_MAP;
    map_req.header.seq_id = get_next_seq_id();
    map_req.handle = alloc_resp.handle;
    map_req.offset = 0;
    map_req.size = TEST_ALLOC_SIZE;
    map_req.gpu_va = 0x200000000ULL;
    map_req.flags = GSP_MAP_FLAG_READ | GSP_MAP_FLAG_WRITE;

    ret = gsp_rpc_map_emulation(test_client, &map_req, &map_resp);
    if (ret != 0 || map_resp.status != 0) {
        pr_err("GSP_EMU_TEST: Pre-map for PIN failed\n");
        return -1;
    }

    pin_req.header.type = GSP_RPC_PIN;
    pin_req.header.seq_id = get_next_seq_id();
    pin_req.gpu_va = map_resp.gpu_va;
    pin_req.size = TEST_ALLOC_SIZE;
    pin_req.flags = GSP_PIN_FLAG_DMA;

    ret = gsp_rpc_pin_emulation(test_client, &pin_req, &pin_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: PIN failed with ret=%d\n", ret);
        return -1;
    }
    if (pin_resp.status != 0) {
        pr_err("GSP_EMU_TEST: PIN returned status=%d\n", pin_resp.status);
        return -1;
    }
    if (pin_resp.dma_addr == 0) {
        pr_err("GSP_EMU_TEST: PIN returned zero DMA address\n");
        return -1;
    }

    pr_info("GSP_EMU_TEST: PIN succeeded, dma_addr=0x%llx\n", pin_resp.dma_addr);

    unpin_req.header.type = GSP_RPC_UNPIN;
    unpin_req.header.seq_id = get_next_seq_id();
    unpin_req.gpu_va = map_resp.gpu_va;
    unpin_req.size = TEST_ALLOC_SIZE;

    ret = gsp_rpc_unpin_emulation(test_client, &unpin_req, &unpin_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: UNPIN failed with ret=%d\n", ret);
        return -1;
    }
    if (unpin_resp.status != 0) {
        pr_err("GSP_EMU_TEST: UNPIN returned status=%d\n", unpin_resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: UNPIN succeeded\n");

    struct gsp_rpc_unmap_request unmap_req = {0};
    struct gsp_rpc_unmap_response unmap_resp = {0};
    unmap_req.header.type = GSP_RPC_UNMAP;
    unmap_req.header.seq_id = get_next_seq_id();
    unmap_req.gpu_va = map_resp.gpu_va;
    unmap_req.size = TEST_ALLOC_SIZE;
    gsp_rpc_unmap_emulation(test_client, &unmap_req, &unmap_resp);

    struct gsp_rpc_free_request free_req = {0};
    struct gsp_rpc_free_response free_resp = {0};
    free_req.header.type = GSP_RPC_FREE;
    free_req.header.seq_id = get_next_seq_id();
    free_req.handle = alloc_resp.handle;
    gsp_rpc_free_emulation(test_client, &free_req, &free_resp);

    return 0;
}

static int test_schedule_yield(void)
{
    struct gsp_rpc_schedule_request req = {0};
    struct gsp_rpc_schedule_response resp = {0};
    struct gsp_rpc_yield_request yield_req = {0};
    struct gsp_rpc_yield_response yield_resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing SCHEDULE/YIELD...\n");

    req.header.type = GSP_RPC_SCHEDULE;
    req.header.seq_id = get_next_seq_id();
    req.work_id = 0x1234;
    req.priority = 10;
    req.flags = 0;

    ret = gsp_rpc_schedule_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: SCHEDULE failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: SCHEDULE returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: SCHEDULE succeeded\n");

    yield_req.header.type = GSP_RPC_YIELD;
    yield_req.header.seq_id = get_next_seq_id();
    yield_req.work_id = 0x1234;

    ret = gsp_rpc_yield_emulation(test_client, &yield_req, &yield_resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: YIELD failed with ret=%d\n", ret);
        return -1;
    }
    if (yield_resp.status != 0) {
        pr_err("GSP_EMU_TEST: YIELD returned status=%d\n", yield_resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: YIELD succeeded\n");
    return 0;
}

static int test_bar_trap(void)
{
    struct gsp_rpc_bar_trap_request req = {0};
    struct gsp_rpc_bar_trap_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing BAR_TRAP...\n");

    req.header.type = GSP_RPC_BAR_TRAP;
    req.header.seq_id = get_next_seq_id();
    req.bar_index = 0;
    req.offset = 0x1000;
    req.size = 4;
    req.is_write = false;
    req.value = 0;

    ret = gsp_rpc_bar_trap_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: BAR_TRAP read failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: BAR_TRAP read returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: BAR_TRAP read succeeded, value=0x%llx\n", resp.value);

    req.header.seq_id = get_next_seq_id();
    req.is_write = true;
    req.value = 0xDEADBEEF;

    ret = gsp_rpc_bar_trap_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: BAR_TRAP write failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: BAR_TRAP write returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: BAR_TRAP write succeeded\n");

    req.header.seq_id = get_next_seq_id();
    req.is_write = false;
    req.value = 0;

    ret = gsp_rpc_bar_trap_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: BAR_TRAP read-back failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.value != 0xDEADBEEF) {
        pr_err("GSP_EMU_TEST: BAR_TRAP read-back mismatch: expected 0xDEADBEEF, got 0x%llx\n", resp.value);
        return -1;
    }

    pr_info("GSP_EMU_TEST: BAR_TRAP read-back verified\n");
    return 0;
}

static int test_notify(void)
{
    struct gsp_rpc_notify_request req = {0};
    struct gsp_rpc_notify_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing NOTIFY...\n");

    req.header.type = GSP_RPC_NOTIFY;
    req.header.seq_id = get_next_seq_id();
    req.event_type = GSP_NOTIFY_EVENT_COMPLETION;
    req.flags = GSP_NOTIFY_FLAG_ENABLE;

    ret = gsp_rpc_notify_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: NOTIFY register failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: NOTIFY register returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: NOTIFY register succeeded\n");

    req.header.seq_id = get_next_seq_id();
    req.flags = GSP_NOTIFY_FLAG_DISABLE;

    ret = gsp_rpc_notify_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: NOTIFY unregister failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: NOTIFY unregister returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: NOTIFY unregister succeeded\n");
    return 0;
}

static int test_migrate(void)
{
    struct gsp_rpc_migrate_request req = {0};
    struct gsp_rpc_migrate_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing MIGRATE...\n");

    req.header.type = GSP_RPC_MIGRATE;
    req.header.seq_id = get_next_seq_id();
    req.phase = GSP_MIGRATE_QUERY;
    req.flags = 0;
    req.data = 0;
    req.data_size = 0;

    ret = gsp_rpc_migrate_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE QUERY failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE QUERY returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: MIGRATE QUERY succeeded\n");

    req.header.seq_id = get_next_seq_id();
    req.phase = GSP_MIGRATE_PREPARE;

    ret = gsp_rpc_migrate_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE PREPARE failed with ret=%d\n", ret);
        return -1;
    }

    req.header.seq_id = get_next_seq_id();
    req.phase = GSP_MIGRATE_TRANSFER;

    ret = gsp_rpc_migrate_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE TRANSFER failed with ret=%d\n", ret);
        return -1;
    }

    req.header.seq_id = get_next_seq_id();
    req.phase = GSP_MIGRATE_COMMIT;

    ret = gsp_rpc_migrate_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE COMMIT failed with ret=%d\n", ret);
        return -1;
    }

    req.header.seq_id = get_next_seq_id();
    req.phase = GSP_MIGRATE_CLEANUP;

    ret = gsp_rpc_migrate_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: MIGRATE CLEANUP failed with ret=%d\n", ret);
        return -1;
    }

    pr_info("GSP_EMU_TEST: MIGRATE 4-phase completed\n");
    return 0;
}

static int test_debug(void)
{
    struct gsp_rpc_debug_request req = {0};
    struct gsp_rpc_debug_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing DEBUG...\n");

    req.header.type = GSP_RPC_DEBUG;
    req.header.seq_id = get_next_seq_id();
    req.sub_command = GSP_DEBUG_GET_VERSION;
    req.data = 0;
    req.data_size = 0;

    ret = gsp_rpc_debug_emulation(test_client, &req, &resp);
    if (ret != 0) {
        pr_err("GSP_EMU_TEST: DEBUG failed with ret=%d\n", ret);
        return -1;
    }
    if (resp.status != 0) {
        pr_err("GSP_EMU_TEST: DEBUG returned status=%d\n", resp.status);
        return -1;
    }

    pr_info("GSP_EMU_TEST: DEBUG succeeded, version=0x%llx\n", resp.data);
    return 0;
}

static int test_concurrent_rpcs(void)
{
    struct gsp_rpc_alloc_request req = {0};
    struct gsp_rpc_alloc_response resp = {0};
    int ret, passed = 0, failed = 0;

    pr_info("GSP_EMU_TEST: Testing concurrent RPCs...\n");

    for (int i = 0; i < 10; i++) {
        req.header.type = GSP_RPC_ALLOC;
        req.header.seq_id = get_next_seq_id();
        req.size = 4096;
        req.alignment = PAGE_SIZE;
        req.flags = GSP_ALLOC_FLAG_VRAM;

        ret = gsp_rpc_alloc_emulation(test_client, &req, &resp);
        if (ret == 0 && resp.status == 0 && resp.handle != 0) {
            passed++;
            struct gsp_rpc_free_request free_req = {0};
            struct gsp_rpc_free_response free_resp = {0};
            free_req.header.type = GSP_RPC_FREE;
            free_req.header.seq_id = get_next_seq_id();
            free_req.handle = resp.handle;
            gsp_rpc_free_emulation(test_client, &free_req, &free_resp);
        } else {
            failed++;
        }
    }

    pr_info("GSP_EMU_TEST: Concurrent test: %d passed, %d failed\n", passed, failed);
    return (failed == 0) ? 0 : -1;
}

static int test_error_handling(void)
{
    struct gsp_rpc_alloc_request req = {0};
    struct gsp_rpc_alloc_response resp = {0};
    int ret;

    pr_info("GSP_EMU_TEST: Testing error handling...\n");

    req.header.type = GSP_RPC_ALLOC;
    req.header.seq_id = 0;
    req.size = 4096;
    req.alignment = PAGE_SIZE;
    req.flags = GSP_ALLOC_FLAG_VRAM;

    ret = gsp_rpc_alloc_emulation(test_client, &req, &resp);
    if (ret == 0 && resp.status != 0) {
        pr_info("GSP_EMU_TEST: Zero seq_id correctly rejected\n");
    } else {
        pr_err("GSP_EMU_TEST: Zero seq_id not handled correctly\n");
        return -1;
    }

    req.header.seq_id = get_next_seq_id();
    req.size = 0;

    ret = gsp_rpc_alloc_emulation(test_client, &req, &resp);
    if (ret == 0 && resp.status != 0) {
        pr_info("GSP_EMU_TEST: Zero size correctly rejected\n");
    } else {
        pr_err("GSP_EMU_TEST: Zero size not handled correctly\n");
        return -1;
    }

    req.header.seq_id = get_next_seq_id();
    req.size = 4096;
    req.alignment = 123;

    ret = gsp_rpc_alloc_emulation(test_client, &req, &resp);
    if (ret == 0 && resp.status != 0) {
        pr_info("GSP_EMU_TEST: Invalid alignment correctly rejected\n");
    } else {
        pr_err("GSP_EMU_TEST: Invalid alignment not handled correctly\n");
        return -1;
    }

    return 0;
}

static int run_all_tests(void)
{
    int failed = 0;

    init_test_context();
    if (!test_ctx || !test_client) {
        pr_err("GSP_EMU_TEST: Failed to initialize test context\n");
        return -1;
    }

    if (test_alloc_free() != 0) failed++;
    if (test_map_unmap() != 0) failed++;
    if (test_pin_unpin() != 0) failed++;
    if (test_schedule_yield() != 0) failed++;
    if (test_bar_trap() != 0) failed++;
    if (test_notify() != 0) failed++;
    if (test_migrate() != 0) failed++;
    if (test_debug() != 0) failed++;
    if (test_concurrent_rpcs() != 0) failed++;
    if (test_error_handling() != 0) failed++;

    cleanup_test_context();

    if (failed == 0) {
        pr_info("GSP_EMU_TEST: ALL TESTS PASSED\n");
        return 0;
    } else {
        pr_err("GSP_EMU_TEST: %d TEST(S) FAILED\n", failed);
        return -1;
    }
}

static int __init gsp_emulation_test_init(void)
{
    pr_info("GSP_EMU_TEST: Starting GSP emulation unit tests...\n");
    return run_all_tests();
}

static void __exit gsp_emulation_test_exit(void)
{
    pr_info("GSP_EMU_TEST: Module unloaded\n");
}

module_init(gsp_emulation_test_init);
module_exit(gsp_emulation_test_exit);