/*
 * GSP vGPU-IO Userspace Test
 * Tests VFIO interface and driver functionality from userspace
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/vfio.h>
#include <errno.h>
#include <stdint.h>
#include <inttypes.h>

#define GSP_VGPU_DEVICE "/dev/vfio/gsp-vgpu"

#define TEST_PASS 0
#define TEST_FAIL 1

static int test_count = 0;
static int pass_count = 0;
static int fail_count = 0;

#define TEST_ASSERT(cond, msg) \
    do { \
        test_count++; \
        if (cond) { \
            pass_count++; \
            printf("PASS: %s\n", msg); \
        } else { \
            fail_count++; \
            fprintf(stderr, "FAIL: %s\n", msg); \
        } \
    } while (0)

#define TEST_ASSERT_EQ(a, b, msg) \
    TEST_ASSERT((a) == (b), msg " (expected " #b ", got " #a ")")

/* IOCTL definitions matching kernel driver */
#define GSP_IOCTL_MAGIC 'G'
#define GSP_IOCTL_CREATE_VF _IOW(GSP_IOCTL_MAGIC, 0x01, struct gsp_vf_config)
#define GSP_IOCTL_DESTROY_VF _IOW(GSP_IOCTL_MAGIC, 0x02, uint32_t)
#define GSP_IOCTL_ALLOC _IOWR(GSP_IOCTL_MAGIC, 0x10, struct gsp_alloc_args)
#define GSP_IOCTL_FREE _IOW(GSP_IOCTL_MAGIC, 0x11, struct gsp_free_args)
#define GSP_IOCTL_MAP _IOWR(GSP_IOCTL_MAGIC, 0x12, struct gsp_map_args)
#define GSP_IOCTL_UNMAP _IOW(GSP_IOCTL_MAGIC, 0x13, struct gsp_unmap_args)
#define GSP_IOCTL_PIN _IOWR(GSP_IOCTL_MAGIC, 0x14, struct gsp_pin_args)
#define GSP_IOCTL_UNPIN _IOW(GSP_IOCTL_MAGIC, 0x15, struct gsp_unpin_args)
#define GSP_IOCTL_BAR_TRAP _IOWR(GSP_IOCTL_MAGIC, 0x16, struct gsp_bar_trap_args)
#define GSP_IOCTL_MIGRATE _IOWR(GSP_IOCTL_MAGIC, 0x17, struct gsp_migrate_args)
#define GSP_IOCTL_WPR_PROTECT _IOWR(GSP_IOCTL_MAGIC, 0x18, struct gsp_wpr_args)
#define GSP_IOCTL_WPR_UNPROTECT _IOW(GSP_IOCTL_MAGIC, 0x19, struct gsp_wpr_args)

struct gsp_vf_config {
    uint32_t id;
    uint64_t bar1_offset;
    uint64_t bar1_size;
    uint64_t pt_base;
    uint64_t pt_size;
    uint64_t cmd_base;
    uint64_t cmd_size;
    uint32_t max_memory_mb;
    uint32_t num_engines;
};

struct gsp_alloc_args {
    uint32_t vf_id;
    uint64_t size;
    uint32_t type;
    uint64_t alignment;
    uint64_t va_hint;
    uint32_t flags;
    uint32_t pid;
    uint64_t handle;
    uint64_t va;
    uint64_t actual_size;
};

struct gsp_free_args {
    uint32_t vf_id;
    uint64_t handle;
};

struct gsp_map_args {
    uint32_t vf_id;
    uint64_t handle;
    uint64_t gpu_va;
    uint64_t size;
    uint32_t prot;
    uint32_t flags;
    uint64_t offset;
    uint64_t mapped_va;
};

struct gsp_unmap_args {
    uint32_t vf_id;
    uint64_t handle;
    uint64_t gpu_va;
    uint64_t size;
};

struct gsp_pin_args {
    uint32_t vf_id;
    uint64_t handle;
    uint64_t offset;
    uint64_t size;
    uint32_t flags;
    uint64_t dma_addr;
    uint32_t num_pages;
};

struct gsp_unpin_args {
    uint32_t vf_id;
    uint64_t handle;
    uint64_t offset;
    uint64_t size;
};

struct gsp_bar_trap_args {
    uint32_t vf_id;
    uint32_t bar_index;
    uint32_t trap_type;
    uint64_t offset;
    uint64_t value;
    uint32_t size;
    uint64_t ret_value;
};

struct gsp_migrate_args {
    uint32_t vf_id;
    uint32_t phase;
    uint32_t subcmd;
    uint64_t src_vf;
    uint64_t dst_vf;
    uint64_t data_ptr;
    uint64_t data_size;
    uint32_t flags;
    uint64_t token;
};

