#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/pci.h>
#include <linux/vfio.h>
#include <linux/iommufd.h>
#include <linux/vfio_pci_core.h>
#include <linux/mdev.h>
#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/list.h>
#include <linux/workqueue.h>
#include <linux/device.h>
#include <linux/dma-mapping.h>
#include <linux/iommu.h>
#include <linux/vmalloc.h>
#include <linux/version.h>
#include <linux/delay.h>

#include "nvidia-vgpu-vfio.h"
#include "bar_emulation.h"
#include "gsp_comm.h"
#include "gsp_emulation.h"
#include "memory.h"
#include "irq.h"
#include "migration.h"
#include "ioctl.h"

#include "nv-vgpu-vfio-interface.h"

#define NV_VERSION_STRING "550.163.01"

/* ============================================================
 * Forward declarations
 * ============================================================ */

static int nv_vgpu_mdev_create(struct mdev_device *mdev);
static void nv_vgpu_mdev_remove(struct mdev_device *mdev);
static void nv_vgpu_device_init_work(struct work_struct *work);
static void nv_vgpu_device_cleanup_work(struct work_struct *work);

/* VFIO device operations */
static int nv_vgpu_vfio_dev_match(struct vfio_device *vdev, char *buf);
static int nv_vgpu_vfio_dev_init(struct vfio_device *vdev);
static void nv_vgpu_vfio_dev_release(struct vfio_device *vdev);
static int nv_vgpu_vfio_dev_open(struct vfio_device *vdev);
static void nv_vgpu_vfio_dev_close(struct vfio_device *vdev);
static ssize_t nv_vgpu_vfio_dev_read(struct vfio_device *vdev, char __user *buf, size_t count, loff_t *ppos);
static ssize_t nv_vgpu_vfio_dev_write(struct vfio_device *vdev, const char __user *buf, size_t count, loff_t *ppos);
static long nv_vgpu_vfio_dev_ioctl(struct vfio_device *vdev, unsigned int cmd, unsigned long arg);
static int nv_vgpu_vfio_dev_mmap(struct vfio_device *vdev, struct vm_area_struct *vma);
static void nv_vgpu_vfio_dev_dma_unmap(struct vfio_device *vdev, u64 iova, u64 length);

/* PCI Config Space handlers for VFIO mdev */
static int nv_vgpu_config_space_init(struct nv_vgpu_device *vgpu, struct pci_dev *pdev);
static ssize_t nv_vgpu_config_read(struct nv_vgpu_device *vgpu, char __user *buf, size_t count, loff_t *ppos);
static ssize_t nv_vgpu_config_write(struct nv_vgpu_device *vgpu, const char __user *buf, size_t count, loff_t *ppos);

/* vGPU type management */
int nv_vgpu_register_type_info(struct nv_vgpu_type_info *info);
struct nv_vgpu_type_info *nv_vgpu_find_type_info(const char *name);

struct mdev_type *nv_vgpu_mdev_types[16];
struct nv_vgpu_type_info g_vgpu_type_infos[16];
int g_num_vgpu_types = 0;
static DEFINE_MUTEX(g_vgpu_type_lock);

/*
 * get_available - report available instances for an mdev type.
 *
 * The mdev core's shared parent->available_instances counter can drift
 * (e.g. if a remove path ever fails to complete). Compute availability
 * from our own authoritative device list instead.
 */
static unsigned int nv_vgpu_get_available(struct mdev_type *mtype)
{
	struct nv_vgpu_type_info *info;
	struct nv_vgpu_device *vgpu;
	const char *name = mtype->sysfs_name;
	unsigned int active = 0, max = 1;

	/* Type dirs may carry an extra "nvidia-" prefix vs our type names */
	if (str_has_prefix(name, "nvidia-"))
		name += strlen("nvidia-");
	info = nv_vgpu_find_type_info(name);
	if (info && info->max_instances)
		max = info->max_instances;

	mutex_lock(&nv_vgpu_driver.lock);
	list_for_each_entry(vgpu, &nv_vgpu_driver.devices, list) {
		if (vgpu->mdev && vgpu->mdev->type == mtype)
			active++;
	}
	mutex_unlock(&nv_vgpu_driver.lock);

	return active >= max ? 0 : max - active;
}

struct mdev_driver nv_vgpu_mdev_driver = {
	.device_api = VFIO_DEVICE_API_PCI_STRING,
	.max_instances = 8,
	.probe = nv_vgpu_mdev_create,
	.remove = nv_vgpu_mdev_remove,
	.get_available = nv_vgpu_get_available,
	.driver = {
		.name = "nvidia-vgpu-vfio-mdev",
	},
};

/* VFIO device operations for vGPU */
static const struct vfio_device_ops nv_vgpu_vfio_dev_ops = {
	.name		= "nvidia-vgpu-vfio",
	.init		= nv_vgpu_vfio_dev_init,
	.release	= nv_vgpu_vfio_dev_release,
	.open_device	= nv_vgpu_vfio_dev_open,
	.close_device	= nv_vgpu_vfio_dev_close,
	.read		= nv_vgpu_vfio_dev_read,
	.write		= nv_vgpu_vfio_dev_write,
	.ioctl		= nv_vgpu_vfio_dev_ioctl,
	.mmap		= nv_vgpu_vfio_dev_mmap,
	.dma_unmap	= nv_vgpu_vfio_dev_dma_unmap,
	.match		= nv_vgpu_vfio_dev_match,
	.bind_iommufd	= vfio_iommufd_emulated_bind,
	.unbind_iommufd = vfio_iommufd_emulated_unbind,
	.attach_ioas	= vfio_iommufd_emulated_attach_ioas,
	.detach_ioas	= vfio_iommufd_emulated_detach_ioas,
};

/* ============================================================
 * Symbols we provide to nvidia.ko
 * ============================================================ */

static vgpu_vfio_ops_t *g_vgpu_vfio_ops = NULL;
static rm_vgpu_vfio_ops_t g_rm_vgpu_vfio_ops;

/* mdev parent structure for our physical device */
static struct mdev_parent *g_mdev_parent = NULL;

/* VFIO ops implementation */
static NV_STATUS nv_vgpu_vfio_probe(struct pci_dev *pdev, NvU32 vgpu_type_id, NvU32 *vgpu_id);
static void nv_vgpu_vfio_remove(struct pci_dev *pdev);
static NV_STATUS nv_vgpu_vfio_inject_interrupt(void *vgpu_dev);

static vgpu_vfio_ops_t nv_vgpu_vfio_ops = {
	.probe = nv_vgpu_vfio_probe,
	.remove = nv_vgpu_vfio_remove,
	.inject_interrupt = nv_vgpu_vfio_inject_interrupt,
};

/* RM ops implementation - callbacks from nvidia.ko */
static NV_STATUS nv_vgpu_create_callback(struct pci_dev *pdev, const NvU8 *uuid,
                                         NvU32 type_id, NvU16 *vgpu_id,
                                         NvU32 *gpu_pci_id, NvU32 gpu_pci_bdf)
{
	pr_info("nvidia-vgpu-vfio: vgpu_create callback: type_id=%u\n", type_id);
	if (vgpu_id)
		*vgpu_id = 1;
	if (gpu_pci_id)
		*gpu_pci_id = 0x10de1ed3;
	return NV_OK;
}

static NV_STATUS nv_vgpu_delete_callback(const NvU8 *uuid, NvU16 vgpu_id)
{
	pr_info("nvidia-vgpu-vfio: vgpu_delete callback: vgpu_id=%u\n", vgpu_id);
	return NV_OK;
}

static NV_STATUS nv_vgpu_bar_info_callback(struct pci_dev *pdev, const NvU8 *uuid,
                                           NvU64 *size, NvU32 bar_index, void *vgpu_dev)
{
	pr_info("nvidia-vgpu-vfio: vgpu_bar_info callback: bar_index=%u\n", bar_index);
	if (size) {
		switch (bar_index) {
		case 0:
			*size = NV_VGPU_BAR0_SIZE;  /* 16MB */
			break;
		case 1:
			*size = NV_VGPU_BAR1_SIZE;  /* 256MB */
			break;
		case 2:
			*size = NV_VGPU_BAR2_SIZE;  /* 32MB */
			break;
		case 3:
			*size = NV_VGPU_BAR3_SIZE;  /* 32MB */
			break;
		default:
			*size = 0;
			break;
		}
	}
	return NV_OK;
}

