#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/pci.h>
#include <linux/vfio.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 "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);

/* 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);

static struct mdev_type *nv_vgpu_mdev_types[4];

static 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,
	.driver = {
		.name = "nvidia-vgpu-vfio-mdev",
	},
};

/* ============================================================
 * 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;
static struct pci_dev *g_pdev = 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;
}

/* 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;
}

/* 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;
}

extern NV_STATUS nv_vgpu_mdev_register_parent(struct pci_dev *pdev);
extern void nv_vgpu_mdev_unregister_parent(struct pci_dev *pdev);
extern int nv_vgpu_mdev_driver_register(void);
extern void nv_vgpu_mdev_driver_unregister(void);

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;
	int ret;

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

	dev_set_drvdata(&mdev->dev, 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);

	parent = mdev->type->parent;
	if (!parent || !parent->dev) {
		ret = -ENODEV;
		goto err_free;
	}
	pdev = to_pci_dev(parent->dev);
	if (!pdev) {
		ret = -ENODEV;
		goto err_free;
	}
	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);

	schedule_work(&vgpu->init_work);

	return 0;

err_put:
	pci_dev_put(vgpu->pdev);
err_free:
	kfree(vgpu);
	return ret;
}

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

	cancel_work_sync(&vgpu->init_work);
	schedule_work(&vgpu->cleanup_work);
	flush_work(&vgpu->cleanup_work);

	mutex_lock(&nv_vgpu_driver.lock);
	list_del(&vgpu->list);
	nv_vgpu_driver.device_count--;
	mutex_unlock(&nv_vgpu_driver.lock);

	pci_dev_put(vgpu->pdev);
	kfree(vgpu);
}

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

	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;
	}

	for (i = 0; i < 4; 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[0]->sysfs_name = "nvidia-vgpu-ampere";
	nv_vgpu_mdev_types[0]->pretty_name = "NVIDIA Ampere vGPU";
	nv_vgpu_mdev_types[1]->sysfs_name = "nvidia-vgpu-hopper";
	nv_vgpu_mdev_types[1]->pretty_name = "NVIDIA Hopper vGPU";
	nv_vgpu_mdev_types[2]->sysfs_name = "nvidia-vgpu-turing";
	nv_vgpu_mdev_types[2]->pretty_name = "NVIDIA Turing vGPU";
	nv_vgpu_mdev_types[3]->sysfs_name = "nvidia-vgpu-pascal";
	nv_vgpu_mdev_types[3]->pretty_name = "NVIDIA Pascal vGPU";

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

	/* 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_driver;
	}

	/* Register mdev parent via nvidia.ko */
	ret = nv_vgpu_mdev_register_parent(pdev);
	pci_dev_put(pdev);
	if (ret) {
		pr_err("Failed to register mdev parent: %d\n", ret);
		goto err_driver;
	}

	/* 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 driver loaded (version %s)\n", NVIDIA_VGPU_VFIO_VERSION);
	return 0;

err_driver:
	mdev_unregister_driver(&nv_vgpu_mdev_driver);
err_types:
	for (i = 0; i < 4; i++) {
		kfree(nv_vgpu_mdev_types[i]);
		nv_vgpu_mdev_types[i] = NULL;
	}
	class_destroy(nv_vgpu_driver.class);
	return ret;
}

static void __exit nv_vgpu_exit(void)
{
	int i;
	struct pci_dev *pdev = 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);
	}

	/* Release host vGPU ops */
	if (pdev) {
		rm_vgpu_vfio_ops_t ops = {0};
		ops.version_string = NV_VERSION_STRING;
		nvidia_vgpu_vfio_get_ops(&ops);
	}

	/* Unregister mdev parent via nvidia.ko */
	if (pdev) {
		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);
	}

	mdev_unregister_driver(&nv_vgpu_mdev_driver);

	class_destroy(nv_vgpu_driver.class);

	for (i = 0; i < 4; i++) {
		kfree(nv_vgpu_mdev_types[i]);
		nv_vgpu_mdev_types[i] = NULL;
	}

	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;
    }

    ret = nv_vgpu_gsp_init(vgpu);
    if (ret) {
        nv_vgpu_err(vgpu, "GSP init failed: %d", ret);
        goto err_memory;
    }

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

    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_gsp:
    nv_vgpu_gsp_cleanup(vgpu);
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_gsp_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");
}

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 */
	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;

	memset(cfg, 0, 256);

	/* 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 (prefetchable) */
	cfg[0x10] = 0x04;  /* 64-bit, prefetchable */
	cfg[0x11] = 0x00;
	cfg[0x12] = 0x00;
	cfg[0x13] = 0x00;
	cfg[0x14] = 0x00;
	cfg[0x15] = 0x00;
	cfg[0x16] = 0x00;
	cfg[0x17] = 0x00;

	/* BAR1 - 256MB Memory space (prefetchable) */
	cfg[0x18] = 0x04;
	cfg[0x19] = 0x00;
	cfg[0x1a] = 0x00;
	cfg[0x1b] = 0x00;
	cfg[0x1c] = 0x00;
	cfg[0x1d] = 0x00;
	cfg[0x1e] = 0x00;
	cfg[0x1f] = 0x00;

	/* BAR2 - 32MB Memory space */
	cfg[0x20] = 0x04;
	cfg[0x21] = 0x00;
	cfg[0x22] = 0x00;
	cfg[0x23] = 0x00;
	cfg[0x24] = 0x00;
	cfg[0x25] = 0x00;
	cfg[0x26] = 0x00;
	cfg[0x27] = 0x00;

	/* BAR3 - 32MB Memory space */
	cfg[0x28] = 0x04;
	cfg[0x29] = 0x00;
	cfg[0x2a] = 0x00;
	cfg[0x2b] = 0x00;
	cfg[0x2c] = 0x00;
	cfg[0x2d] = 0x00;
	cfg[0x2e] = 0x00;
	cfg[0x2f] = 0x00;

	/* BAR4 - Not used */
	cfg[0x30] = 0x00;
	cfg[0x31] = 0x00;
	cfg[0x32] = 0x00;
	cfg[0x33] = 0x00;

	/* BAR5 - Not used */
	cfg[0x34] = 0x00;
	cfg[0x35] = 0x00;
	cfg[0x36] = 0x00;
	cfg[0x37] = 0x00;

	/* CardBus CIS Pointer (0x38-0x3B) - Not used */
	cfg[0x38] = 0x00;
	cfg[0x39] = 0x00;
	cfg[0x3a] = 0x00;
	cfg[0x3b] = 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] = 0x00;  /* Next capability: None */
	cfg[0x72] = 0x00;  /* Table size = 0 (1 vector), 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)", vendor_id, device_id);
	return 0;
}

static ssize_t nv_vgpu_config_read(struct nv_vgpu_device *vgpu, char __user *buf, size_t count, loff_t *ppos)
{
	size_t bytes_read = 0;
	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;
}

module_init(nv_vgpu_init);
module_exit(nv_vgpu_exit);