struct gsp_wpr_args {
    uint32_t vf_id;
    uint64_t base;
    uint64_t size;
    uint32_t flags;
};

int main(int argc, char *argv[])
{
    int fd, ret;
    uint32_t vf_id = 1;

    printf("========================================\n");
    printf("GSP vGPU-IO Userspace Test Suite\n");
    printf("========================================\n");

    /* Open device */
    fd = open(GSP_VGPU_DEVICE, O_RDWR);
    if (fd < 0) {
        fprintf(stderr, "Failed to open %s: %s\n", GSP_VGPU_DEVICE, strerror(errno));
        fprintf(stderr, "Make sure driver is loaded and VFIO device created\n");
        return TEST_FAIL;
    }
    printf("Opened device: %s (fd=%d)\n", GSP_VGPU_DEVICE, fd);

    /* Test 1: Create VF */
    {
        struct gsp_vf_config config = {
            .id = vf_id,
            .bar1_offset = 0x10000000,
            .bar1_size = 0x100000,
            .pt_base = 0x20000000,
            .pt_size = 0x100000,
            .cmd_base = 0x30000000,
            .cmd_size = 0x10000,
            .max_memory_mb = 1024,
            .num_engines = 4,
        };

        ret = ioctl(fd, GSP_IOCTL_CREATE_VF, &config);
        TEST_ASSERT(ret == 0, "Create VF");
        if (ret == 0)
            printf("  VF %u created successfully\n", vf_id);
    }

    /* Test 2: Allocate memory */
    {
        struct gsp_alloc_args args = {
            .vf_id = vf_id,
            .size = 0x10000,  // 64KB
            .type = 0,        // VRAM
            .alignment = 0x1000,
            .va_hint = 0,
            .flags = 0,
            .pid = getpid(),
        };

        ret = ioctl(fd, GSP_IOCTL_ALLOC, &args);
        TEST_ASSERT(ret == 0, "Allocate memory");
        if (ret == 0) {
            printf("  Allocated: handle=0x%" PRIx64 ", va=0x%" PRIx64 ", size=0x%" PRIx64 "\n",
                   args.handle, args.va, args.actual_size);
        }
    }

    /* Test 3: Pin memory */
    {
        struct gsp_pin_args args = {
            .vf_id = vf_id,
            .handle = 0,  // Would use handle from alloc
            .offset = 0,
            .size = 0x1000,
            .flags = 3,  // READ | WRITE
        };

        ret = ioctl(fd, GSP_IOCTL_PIN, &args);
        TEST_ASSERT(ret == 0 || errno == ENODEV, "Pin memory (may fail without real allocation)");
    }

    /* Test 4: BAR Trap */
    {
        struct gsp_bar_trap_args args = {
            .vf_id = vf_id,
            .bar_index = 1,
            .trap_type = 1,  // WRITE
            .offset = 0x1000,
            .value = 0xDEADBEEF,
            .size = 4,
        };

        ret = ioctl(fd, GSP_IOCTL_BAR_TRAP, &args);
        TEST_ASSERT(ret == 0 || errno == ENODEV, "BAR trap (may fail without real HW)");
    }

    /* Test 5: WPR Protect */
    {
        struct gsp_wpr_args args = {
            .vf_id = vf_id,
            .base = 0x10000000,
            .size = 0x10000,
            .flags = 3,  // WRITE | READ
        };

        ret = ioctl(fd, GSP_IOCTL_WPR_PROTECT, &args);
        TEST_ASSERT(ret == 0 || errno == ENODEV, "WPR protect (may fail without real HW)");
    }

    /* Test 6: Migration */
    {
        struct gsp_migrate_args args = {
            .vf_id = vf_id,
            .phase = 1,  // PREPARE
            .subcmd = 3,
            .src_vf = vf_id,
            .dst_vf = 2,
            .flags = 1,  // LIVE
        };

        ret = ioctl(fd, GSP_IOCTL_MIGRATE, &args);
        TEST_ASSERT(ret == 0 || errno == ENODEV, "Migration (may fail without real HW)");
    }

    /* Test 7: Destroy VF */
    {
        ret = ioctl(fd, GSP_IOCTL_DESTROY_VF, &vf_id);
        TEST_ASSERT(ret == 0, "Destroy VF");
        if (ret == 0)
            printf("  VF %u destroyed successfully\n", vf_id);
    }

    close(fd);

    printf("========================================\n");
    printf("Userspace Test Results: %d total, %d passed, %d failed\n",
           test_count, pass_count, fail_count);
    printf("========================================\n");

    return (fail_count > 0) ? TEST_FAIL : TEST_PASS;
}