static NV_STATUS nv_vgpu_get_name_callback(struct pci_dev *pdev, NvU32 type_id, char *buf)
{
	pr_info("nvidia-vgpu-vfio: get_name callback: type_id=%u\n", type_id);
	if (buf) {
		switch (type_id) {
		case 0: /* Ampere */
			strcpy(buf, "NVIDIA A100");
			break;
		case 1: /* Hopper */
			strcpy(buf, "NVIDIA H100");
			break;
		case 2: /* Turing */
			strcpy(buf, "NVIDIA T4");
			break;
		case 3: /* Pascal */
			strcpy(buf, "NVIDIA P4");
			break;
		default:
			strcpy(buf, "NVIDIA vGPU");
			break;
		}
	}
	return NV_OK;
}

static NV_STATUS nv_vgpu_get_description_callback(struct pci_dev *pdev, NvU32 type_id, char *buf)
{
	pr_info("nvidia-vgpu-vfio: get_description callback: type_id=%u\n", type_id);
	if (buf) {
		switch (type_id) {
		case 0:
			strcpy(buf, "NVIDIA Ampere Architecture vGPU");
			break;
		case 1:
			strcpy(buf, "NVIDIA Hopper Architecture vGPU");
			break;
		case 2:
			strcpy(buf, "NVIDIA Turing Architecture vGPU");
			break;
		case 3:
			strcpy(buf, "NVIDIA Pascal Architecture vGPU");
			break;
		default:
			strcpy(buf, "NVIDIA vGPU Description");
			break;
		}
	}
	return NV_OK;
}

static NV_STATUS nv_vgpu_get_instances_callback(struct pci_dev *pdev, NvU32 type_id, char *buf)
{
	pr_info("nvidia-vgpu-vfio: get_instances callback: type_id=%u\n", type_id);
	if (buf)
		strcpy(buf, "1");
	return NV_OK;
}

static void nv_vgpu_update_request_callback(const NvU8 *uuid, NvU64 *vgpu_dev,
                                            NvU64 *vgpu_vm, VGPU_DEVICE_STATE device_state,
                                            const char *params)
{
	pr_info("nvidia-vgpu-vfio: update_request callback: state=%u\n", device_state);
}

static NV_STATUS nv_vgpu_start_callback(const NvU8 *uuid, void *wait_queue, NvS32 *return_status,
                                        NvU8 *vm_name, NvU32 qemu_pid)
{
	pr_info("nvidia-vgpu-vfio: vgpu_start callback: qemu_pid=%u\n", qemu_pid);
	if (return_status)
		*return_status = 0;
	return NV_OK;
}

static NV_STATUS nv_vgpu_sparse_mmap_callback(struct pci_dev *pdev, const NvU8 *uuid,
                                              NvU64 **pages, NvU64 **offsets, NvU32 *count)
{
	pr_info("nvidia-vgpu-vfio: sparse_mmap callback\n");
	if (count)
		*count = 0;
	return NV_OK;
}

/* Build RM ops structure */
static void populate_rm_ops(void)
{
	memset(&g_rm_vgpu_vfio_ops, 0, sizeof(g_rm_vgpu_vfio_ops));
	g_rm_vgpu_vfio_ops.version_string = NV_VERSION_STRING;
	g_rm_vgpu_vfio_ops.vgpu_create = nv_vgpu_create_callback;
	g_rm_vgpu_vfio_ops.vgpu_delete = nv_vgpu_delete_callback;
	g_rm_vgpu_vfio_ops.vgpu_bar_info = nv_vgpu_bar_info_callback;
	g_rm_vgpu_vfio_ops.vgpu_start = nv_vgpu_start_callback;
	g_rm_vgpu_vfio_ops.get_name = nv_vgpu_get_name_callback;
	g_rm_vgpu_vfio_ops.get_description = nv_vgpu_get_description_callback;
	g_rm_vgpu_vfio_ops.get_instances = nv_vgpu_get_instances_callback;
	g_rm_vgpu_vfio_ops.get_sparse_mmap = nv_vgpu_sparse_mmap_callback;
	g_rm_vgpu_vfio_ops.update_request = nv_vgpu_update_request_callback;
}

/* ============================================================
 * vGPU Type Management (userspace provides info via ioctl)
 * ============================================================ */

int nv_vgpu_register_type_info(struct nv_vgpu_type_info *info)
{
    int i;
    
    mutex_lock(&g_vgpu_type_lock);
    
    /* Check if already registered */
    for (i = 0; i < g_num_vgpu_types; i++) {
        if (strcmp(g_vgpu_type_infos[i].name, info->name) == 0) {
            /* Update existing */
            g_vgpu_type_infos[i] = *info;
            mutex_unlock(&g_vgpu_type_lock);
            return 0;
        }
    }
    
    /* Add new */
    if (g_num_vgpu_types >= 16) {
        mutex_unlock(&g_vgpu_type_lock);
        return -ENOSPC;
    }
    
    g_vgpu_type_infos[g_num_vgpu_types] = *info;
    g_num_vgpu_types++;
    
    mutex_unlock(&g_vgpu_type_lock);
    
    pr_info("Registered vGPU type: name='%s', vdev=0x%llx, pdev=0x%llx, max_inst=%u\n",
            info->name, info->vdev_id, info->pdev_id, info->max_instances);
    return 0;
}
EXPORT_SYMBOL(nv_vgpu_register_type_info);

struct nv_vgpu_type_info *nv_vgpu_find_type_info(const char *name)
{
    int i;
    
    mutex_lock(&g_vgpu_type_lock);
    for (i = 0; i < g_num_vgpu_types; i++) {
        if (strcmp(g_vgpu_type_infos[i].name, name) == 0) {
            mutex_unlock(&g_vgpu_type_lock);
            return &g_vgpu_type_infos[i];
        }
    }
    mutex_unlock(&g_vgpu_type_lock);
    return NULL;
}
EXPORT_SYMBOL(nv_vgpu_find_type_info);

static void nv_vgpu_create_static_types(void)
{
    int i;
    struct nv_vgpu_type_info static_types[4] = {
        { .name = "nvidia-vgpu-ampere", .description = "NVIDIA Ampere vGPU", .vdev_id = 0x10de2235, .pdev_id = 0x10de2235, .max_instances = 1 },
        { .name = "nvidia-vgpu-hopper", .description = "NVIDIA Hopper vGPU", .vdev_id = 0x10de2331, .pdev_id = 0x10de2331, .max_instances = 1 },
        { .name = "nvidia-vgpu-turing", .description = "NVIDIA Turing vGPU", .vdev_id = 0x10de1ed3, .pdev_id = 0x10de1ed3, .max_instances = 1 },
        { .name = "nvidia-vgpu-pascal", .description = "NVIDIA Pascal vGPU", .vdev_id = 0x10de1db1, .pdev_id = 0x10de1db1, .max_instances = 1 },
    };
    
    for (i = 0; i < 4; i++) {
        nv_vgpu_register_type_info(&static_types[i]);
    }
}

int nv_vgpu_create_mdev_types(void)
{
    int i, ret;
    
    for (i = 0; i < g_num_vgpu_types; i++) {
        nv_vgpu_mdev_types[i] = kzalloc(sizeof(struct mdev_type), GFP_KERNEL);
        if (!nv_vgpu_mdev_types[i]) {
            ret = -ENOMEM;
            goto err_types;
        }
        nv_vgpu_mdev_types[i]->sysfs_name = kstrdup(g_vgpu_type_infos[i].name, GFP_KERNEL);
        nv_vgpu_mdev_types[i]->pretty_name = kstrdup(g_vgpu_type_infos[i].description[0] ? g_vgpu_type_infos[i].description : g_vgpu_type_infos[i].name, GFP_KERNEL);
        if (!nv_vgpu_mdev_types[i]->sysfs_name || !nv_vgpu_mdev_types[i]->pretty_name) {
            ret = -ENOMEM;
            goto err_types;
        }
    }
    
    return 0;
    
err_types:
    for (i = 0; i < g_num_vgpu_types; i++) {
        kfree(nv_vgpu_mdev_types[i]->sysfs_name);
        kfree(nv_vgpu_mdev_types[i]->pretty_name);
        kfree(nv_vgpu_mdev_types[i]);
        nv_vgpu_mdev_types[i] = NULL;
    }
    return ret;
}
EXPORT_SYMBOL(nv_vgpu_create_mdev_types);

