/*
 * Test for guest vGPU driver initialization
 * This test verifies that the guest vGPU driver can initialize properly
 * when passed through to a VM via VFIO/mdev
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/vfio.h>
#include <linux/mdev.h>
#include <errno.h>

#define NV_VGPU_IOCTL_MAGIC 'N'
#define NV_VGPU_IOCTL_GET_INFO \
	_IOWR(NV_VGPU_IOCTL_MAGIC, 0x01, struct nv_vgpu_ioctl_get_info)
#define NV_VGPU_IOCTL_MAP_BAR \
	_IOWR(NV_VGPU_IOCTL_MAGIC, 0x02, struct nv_vgpu_ioctl_map_bar)
#define NV_VGPU_IOCTL_GSP_RPC \
	_IOWR(NV_VGPU_IOCTL_MAGIC, 0x06, struct nv_vgpu_ioctl_gsp_rpc)

struct nv_vgpu_ioctl_get_info {
    uint32_t vgpu_type;
    uint32_t vgpu_id;
    uint32_t gpu_instance_id;
    uint32_t compute_instance_id;
    uint32_t state;
    uint64_t fb_base;
    uint64_t fb_size;
    uint64_t bar0_config_base;
    uint32_t num_bars;
    struct {
        uint32_t index;
        uint64_t gpa;
        uint64_t size;
        uint64_t iova;
        bool mapped;
        bool trapped;
    } bars[8];
};

struct nv_vgpu_ioctl_map_bar {
    uint32_t bar_index;
    uint64_t gpa;
    uint64_t size;
    uint64_t iova;
    uint32_t flags;
};

struct nv_vgpu_ioctl_gsp_rpc {
    uint32_t type;
    uint32_t flags;
    uint64_t data_ptr;
    uint32_t data_size;
    uint32_t result;
};

int main(int argc, char *argv[])
{
    int ret = 0;
    int vfio_fd = -1;
    int mdev_fd = -1;
    char *mdev_path = NULL;
    struct nv_vgpu_ioctl_get_info info = {0};
    struct nv_vgpu_ioctl_map_bar map_bar = {0};
    struct vfio_device_info dev_info = {0};

    printf("Testing guest vGPU driver initialization...\n");

    if (argc < 2) {
        printf("Usage: %s <mdev_uuid>\n", argv[0]);
        printf("Example: %s 00000000-0000-0000-0000-000000000000\n", argv[0]);
        return 1;
    }

    mdev_path = argv[1];

    vfio_fd = open("/dev/vfio/vfio", O_RDWR);
    if (vfio_fd < 0) {
        printf("FAIL: Cannot open /dev/vfio/vfio: %s\n", strerror(errno));
        return 1;
    }

    dev_info.argsz = sizeof(dev_info);
    if (ioctl(vfio_fd, VFIO_GET_API_VERSION, &dev_info) < 0) {
        printf("FAIL: VFIO_GET_API_VERSION failed: %s\n", strerror(errno));
        ret = 1;
        goto cleanup;
    }

    if (dev_info.version != VFIO_API_VERSION) {
        printf("FAIL: VFIO API version mismatch\n");
        ret = 1;
        goto cleanup;
    }

    printf("PASS: VFIO API version %d\n", dev_info.version);

    char mdev_dev_path[256];
    snprintf(mdev_dev_path, sizeof(mdev_dev_path), "/dev/vfio/%s", mdev_path);

    mdev_fd = open(mdev_dev_path, O_RDWR);
    if (mdev_fd < 0) {
        printf("FAIL: Cannot open mdev device %s: %s\n", mdev_dev_path, strerror(errno));
        ret = 1;
        goto cleanup;
    }

    printf("PASS: Opened mdev device %s\n", mdev_dev_path);

    dev_info.argsz = sizeof(dev_info);
    if (ioctl(mdev_fd, VFIO_DEVICE_GET_INFO, &dev_info) < 0) {
        printf("FAIL: VFIO_DEVICE_GET_INFO failed: %s\n", strerror(errno));
        ret = 1;
        goto cleanup;
    }

    printf("PASS: Device info - flags: 0x%x, num_regions: %d, num_irqs: %d\n",
           dev_info.flags, dev_info.num_regions, dev_info.num_irqs);

    if (ioctl(mdev_fd, NV_VGPU_IOCTL_GET_INFO, &info) < 0) {
        printf("FAIL: NV_VGPU_IOCTL_GET_INFO failed: %s\n", strerror(errno));
        ret = 1;
        goto cleanup;
    }

    printf("PASS: vGPU info - type: %d, id: %d, state: %d\n",
           info.vgpu_type, info.vgpu_id, info.state);
    printf("       fb_base: 0x%lx, fb_size: 0x%lx\n", info.fb_base, info.fb_size);
    printf("       bar0_config_base: 0x%lx\n", info.bar0_config_base);
    printf("       num_bars: %d\n", info.num_bars);

    for (int i = 0; i < info.num_bars; i++) {
        printf("       BAR%d: gpa=0x%lx, size=0x%lx, iova=0x%lx, mapped=%d, trapped=%d\n",
               info.bars[i].index, info.bars[i].gpa, info.bars[i].size,
               info.bars[i].iova, info.bars[i].mapped, info.bars[i].trapped);
    }

    if (info.num_bars > 0) {
        map_bar.bar_index = 0;
        map_bar.gpa = info.bars[0].gpa;
        map_bar.size = info.bars[0].size;
        map_bar.flags = 0x1;

        if (ioctl(mdev_fd, NV_VGPU_IOCTL_MAP_BAR, &map_bar) < 0) {
            printf("FAIL: NV_VGPU_IOCTL_MAP_BAR failed: %s\n", strerror(errno));
            ret = 1;
            goto cleanup;
        }

        printf("PASS: Mapped BAR0 at IOVA 0x%lx\n", map_bar.iova);
    }

    struct nv_vgpu_ioctl_gsp_rpc rpc = {0};
    rpc.type = 0x01;
    rpc.data_size = 0;

    if (ioctl(mdev_fd, NV_VGPU_IOCTL_GSP_RPC, &rpc) < 0) {
        printf("WARN: NV_VGPU_IOCTL_GSP_RPC failed: %s (may not be implemented yet)\n", strerror(errno));
    } else {
        printf("PASS: GSP RPC test completed\n");
    }

    printf("\nAll guest initialization tests PASSED\n");

cleanup:
    if (mdev_fd >= 0)
        close(mdev_fd);
    if (vfio_fd >= 0)
        close(vfio_fd);

    return ret;
}