void nv_vgpu_free_type_info(void)
{
    int i;
    for (i = 0; i < g_num_vgpu_types; i++) {
        kfree(nv_vgpu_mdev_types[i]->sysfs_name);
        kfree(nv_vgpu_mdev_types[i]->pretty_name);
        kfree(nv_vgpu_mdev_types[i]);
        nv_vgpu_mdev_types[i] = NULL;
    }
    g_num_vgpu_types = 0;
}
EXPORT_SYMBOL(nv_vgpu_free_type_info);

/* get_ops: called by nvidia-vgpu-vfio to get RM ops */
NV_STATUS nvidia_vgpu_vfio_get_ops(rm_vgpu_vfio_ops_t *ops)
{
	if (!ops)
		return NV_ERR_INVALID_ARGUMENT;
	
	populate_rm_ops();
	memcpy(ops, &g_rm_vgpu_vfio_ops, sizeof(rm_vgpu_vfio_ops_t));
	pr_info("nvidia-vgpu-vfio: get_ops provided RM ops\n");
	return NV_OK;
}
EXPORT_SYMBOL(nvidia_vgpu_vfio_get_ops);

/* set_ops: called by nvidia-vgpu-vfio to provide VFIO ops */
NV_STATUS nvidia_vgpu_vfio_set_ops(vgpu_vfio_ops_t *ops)
{
	if (!ops)
		return NV_ERR_INVALID_ARGUMENT;
	
	g_vgpu_vfio_ops = ops;
	pr_info("nvidia-vgpu-vfio: set_ops received VFIO ops\n");
	return NV_OK;
}
EXPORT_SYMBOL(nvidia_vgpu_vfio_set_ops);

/* mdev driver registration - called by nvidia.ko */
int nv_vgpu_mdev_driver_register(void)
{
	pr_info("nvidia-vgpu-vfio: mdev driver register\n");
	return mdev_register_driver(&nv_vgpu_mdev_driver);
}
EXPORT_SYMBOL(nv_vgpu_mdev_driver_register);

void nv_vgpu_mdev_driver_unregister(void)
{
	pr_info("nvidia-vgpu-vfio: mdev driver unregister\n");
	mdev_unregister_driver(&nv_vgpu_mdev_driver);
}
EXPORT_SYMBOL(nv_vgpu_mdev_driver_unregister);

/* mdev parent device registration - called by nvidia.ko */

int nv_vgpu_mdev_register_parent(struct pci_dev *pdev)
{
	struct mdev_parent *parent;
	int ret;

	pr_info("nvidia-vgpu-vfio: mdev parent register for %s\n", pci_name(pdev));

	if (g_mdev_parent) {
		pr_warn("mdev parent already registered\n");
		return 0;
	}

	parent = kzalloc(sizeof(*parent), GFP_KERNEL);
	if (!parent)
		return -ENOMEM;

	parent->dev = &pdev->dev;

	ret = mdev_register_parent(parent, &pdev->dev, &nv_vgpu_mdev_driver,
				   nv_vgpu_mdev_types, g_num_vgpu_types);
	if (ret) {
		pr_err("Failed to register mdev parent: %d\n", ret);
		kfree(parent);
		return ret;
	}

	g_mdev_parent = parent;
	pr_info("mdev parent registered successfully with %d types\n", g_num_vgpu_types);
	return 0;
}
EXPORT_SYMBOL(nv_vgpu_mdev_register_parent);

void nv_vgpu_mdev_unregister_parent(struct pci_dev *pdev)
{
	pr_info("nvidia-vgpu-vfio: mdev parent unregister for %s\n", pci_name(pdev));

	if (g_mdev_parent) {
		mdev_unregister_parent(g_mdev_parent);
		kfree(g_mdev_parent);
		g_mdev_parent = NULL;
	}
}
EXPORT_SYMBOL(nv_vgpu_mdev_unregister_parent);

MODULE_SOFTDEP("pre: nvidia");

MODULE_AUTHOR("Open-IOV Project");
MODULE_DESCRIPTION("Open-IOV vGPU VFIO Mediated Device Driver");
MODULE_LICENSE("GPL v2");
MODULE_VERSION(OPEN_IOV_VFIO_VERSION);

struct nv_vgpu_driver_data nv_vgpu_driver = {
	.lock = __MUTEX_INITIALIZER(nv_vgpu_driver.lock),
	.devices = LIST_HEAD_INIT(nv_vgpu_driver.devices),
	.device_count = 0,
};

static int nv_vgpu_mdev_create(struct mdev_device *mdev)
{
	struct nv_vgpu_device *vgpu;
	struct pci_dev *pdev;
	struct mdev_parent *parent;
	struct iommu_group *iommu_group;
	int ret;

	vgpu = vfio_alloc_device(nv_vgpu_device, vfio_dev, &mdev->dev,
				 &nv_vgpu_vfio_dev_ops);
	if (IS_ERR(vgpu))
		return PTR_ERR(vgpu);

 	vgpu->mdev = mdev;
 	vgpu->state = NV_VGPU_STATE_UNINITIALIZED;
 	mutex_init(&vgpu->lock);
 	INIT_WORK(&vgpu->init_work, nv_vgpu_device_init_work);
 	INIT_WORK(&vgpu->cleanup_work, nv_vgpu_device_cleanup_work);

 	/* Initialize VFIO device registration state */
 	vgpu->vfio_dev_registered = false;

 	/* Initialize VFIO device open/close tracking */
 	atomic_set(&vgpu->vfio_dev_open_count, 0);
 	init_waitqueue_head(&vgpu->vfio_dev_close_wait);

/* Set drvdata so we can retrieve it in remove */
	dev_set_drvdata(&mdev->dev, vgpu);

	parent = mdev->type->parent;
	if (!parent || !parent->dev) {
		ret = -ENODEV;
		goto err_put_vdev;
	}
	pdev = to_pci_dev(parent->dev);
	if (!pdev) {
		ret = -ENODEV;
		goto err_put_vdev;
	}
vgpu->pdev = pci_dev_get(pdev);

	/* Initialize PCI config space for VFIO mdev */
	ret = nv_vgpu_config_space_init(vgpu, pdev);
	if (ret) {
		nv_vgpu_err(vgpu, "PCI config space init failed: %d", ret);
		goto err_put;
	}

	mutex_lock(&nv_vgpu_driver.lock);
	list_add_tail(&vgpu->list, &nv_vgpu_driver.devices);
	nv_vgpu_driver.device_count++;
	mutex_unlock(&nv_vgpu_driver.lock);

	/* Register VFIO device (required for VFIO_DEVICE_API_PCI_STRING).
	 * This creates a new IOMMU group for the emulated device.
	 * Do NOT add mdev to parent's IOMMU group first - vfio_register_emulated_iommu_dev
	 * creates its own IOMMU group for the emulated device. */
	ret = vfio_register_emulated_iommu_dev(&vgpu->vfio_dev);
	if (ret) {
		nv_vgpu_err(vgpu, "Failed to register VFIO device: %d", ret);
		goto err_list_del;
	}
	vgpu->vfio_dev_registered = true;

	schedule_work(&vgpu->init_work);

	return 0;

 err_list_del:
 	mutex_lock(&nv_vgpu_driver.lock);
 	list_del(&vgpu->list);
 	nv_vgpu_driver.device_count--;
 	mutex_unlock(&nv_vgpu_driver.lock);
 err_put:
 	pci_dev_put(vgpu->pdev);
 err_put_vdev:
 	vfio_put_device(&vgpu->vfio_dev);
 	return ret;
}

static void nv_vgpu_mdev_remove(struct mdev_device *mdev)
{
	struct nv_vgpu_device *vgpu = dev_get_drvdata(&mdev->dev);
	long timeout;

	if (!vgpu) {
		printk(KERN_ERR "nvidia-vgpu-vfio: No vGPU device data found for mdev %s\n",
		       dev_name(&mdev->dev));
		return;
	}

	pr_info("nvidia-vgpu-vfio: mdev remove started for %s\n", dev_name(&mdev->dev));

	/* Cancel init work first */
	cancel_work_sync(&vgpu->init_work);

	/* Cancel cleanup work to prevent race */
	cancel_work_sync(&vgpu->cleanup_work);

	/* Wait for any open file descriptors to close before proceeding.
	 * The VFIO device must not be unregistered while userspace has it open,
	 * otherwise vfio_device_group_unregister can crash with NULL pointer deref. */
	pr_info("nvidia-vgpu-vfio: waiting for vfio_dev_open_count to reach 0 (current: %d)\n",
		atomic_read(&vgpu->vfio_dev_open_count));
	
	timeout = wait_event_timeout(vgpu->vfio_dev_close_wait,
				     atomic_read(&vgpu->vfio_dev_open_count) == 0,
				     msecs_to_jiffies(5000));
	if (timeout == 0) {
		pr_warn("nvidia-vgpu-vfio: timeout waiting for vfio device close (open_count=%d), proceeding anyway\n",
			atomic_read(&vgpu->vfio_dev_open_count));
	} else {
		pr_info("nvidia-vgpu-vfio: all vfio device file descriptors closed\n");
	}

	/* Mark device as being removed - actual resource cleanup happens in
	 * vfio_device_release callback which is called when refcount reaches 0
	 * after mdev framework unregisters the vfio_device. */
	vgpu->vfio_dev_registered = false;

	pr_info("nvidia-vgpu-vfio: mdev remove callback completed for %s\n", dev_name(&mdev->dev));
}

/* VFIO device operations implementations */
static int nv_vgpu_vfio_dev_init(struct vfio_device *vdev)
{
	pr_info("nvidia-vgpu-vfio: VFIO device init\n");
	return 0;
}

static void nv_vgpu_vfio_dev_release(struct vfio_device *vdev)
{
	pr_info("nvidia-vgpu-vfio: VFIO device release\n");
}

static int nv_vgpu_vfio_dev_open(struct vfio_device *vdev)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	
	pr_info("nvidia-vgpu-vfio: VFIO device open\n");
	atomic_inc(&vgpu->vfio_dev_open_count);
	return 0;
}

static void nv_vgpu_vfio_dev_close(struct vfio_device *vdev)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	
	pr_info("nvidia-vgpu-vfio: VFIO device close\n");
	atomic_dec(&vgpu->vfio_dev_open_count);
	wake_up(&vgpu->vfio_dev_close_wait);
}

static ssize_t nv_vgpu_vfio_dev_read(struct vfio_device *vdev, char __user *buf,
				     size_t count, loff_t *ppos)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	loff_t pos = *ppos;
	
	pr_info("nvidia-vgpu-vfio: VFIO device read count=%zu pos=%lld (0x%llx)\n", count, pos, pos);

	/* Config space region: VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_CONFIG_REGION_INDEX) = 0x70000000000 */
	if (pos >= 0x70000000000ULL && pos < 0x70000000000ULL + 4096) {
		loff_t cfg_pos = pos - 0x70000000000ULL;
		size_t to_read = min(count, (size_t)(4096 - cfg_pos));
		if (copy_to_user(buf, vgpu->pci_config + cfg_pos, to_read))
			return -EFAULT;
		*ppos += to_read;
		return to_read;
	}

	/* BAR0 region: VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR0_REGION_INDEX) = 0x0 */
	if (pos >= 0x0 && pos < 0x0 + (256 * 1024 * 1024)) {
		/* This would be handled by mmap, not read */
		return -EINVAL;
	}

	return 0;
}

static ssize_t nv_vgpu_vfio_dev_write(struct vfio_device *vdev, const char __user *buf,
				      size_t count, loff_t *ppos)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	loff_t pos = *ppos;
	
	pr_info("nvidia-vgpu-vfio: VFIO device write count=%zu pos=%lld (0x%llx)\n", count, pos, pos);

	/* Config space region */
	if (pos >= 0x70000000000ULL && pos < 0x70000000000ULL + 4096) {
		loff_t cfg_pos = pos - 0x70000000000ULL;
		size_t to_write = min(count, (size_t)(4096 - cfg_pos));
		if (copy_from_user(vgpu->pci_config + cfg_pos, buf, to_write))
			return -EFAULT;
		*ppos += to_write;
		return to_write;
	}

	return 0;
}

static long nv_vgpu_vfio_dev_ioctl(struct vfio_device *vdev, unsigned int cmd,
				   unsigned long arg)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	unsigned long minsz;
	int ret;
	
	pr_info("nvidia-vgpu-vfio: VFIO device ioctl cmd=0x%x\n", cmd);

	if (cmd == VFIO_DEVICE_GET_INFO) {
		struct vfio_device_info info;

		minsz = offsetofend(struct vfio_device_info, num_irqs);

		pr_info("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO minsz=%lu argsz=%u\n", minsz, info.argsz);

		if (copy_from_user(&info, (void __user *)arg, minsz)) {
			pr_err("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO copy_from_user failed\n");
			return -EFAULT;
		}

		pr_info("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO argsz=%u minsz=%lu\n", info.argsz, minsz);

		if (info.argsz < minsz) {
			pr_err("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO argsz too small\n");
			return -EINVAL;
		}

		info.flags = VFIO_DEVICE_FLAGS_PCI | VFIO_DEVICE_FLAGS_RESET;
		info.num_regions = VFIO_PCI_NUM_REGIONS;
		info.num_irqs = VFIO_PCI_NUM_IRQS;

		pr_info("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO returning flags=0x%x num_regions=%u num_irqs=%u\n",
			info.flags, info.num_regions, info.num_irqs);

		ret = copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		pr_info("nvidia-vgpu-vfio: VFIO_DEVICE_GET_INFO copy_to_user ret=%d\n", ret);
		return ret;
	}

	if (cmd == VFIO_DEVICE_GET_REGION_INFO) {
		struct vfio_region_info info;

		minsz = offsetofend(struct vfio_region_info, offset);

		if (copy_from_user(&info, (void __user *)arg, minsz))
			return -EFAULT;

		if (info.argsz < minsz)
			return -EINVAL;

		/* PCI config space region */
		if (info.index == VFIO_PCI_CONFIG_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = 4096; /* 4KB config space */
			info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR0 - GPU registers */
		if (info.index == VFIO_PCI_BAR0_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[0].size ? vgpu->bars[0].size : (256 * 1024 * 1024); /* 256MB default */
			info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR1 - VRAM (if present) */
		if (info.index == VFIO_PCI_BAR1_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[1].size ? vgpu->bars[1].size : 0;
			if (info.size) {
				info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			} else {
				info.flags = 0;
			}
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR2 - Additional BAR (if present) */
		if (info.index == VFIO_PCI_BAR2_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[2].size ? vgpu->bars[2].size : 0;
			if (info.size) {
				info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			} else {
				info.flags = 0;
			}
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR3 - Additional BAR (if present) */
		if (info.index == VFIO_PCI_BAR3_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[3].size ? vgpu->bars[3].size : 0;
			if (info.size) {
				info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			} else {
				info.flags = 0;
			}
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR4 - Additional BAR (if present) */
		if (info.index == VFIO_PCI_BAR4_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[4].size ? vgpu->bars[4].size : 0;
			if (info.size) {
				info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			} else {
				info.flags = 0;
			}
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* BAR5 - Additional BAR (if present) */
		if (info.index == VFIO_PCI_BAR5_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = vgpu->bars[5].size ? vgpu->bars[5].size : 0;
			if (info.size) {
				info.flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE | VFIO_REGION_INFO_FLAG_MMAP;
			} else {
				info.flags = 0;
			}
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* ROM region */
		if (info.index == VFIO_PCI_ROM_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = 0;
			info.flags = 0;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* VGA region */
		if (info.index == VFIO_PCI_VGA_REGION_INDEX) {
			info.offset = VFIO_PCI_INDEX_TO_OFFSET(info.index);
			info.size = 0;
			info.flags = 0;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		return -EINVAL;
	}

	if (cmd == VFIO_DEVICE_GET_IRQ_INFO) {
		struct vfio_irq_info info;

		minsz = offsetofend(struct vfio_irq_info, count);

		if (copy_from_user(&info, (void __user *)arg, minsz))
			return -EFAULT;

		if (info.argsz < minsz)
			return -EINVAL;

		/* INTx */
		if (info.index == VFIO_PCI_INTX_IRQ_INDEX) {
			info.flags = VFIO_IRQ_INFO_EVENTFD;
			info.count = 1;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* MSI */
		if (info.index == VFIO_PCI_MSI_IRQ_INDEX) {
			info.flags = VFIO_IRQ_INFO_EVENTFD;
			info.count = 1;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* MSI-X */
		if (info.index == VFIO_PCI_MSIX_IRQ_INDEX) {
			info.flags = VFIO_IRQ_INFO_EVENTFD | VFIO_IRQ_INFO_AUTOMASKED;
			info.count = 32; /* 32 MSI-X vectors */
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		/* ERR IRQ */
		if (info.index == VFIO_PCI_ERR_IRQ_INDEX) {
			info.flags = VFIO_IRQ_INFO_EVENTFD;
			info.count = 1;
			return copy_to_user((void __user *)arg, &info, minsz) ? -EFAULT : 0;
		}

		return -EINVAL;
	}

	if (cmd == VFIO_DEVICE_SET_IRQS) {
		struct vfio_irq_set *irq_set_ptr;
		struct vfio_irq_set irq_set_hdr;
		unsigned long minsz = offsetof(struct vfio_irq_set, data);
		int *fds = NULL;

		pr_info("nvidia-vgpu-vfio: VFIO_DEVICE_SET_IRQS\n");

		/* First copy the fixed part to get argsz */
		if (copy_from_user(&irq_set_hdr, (void __user *)arg, minsz)) {
			pr_err("nvidia-vgpu-vfio: SET_IRQS copy_from_user hdr failed\n");
			return -EFAULT;
		}

		pr_info("nvidia-vgpu-vfio: SET_IRQS hdr: argsz=%u flags=0x%x index=%u start=%u count=%u\n",
			irq_set_hdr.argsz, irq_set_hdr.flags, irq_set_hdr.index, irq_set_hdr.start, irq_set_hdr.count);

		if (irq_set_hdr.argsz < minsz) {
			pr_err("nvidia-vgpu-vfio: SET_IRQS argsz too small: %u < %lu\n", irq_set_hdr.argsz, minsz);
			return -EINVAL;
		}

		/* Allocate buffer for full struct including data array */
		irq_set_ptr = kmalloc(irq_set_hdr.argsz, GFP_KERNEL);
		if (!irq_set_ptr) {
			pr_err("nvidia-vgpu-vfio: SET_IRQS kmalloc failed\n");
			return -ENOMEM;
		}

		if (copy_from_user(irq_set_ptr, (void __user *)arg, irq_set_hdr.argsz)) {
			pr_err("nvidia-vgpu-vfio: SET_IRQS copy_from_user full failed\n");
			kfree(irq_set_ptr);
			return -EFAULT;
		}

		pr_info("nvidia-vgpu-vfio: SET_IRQS full: argsz=%u flags=0x%x index=%u start=%u count=%u\n",
			irq_set_ptr->argsz, irq_set_ptr->flags, irq_set_ptr->index, irq_set_ptr->start, irq_set_ptr->count);

		/* Validate data type */
		if ((irq_set_ptr->flags & VFIO_IRQ_SET_DATA_TYPE_MASK) == 0 ||
		    (irq_set_ptr->flags & VFIO_IRQ_SET_ACTION_TYPE_MASK) == 0) {
			pr_err("nvidia-vgpu-vfio: SET_IRQS invalid flags\n");
			kfree(irq_set_ptr);
			return -EINVAL;
		}

		/* Handle DATA_EVENTFD */
		if (irq_set_ptr->flags & VFIO_IRQ_SET_DATA_EVENTFD) {
			if (irq_set_ptr->count == 0) {
				kfree(irq_set_ptr);
				return -EINVAL;
			}

			fds = kcalloc(irq_set_ptr->count, sizeof(int), GFP_KERNEL);
			if (!fds) {
				kfree(irq_set_ptr);
				return -ENOMEM;
			}

			/* Copy eventfds directly from userspace (after the fixed header) */
			if (copy_from_user(fds, (void __user *)arg + minsz, irq_set_ptr->count * sizeof(int))) {
				pr_err("nvidia-vgpu-vfio: SET_IRQS copy_from_user data failed\n");
				kfree(fds);
				kfree(irq_set_ptr);
				return -EFAULT;
			}

			/* For now, just store the eventfds for INTx and MSI-X */
			if (irq_set_ptr->index == VFIO_PCI_INTX_IRQ_INDEX) {
				pr_info("nvidia-vgpu-vfio: SET_IRQS INTx trigger fds[%d]\n", irq_set_ptr->count);
				/* TODO: Store eventfds for INTx */
			} else if (irq_set_ptr->index == VFIO_PCI_MSIX_IRQ_INDEX) {
				pr_info("nvidia-vgpu-vfio: SET_IRQS MSI-X trigger fds[%d] start=%d\n", irq_set_ptr->count, irq_set_ptr->start);
				/* TODO: Store eventfds for MSI-X vectors */
			} else if (irq_set_ptr->index == VFIO_PCI_MSI_IRQ_INDEX) {
				pr_info("nvidia-vgpu-vfio: SET_IRQS MSI trigger fds[%d]\n", irq_set_ptr->count);
				/* TODO: Store eventfds for MSI */
			} else if (irq_set_ptr->index == VFIO_PCI_ERR_IRQ_INDEX) {
				pr_info("nvidia-vgpu-vfio: SET_IRQS ERR trigger fds[%d]\n", irq_set_ptr->count);
				/* TODO: Store eventfds for ERR */
			}

			kfree(fds);
			kfree(irq_set_ptr);
			return 0;
		}

		/* Handle DATA_NONE (mask/unmask) */
		if (irq_set_ptr->flags & VFIO_IRQ_SET_DATA_NONE) {
			pr_info("nvidia-vgpu-vfio: SET_IRQS DATA_NONE index=%d start=%d count=%d action=%x\n",
				irq_set_ptr->index, irq_set_ptr->start, irq_set_ptr->count, irq_set_ptr->flags & VFIO_IRQ_SET_ACTION_TYPE_MASK);
			kfree(irq_set_ptr);
			return 0;
		}

		/* Handle DATA_BOOL */
		if (irq_set_ptr->flags & VFIO_IRQ_SET_DATA_BOOL) {
			pr_info("nvidia-vgpu-vfio: SET_IRQS DATA_BOOL index=%d start=%d count=%d\n",
				irq_set_ptr->index, irq_set_ptr->start, irq_set_ptr->count);
			kfree(irq_set_ptr);
			return 0;
		}

		kfree(irq_set_ptr);
		return -EINVAL;
	}

	return -ENOTTY;
}

static int nv_vgpu_vfio_dev_mmap(struct vfio_device *vdev, struct vm_area_struct *vma)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	unsigned long offset = vma->vm_pgoff << PAGE_SHIFT;
	unsigned long size = vma->vm_end - vma->vm_start;
	unsigned long pfn;
	void *vaddr;
	struct nv_vgpu_bar_emul_ctx *bar_ctx = (struct nv_vgpu_bar_emul_ctx *)vgpu->bar_emul_ctx;
	unsigned long region_base;
	unsigned long region_size;

	pr_info("nvidia-vgpu-vfio: VFIO device mmap offset=0x%lx size=0x%lx\n", offset, size);

	/* Config space region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_CONFIG_REGION_INDEX);
	region_size = 4096;
	if (offset >= region_base && offset < region_base + region_size) {
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = vgpu->pci_config + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR0 region - map to BAR0 emulation shadow */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR0_REGION_INDEX);
	region_size = bar_ctx ? bar_ctx->bar0_shadow_size : (256 * 1024 * 1024);
	if (offset >= region_base && offset < region_base + region_size) {
		if (!bar_ctx || !bar_ctx->bar0_shadow)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = bar_ctx->bar0_shadow + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR1 region - VRAM (use physical address directly) */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR1_REGION_INDEX);
	region_size = vgpu->bars[1].size;
	if (offset >= region_base && offset < region_base + region_size) {
		if (region_size == 0)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		/* Use physical address directly since vaddr is NULL (managed by nvidia.ko) */
		pfn = (vgpu->bars[1].gpa + (offset - region_base)) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR2 region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR2_REGION_INDEX);
	region_size = vgpu->bars[2].size;
	if (offset >= region_base && offset < region_base + region_size) {
		if (!vgpu->bars[2].vaddr || region_size == 0)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = vgpu->bars[2].vaddr + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR3 region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR3_REGION_INDEX);
	region_size = vgpu->bars[3].size;
	if (offset >= region_base && offset < region_base + region_size) {
		if (!vgpu->bars[3].vaddr || region_size == 0)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = vgpu->bars[3].vaddr + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR4 region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR4_REGION_INDEX);
	region_size = vgpu->bars[4].size;
	if (offset >= region_base && offset < region_base + region_size) {
		if (!vgpu->bars[4].vaddr || region_size == 0)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = vgpu->bars[4].vaddr + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* BAR5 region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_BAR5_REGION_INDEX);
	region_size = vgpu->bars[5].size;
	if (offset >= region_base && offset < region_base + region_size) {
		if (!vgpu->bars[5].vaddr || region_size == 0)
			return -ENODEV;
		if (offset + size > region_base + region_size)
			return -EINVAL;
		vaddr = vgpu->bars[5].vaddr + (offset - region_base);
		pfn = virt_to_phys(vaddr) >> PAGE_SHIFT;
		vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
		return remap_pfn_range(vma, vma->vm_start, pfn, size, vma->vm_page_prot);
	}

	/* ROM region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_ROM_REGION_INDEX);
	region_size = 0;
	if (offset >= region_base && offset < region_base + region_size) {
		return -EINVAL;
	}

	/* VGA region */
	region_base = VFIO_PCI_INDEX_TO_OFFSET(VFIO_PCI_VGA_REGION_INDEX);
	region_size = 0;
	if (offset >= region_base && offset < region_base + region_size) {
		return -EINVAL;
	}

	pr_err("nvidia-vgpu-vfio: mmap failed for unknown offset 0x%lx\n", offset);
	return -EINVAL;
}

static void nv_vgpu_vfio_dev_dma_unmap(struct vfio_device *vdev, u64 iova, u64 length)
{
	pr_info("nvidia-vgpu-vfio: VFIO device dma_unmap iova=0x%llx len=0x%llx\n", iova, length);
}

static int nv_vgpu_vfio_dev_match(struct vfio_device *vdev, char *buf)
{
	struct nv_vgpu_device *vgpu = container_of(vdev, struct nv_vgpu_device, vfio_dev);
	
	/* Match by mdev UUID */
	if (vgpu->mdev && !strcmp(dev_name(&vgpu->mdev->dev), buf)) {
		return 1;
	}
	
	/* Also match by vfio device name (vfio0, vfio1, etc.) */
	if (!strcmp(dev_name(&vdev->device), buf)) {
		return 1;
	}
	
	return 0;
}

static int __init nv_vgpu_init(void)
{
	int ret;
	struct pci_dev *pdev = NULL;

	pr_info("nvidia-vgpu-vfio: init started\n");

	nv_vgpu_driver.class = class_create("nvidia-vgpu");
	if (IS_ERR(nv_vgpu_driver.class)) {
		ret = PTR_ERR(nv_vgpu_driver.class);
		pr_err("Failed to create nvidia-vgpu class: %d\n", ret);
		return ret;
	}
	pr_info("nvidia-vgpu-vfio: class created\n");

	/* Create static vGPU types as fallback (userspace can override via ioctl) */
	nv_vgpu_create_static_types();
	pr_info("nvidia-vgpu-vfio: static types created\n");

	/* Create mdev types from registered vGPU types */
	ret = nv_vgpu_create_mdev_types();
	if (ret) {
		pr_err("Failed to create mdev types: %d\n", ret);
		goto err_class;
	}
	pr_info("nvidia-vgpu-vfio: mdev types created\n");

	/* Find the parent PCI device to register mdev parent */
	pdev = pci_get_device(0x10de, 0x1ed3, NULL);
	if (!pdev) {
		pr_err("Failed to find NVIDIA GPU (10de:1ed3) for mdev parent registration\n");
		ret = -ENODEV;
		goto err_types;
	}
	pr_info("nvidia-vgpu-vfio: found GPU %s\n", pci_name(pdev));

	/* Register mdev parent via nvidia.ko */
	ret = nv_vgpu_mdev_register_parent(pdev);
	if (ret) {
		pr_err("Failed to register mdev parent: %d\n", ret);
		goto err_put_pdev;
	}
	pr_info("nvidia-vgpu-vfio: mdev parent registered\n");

	ret = mdev_register_driver(&nv_vgpu_mdev_driver);
	if (ret) {
		pr_err("Failed to register mdev driver: %d\n", ret);
		goto err_types;
	}
	pr_info("nvidia-vgpu-vfio: mdev driver registered\n");

	pci_dev_put(pdev);
	pdev = NULL;

	/* Provide VFIO ops to host driver (nvidia.ko) */
	ret = nvidia_vgpu_vfio_set_ops(&nv_vgpu_vfio_ops);
	if (ret) {
		pr_err("Failed to set VFIO ops: %d\n", ret);
		goto err_driver;
	}
	pr_info("nvidia-vgpu-vfio: VFIO ops set\n");

	pr_info("NVIDIA vGPU VFIO driver loaded (version %s)\n", NVIDIA_VGPU_VFIO_VERSION);
	return 0;

err_driver:
	mdev_unregister_driver(&nv_vgpu_mdev_driver);
err_types:
	nv_vgpu_free_type_info();
err_put_pdev:
	if (pdev)
		pci_dev_put(pdev);
err_class:
	class_destroy(nv_vgpu_driver.class);
	return ret;
}

static void __exit nv_vgpu_exit(void)
{
	struct pci_dev *pdev = NULL;

	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: starting cleanup\n");

	/* Break circular dependency with nvidia.ko FIRST
	 * by clearing the VFIO ops pointer. This signals nvidia.ko
	 * to release its references to our module. */
	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: clearing VFIO ops to break circular dependency\n");
	g_vgpu_vfio_ops = NULL;

	if (nv_vgpu_driver.device_count > 0) {
		struct nv_vgpu_device *vgpu;
		mutex_lock(&nv_vgpu_driver.lock);
		list_for_each_entry(vgpu, &nv_vgpu_driver.devices, list) {
			if (vgpu->pdev) {
				pdev = vgpu->pdev;
				break;
			}
		}
		mutex_unlock(&nv_vgpu_driver.lock);
	}

	if (!pdev) {
		pdev = pci_get_device(0x10de, 0x1ed3, NULL);
	}

	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: unregistering mdev driver\n");
	/* Unregister mdev driver FIRST, then mdev parent */
	mdev_unregister_driver(&nv_vgpu_mdev_driver);
	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: mdev driver unregistered\n");

	/* Unregister mdev parent via nvidia.ko */
	if (pdev) {
		pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: unregistering mdev parent\n");
		nv_vgpu_mdev_unregister_parent(pdev);
		pr_info("nvidia-vgpu-vfio: mdev parent unregistered for %s\n", dev_name(&pdev->dev));
		pci_dev_put(pdev);
	}

	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: destroying class\n");
	class_destroy(nv_vgpu_driver.class);

	pr_info("nvidia-vgpu-vfio: nv_vgpu_exit: freeing type info\n");
	nv_vgpu_free_type_info();

	pr_info("NVIDIA vGPU VFIO driver unloaded\n");
}

static void nv_vgpu_device_init_work(struct work_struct *work)
{
    struct nv_vgpu_device *vgpu = container_of(work, struct nv_vgpu_device, init_work);
    int ret;

    mutex_lock(&vgpu->lock);
    vgpu->state = NV_VGPU_STATE_INITIALIZING;
    mutex_unlock(&vgpu->lock);

    ret = nv_vgpu_bar_init(vgpu);
    if (ret) {
        nv_vgpu_err(vgpu, "BAR init failed: %d", ret);
        goto err_state;
    }

    ret = nv_vgpu_irq_init(vgpu);
    if (ret) {
        nv_vgpu_err(vgpu, "IRQ init failed: %d", ret);
        goto err_bar;
    }

    ret = nv_vgpu_memory_init(vgpu);
    if (ret) {
        nv_vgpu_err(vgpu, "Memory init failed: %d", ret);
        goto err_irq;
    }

    /* Physical GSP is managed by host nvidia.ko - skip nv_vgpu_gsp_init to avoid conflict */
    nv_vgpu_info(vgpu, "Skipping physical GSP init (host nvidia.ko manages physical GSP)");

    /* Initialize GSP emulation engine (host-side) for guest GSP RPC handling */
    ret = gsp_emu_init(vgpu);
    if (ret) {
        nv_vgpu_err(vgpu, "GSP emulation init failed: %d", ret);
        goto err_memory;
    }

    ret = nv_vgpu_migration_init(vgpu);
    if (ret) {
        nv_vgpu_warn(vgpu, "Migration init failed: %d", ret);
    }

    mutex_lock(&vgpu->lock);
    vgpu->state = NV_VGPU_STATE_READY;
    mutex_unlock(&vgpu->lock);

    nv_vgpu_info(vgpu, "vGPU device initialized successfully");
    return;

err_memory:
    nv_vgpu_memory_cleanup(vgpu);
err_irq:
    nv_vgpu_irq_cleanup(vgpu);
err_bar:
    nv_vgpu_bar_cleanup(vgpu);
err_state:
    mutex_lock(&vgpu->lock);
    vgpu->state = NV_VGPU_STATE_ERROR;
    mutex_unlock(&vgpu->lock);
}

static void nv_vgpu_device_cleanup_work(struct work_struct *work)
{
    struct nv_vgpu_device *vgpu = container_of(work, struct nv_vgpu_device, cleanup_work);

    nv_vgpu_migration_cleanup(vgpu);
    gsp_emu_cleanup(vgpu);
    nv_vgpu_memory_cleanup(vgpu);
    nv_vgpu_irq_cleanup(vgpu);
    nv_vgpu_bar_cleanup(vgpu);

    mutex_lock(&vgpu->lock);
    vgpu->state = NV_VGPU_STATE_UNINITIALIZED;
    mutex_unlock(&vgpu->lock);

    nv_vgpu_info(vgpu, "vGPU device cleaned up");
}

static NV_STATUS nv_vgpu_vfio_probe(struct pci_dev *pdev, NvU32 vgpu_type_id, NvU32 *vgpu_id)
{
	pr_info("nvidia-vgpu-vfio: VFIO probe called for vGPU type %u\n", vgpu_type_id);
	return NV_OK;
}

static void nv_vgpu_vfio_remove(struct pci_dev *pdev)
{
	pr_info("nvidia-vgpu-vfio: VFIO remove called\n");
	/* Clear VFIO ops to break circular dependency with nvidia.ko */
	g_vgpu_vfio_ops = NULL;
}

static NV_STATUS nv_vgpu_vfio_inject_interrupt(void *vgpu_dev)
{
	pr_info("nvidia-vgpu-vfio: VFIO inject interrupt called\n");
	return NV_OK;
}

/* ============================================================
 * PCI Config Space for VFIO mdev
 * ============================================================ */

static int nv_vgpu_config_space_init(struct nv_vgpu_device *vgpu, struct pci_dev *pdev)
{
	u8 *cfg = vgpu->pci_config;
	u16 vendor_id = 0x10de;   /* NVIDIA */
	u16 device_id = 0x1ed3;   /* TU104 vGPU (default) */
	u16 subsys_vendor = 0x10de;
	u16 subsys_device = 0x12b4;
	u8 revision = 0xa1;
	u8 prog_if = 0x00;
	u8 subclass = 0x00;
	u8 class_code = 0x03;     /* Display controller */
	u8 cache_line = 0x10;
	u8 latency = 0x00;
	u8 header_type = 0x00;    /* Normal device */
	u8 bist = 0x00;
	int i;

	memset(cfg, 0, 4096);

	/* Look up vGPU type info from mdev type name */
	if (vgpu->mdev && vgpu->mdev->type && vgpu->mdev->type->sysfs_name) {
		for (i = 0; i < g_num_vgpu_types; i++) {
			if (strcmp(vgpu->mdev->type->sysfs_name, g_vgpu_type_infos[i].name) == 0) {
				vendor_id = (u16)((g_vgpu_type_infos[i].vdev_id >> 16) & 0xFFFF);
				device_id = (u16)(g_vgpu_type_infos[i].vdev_id & 0xFFFF);
				subsys_vendor = (u16)((g_vgpu_type_infos[i].pdev_id >> 16) & 0xFFFF);
				subsys_device = (u16)(g_vgpu_type_infos[i].pdev_id & 0xFFFF);
				nv_vgpu_info(vgpu, "Using vGPU type: name='%s', vendor=0x%04x, device=0x%04x, subsys_vendor=0x%04x, subsys_device=0x%04x",
					g_vgpu_type_infos[i].name, vendor_id, device_id, subsys_vendor, subsys_device);
				break;
			}
		}
	}

	/* PCI Header (offset 0x00-0x3F) */
	cfg[0x00] = vendor_id & 0xff;
	cfg[0x01] = (vendor_id >> 8) & 0xff;
	cfg[0x02] = device_id & 0xff;
	cfg[0x03] = (device_id >> 8) & 0xff;

	/* Command register (0x04) - Memory + IO + BusMaster */
	cfg[0x04] = 0x07;
	cfg[0x05] = 0x00;

	/* Status register (0x06) */
	cfg[0x06] = 0x10;
	cfg[0x07] = 0x00;

	/* Revision ID (0x08) */
	cfg[0x08] = revision;

	/* Class code (0x09-0x0B) */
	cfg[0x09] = prog_if;
	cfg[0x0a] = subclass;
	cfg[0x0b] = class_code;

	/* Cache line size (0x0C) */
	cfg[0x0c] = cache_line;

	/* Latency timer (0x0D) */
	cfg[0x0d] = latency;

	/* Header type (0x0E) */
	cfg[0x0e] = header_type;

	/* BIST (0x0F) */
	cfg[0x0f] = bist;

	/* BAR0 - 16MB Memory space (32-bit, non-prefetchable) - matches TU104 */
	cfg[0x10] = 0x00;  /* 32-bit, non-prefetchable */
	cfg[0x11] = 0x00;
	cfg[0x12] = 0x00;
	cfg[0x13] = 0x00;

	/* BAR1 - 256MB Memory space (64-bit, prefetchable) - VRAM */
	cfg[0x14] = 0x0c;  /* 64-bit, prefetchable */
	cfg[0x15] = 0x00;
	cfg[0x16] = 0x00;
	cfg[0x17] = 0x00;
	cfg[0x18] = 0x00;  /* Upper 32 bits */
	cfg[0x19] = 0x00;
	cfg[0x1a] = 0x00;
	cfg[0x1b] = 0x00;

	/* BAR2 - Not present (consumed by BAR1 upper half) */
	cfg[0x1c] = 0x00;
	cfg[0x1d] = 0x00;
	cfg[0x1e] = 0x00;
	cfg[0x1f] = 0x00;

	/* BAR3 - 32MB Memory space (64-bit, prefetchable) */
	cfg[0x20] = 0x0c;  /* 64-bit, prefetchable */
	cfg[0x21] = 0x00;
	cfg[0x22] = 0x00;
	cfg[0x23] = 0x00;
	cfg[0x24] = 0x00;  /* Upper 32 bits */
	cfg[0x25] = 0x00;
	cfg[0x26] = 0x00;
	cfg[0x27] = 0x00;

	/* BAR4 - Not present (consumed by BAR3 upper half) */
	cfg[0x28] = 0x00;
	cfg[0x29] = 0x00;
	cfg[0x2a] = 0x00;
	cfg[0x2b] = 0x00;

	/* CardBus CIS Pointer (0x2C) - Not used */
	cfg[0x2c] = 0x00;
	cfg[0x2d] = 0x00;
	cfg[0x2e] = 0x00;
	cfg[0x2f] = 0x00;

	/* Subsystem Vendor ID (0x2E) and Subsystem ID (0x2F) - wait, correct offsets: */
	/* Subsystem Vendor ID at 0x2C, Subsystem ID at 0x2E */
	cfg[0x2c] = subsys_vendor & 0xff;
	cfg[0x2d] = (subsys_vendor >> 8) & 0xff;
	cfg[0x2e] = subsys_device & 0xff;
	cfg[0x2f] = (subsys_device >> 8) & 0xff;

	/* Expansion ROM Base Address (0x30-0x33) - Not used */
	cfg[0x30] = 0x00;
	cfg[0x31] = 0x00;
	cfg[0x32] = 0x00;
	cfg[0x33] = 0x00;

	/* Capabilities Pointer (0x34) - Point to Power Management at 0x40 */
	cfg[0x34] = 0x40;
	cfg[0x35] = 0x00;
	cfg[0x36] = 0x00;
	cfg[0x37] = 0x00;

	/* Interrupt Line (0x3C) - Will be set by guest */
	cfg[0x3c] = 0x00;
	/* Interrupt Pin (0x3D) - INTA */
	cfg[0x3d] = 0x01;
	/* Min_Gnt (0x3E) */
	cfg[0x3e] = 0x00;
	/* Max_Lat (0x3F) */
	cfg[0x3f] = 0x00;

	/* Subsystem Vendor ID (0x3C) */
	cfg[0x3c] = subsys_vendor & 0xff;
	cfg[0x3d] = (subsys_vendor >> 8) & 0xff;

	/* Subsystem ID (0x3E) */
	cfg[0x3e] = subsys_device & 0xff;
	cfg[0x3f] = (subsys_device >> 8) & 0xff;

	/* Expansion ROM Base Address (0x40-0x43) - Not used */

	/* Capabilities Pointer (0x34) - Point to Power Management at 0x40 */
	cfg[0x34] = 0x40;
	cfg[0x35] = 0x00;
	cfg[0x36] = 0x00;
	cfg[0x37] = 0x00;

	/* Interrupt Line (0x3C) - Will be set by guest */
	cfg[0x3c] = 0x00;
	/* Interrupt Pin (0x3D) - INTA */
	cfg[0x3d] = 0x01;
	/* Min_Gnt (0x3E) */
	cfg[0x3e] = 0x00;
	/* Max_Lat (0x3F) */
	cfg[0x3f] = 0x00;

	/* Power Management Capability (0x40) */
	cfg[0x40] = 0x01;  /* Capability ID: Power Management */
	cfg[0x41] = 0x50;  /* Next capability: MSI at 0x50 */
	cfg[0x42] = 0x03;  /* PME support from D0, D3hot, D3cold */
	cfg[0x43] = 0x00;
	cfg[0x44] = 0x00;  /* PM Control/Status */
	cfg[0x45] = 0x00;
	cfg[0x46] = 0x00;
	cfg[0x47] = 0x00;

	/* MSI Capability (0x50) */
	cfg[0x50] = 0x05;  /* Capability ID: MSI */
	cfg[0x51] = 0x70;  /* Next capability: MSI-X at 0x70 */
	cfg[0x52] = 0x80;  /* 64-bit capable, 1 vector */
	cfg[0x53] = 0x00;
	cfg[0x54] = 0x00;  /* Message Address (low) - set by guest */
	cfg[0x55] = 0x00;
	cfg[0x56] = 0x00;
	cfg[0x57] = 0x00;
	cfg[0x58] = 0x00;  /* Message Address (high) - 64-bit */
	cfg[0x59] = 0x00;
	cfg[0x5a] = 0x00;
	cfg[0x5b] = 0x00;
	cfg[0x5c] = 0x01;  /* Message Data - set by guest */
	cfg[0x5d] = 0x00;

	/* MSI-X Capability (0x70) */
	cfg[0x70] = 0x11;  /* Capability ID: MSI-X */
	cfg[0x71] = 0x80;  /* Next capability: PCI Express at 0x80 */
	cfg[0x72] = 0x1f;  /* Table size = 31 (32 vectors), 64-bit */
	cfg[0x73] = 0x00;
	/* Table Offset/BIR (0x74) - Table at BAR0 offset 0x2000 */
	cfg[0x74] = 0x00;
	cfg[0x75] = 0x20;
	cfg[0x76] = 0x00;
	cfg[0x77] = 0x00;
	/* PBA Offset/BIR (0x78) - PBA at BAR0 offset 0x3000 */
	cfg[0x78] = 0x00;
	cfg[0x79] = 0x30;
	cfg[0x7a] = 0x00;
	cfg[0x7b] = 0x00;

	/* PCI Express Capability (0x80) */
	cfg[0x80] = 0x10;  /* Capability ID: PCI Express */
	cfg[0x81] = 0x00;  /* Next capability: None */
	cfg[0x82] = 0x02;  /* Device type: PCIe Endpoint, version 2 */
	cfg[0x83] = 0x00;
	cfg[0x84] = 0x00;  /* Device Capabilities */
	cfg[0x85] = 0x00;
	cfg[0x86] = 0x00;
	cfg[0x87] = 0x00;
	cfg[0x88] = 0x00;  /* Device Control */
	cfg[0x89] = 0x00;
	cfg[0x8a] = 0x10;  /* Device Status */
	cfg[0x8b] = 0x00;
	cfg[0x8c] = 0x00;  /* Link Capabilities */
	cfg[0x8d] = 0x00;
	cfg[0x8e] = 0x00;
	cfg[0x8f] = 0x00;
	cfg[0x90] = 0x00;  /* Link Control */
	cfg[0x91] = 0x00;
	cfg[0x92] = 0x00;  /* Link Status */
	cfg[0x93] = 0x00;
	cfg[0x94] = 0x00;  /* Slot Capabilities */
	cfg[0x95] = 0x00;
	cfg[0x96] = 0x00;
	cfg[0x97] = 0x00;
	cfg[0x98] = 0x00;  /* Slot Control */
	cfg[0x99] = 0x00;
	cfg[0x9a] = 0x00;  /* Slot Status */
	cfg[0x9b] = 0x00;
	cfg[0x9c] = 0x00;  /* Root Capabilities */
	cfg[0x9d] = 0x00;
	cfg[0x9e] = 0x00;  /* Root Control */
	cfg[0x9f] = 0x00;

	vgpu->config_space_dirty = false;
	nv_vgpu_info(vgpu, "PCI config space initialized (vendor=0x%04x, device=0x%04x, subsys_vendor=0x%04x, subsys_device=0x%04x)", vendor_id, device_id, subsys_vendor, subsys_device);
	return 0;
}

static ssize_t nv_vgpu_config_read(struct nv_vgpu_device *vgpu, char __user *buf, size_t count, loff_t *ppos)
{
	loff_t pos = *ppos;

	if (pos >= 256)
		return 0;

	if (pos + count > 256)
		count = 256 - pos;

	if (copy_to_user(buf, vgpu->pci_config + pos, count))
		return -EFAULT;

	*ppos += count;
	return count;
}

static ssize_t nv_vgpu_config_write(struct nv_vgpu_device *vgpu, const char __user *buf, size_t count, loff_t *ppos)
{
	loff_t pos = *ppos;
	u8 val;

	if (pos >= 256)
		return -EINVAL;

	if (pos + count > 256)
		count = 256 - pos;

	/* Handle specific writable registers */
	while (count--) {
		if (copy_from_user(&val, buf++, 1))
			return -EFAULT;

		switch (pos) {
		case 0x04: /* Command register */
			/* Allow Memory, IO, BusMaster */
			vgpu->pci_config[pos] = val & 0x07;
			break;
		case 0x0c: /* Cache line size */
			vgpu->pci_config[pos] = val;
			break;
		case 0x0d: /* Latency timer */
			vgpu->pci_config[pos] = val;
			break;
		case 0x3c: /* Interrupt line */
			vgpu->pci_config[pos] = val;
			break;
		case 0x54: case 0x55: case 0x56: case 0x57: /* MSI Message Address Low */
		case 0x58: case 0x59: case 0x5a: case 0x5b: /* MSI Message Address High */
		case 0x5c: case 0x5d: /* MSI Message Data */
			vgpu->pci_config[pos] = val;
			break;
		case 0x72: case 0x73: /* MSI-X Message Control */
			vgpu->pci_config[pos] = val;
			break;
		case 0x74: case 0x75: case 0x76: case 0x77: /* MSI-X Table Offset/BIR */
			vgpu->pci_config[pos] = val;
			break;
		case 0x78: case 0x79: case 0x7a: case 0x7b: /* MSI-X PBA Offset/BIR */
			vgpu->pci_config[pos] = val;
			break;
		case 0x88: case 0x89: /* Device Control */
		case 0x90: case 0x91: /* Link Control */
		case 0x98: case 0x99: /* Slot Control */
		case 0x9e: case 0x9f: /* Root Control */
			vgpu->pci_config[pos] = val;
			break;
		default:
			/* Most config space is read-only for guest */
			break;
		}
		pos++;
	}

	vgpu->config_space_dirty = true;
	*ppos = pos;
	return pos - *ppos + count;
}

/* ============================================================
 * Exported symbols for nvidia.ko (vGPU VFIO interface)
 * ============================================================ */

NV_STATUS nvidia_vgpu_vfio_probe(struct pci_dev *pdev)
{
	pr_info("nvidia-vgpu-vfio: nvidia_vgpu_vfio_probe called for %s\n", pci_name(pdev));
	return NV_OK;
}
EXPORT_SYMBOL(nvidia_vgpu_vfio_probe);

void nvidia_vgpu_vfio_remove(struct pci_dev *pdev, NvBool remove_device)
{
	pr_info("nvidia-vgpu-vfio: nvidia_vgpu_vfio_remove called for %s, remove_device=%d\n",
		pci_name(pdev), remove_device);
}
EXPORT_SYMBOL(nvidia_vgpu_vfio_remove);

module_init(nv_vgpu_init);
module_exit(nv_vgpu_exit);