#ifndef _NVIDIA_VGPU_VFIO_H_
#define _NVIDIA_VGPU_VFIO_H_

#include <linux/types.h>
#include <linux/mutex.h>
#include <linux/list.h>
#include <linux/workqueue.h>
#include <linux/vfio.h>
#include <linux/mdev.h>
#include <linux/pci.h>
#include <linux/dma-mapping.h>
#include <linux/atomic.h>

#define VGSP_VERSION "0.1.0"
#define VGSP_NAME "vgsp"

#define NVIDIA_VGPU_VFIO_VERSION VGSP_VERSION
#define NVIDIA_VGPU_VFIO_NAME VGSP_NAME

#define NVIDIA_VGPU_MAX_VF_COUNT 32
#define NVIDIA_VGPU_MAX_REGIONS 8
#define NVIDIA_VGPU_MAX_IRQS 16

#define NV_VGPU_MAX_REGIONS NVIDIA_VGPU_MAX_REGIONS
#define NV_VGPU_MAX_IRQS NVIDIA_VGPU_MAX_IRQS

#define NV_VGPU_BAR0_SIZE (16 * 1024 * 1024)
/*
 * BAR1 (VRAM aperture) and the emulated framebuffer size. These two move
 * together: the guest derives its BAR1 size from fbLength
 * (kbusInitBarsSize_VF: pciBarSizes[1] = pGpu->fbLength), and its vidmem
 * address space is [0, fbLength). The shadow backing bars[1].vaddr is
 * NV_VGPU_BAR1_SIZE, so any guest vidmem PA the emulation translates must
 * stay inside it. 512 MiB gives UVM enough headroom after its own
 * reservations to satisfy the acceptance program's 3x4 MiB vadd buffers
 * (at 256 MiB the guest reported only ~5 MiB free).
 */
#define NV_VGPU_BAR1_SIZE (512 * 1024 * 1024)
#define NV_VGPU_EMUL_FB_SIZE (512ULL * 1024 * 1024)
#define NV_VGPU_BAR2_SIZE (32 * 1024 * 1024)
#define NV_VGPU_BAR3_SIZE (32 * 1024 * 1024)

enum nv_vgpu_type {
	NV_VGPU_TYPE_UNKNOWN = 0,
	NV_VGPU_TYPE_AMPERE,
	NV_VGPU_TYPE_HOPPER,
	NV_VGPU_TYPE_TURING,
	NV_VGPU_TYPE_PASCAL,
	NV_VGPU_TYPE_MAXWELL,
	NV_VGPU_TYPE_KEPLER,
};

enum nv_vgpu_state {
	NV_VGPU_STATE_UNINITIALIZED = 0,
	NV_VGPU_STATE_INITIALIZING,
	NV_VGPU_STATE_READY,
	NV_VGPU_STATE_RUNNING,
	NV_VGPU_STATE_SUSPENDED,
	NV_VGPU_STATE_ERROR,
	NV_VGPU_STATE_MIGRATING,
};

struct nv_vgpu_bar {
	u64 gpa;
	u64 size;
	u64 iova;
	void *vaddr;
	bool mapped;
	bool trapped;
	/*
	 * True when vaddr is a host-RAM shadow (vzalloc) owned by this module
	 * instead of an ioremap'd physical BAR. Guest mappings must then remap
	 * the shadow's own pages; a real BAR must never be exposed to the
	 * guest outside the supervised physical claim.
	 */
	bool emulated;
	struct mutex lock;
};

struct nv_vgpu_irq {
	int virq;
	int type;
	bool masked;
	struct eventfd_ctx *trigger;
	struct irq_domain *domain;
};

struct nv_vgpu_region {
	u64 offset;
	u64 size;
	u64 iova;
	int fd;
	void *vaddr;
	bool readonly;
};

struct nv_vgpu_migration_ctx;

struct nv_vgpu_migration_state {
	u32 version;
	u32 flags;
	u64 timestamp;
	u8 *data;
	size_t data_size;
	struct completion done;
	int result;
	struct nv_vgpu_migration_ctx *ctx;
};

struct nv_vgpu_device {
	/* VFIO device for VFIO group - must be first for vfio_alloc_device */
	struct vfio_device vfio_dev;
	
	struct mdev_device *mdev;
	struct pci_dev *pdev;
	struct device dev;
	struct mutex lock;
	enum nv_vgpu_type type;
	enum nv_vgpu_state state;
	u32 vgpu_id;
	u32 gpu_instance_id;
	u32 compute_instance_id;
	
	struct nv_vgpu_bar bars[NV_VGPU_MAX_REGIONS];
	struct nv_vgpu_irq irqs[NV_VGPU_MAX_IRQS];
	struct nv_vgpu_region regions[NV_VGPU_MAX_REGIONS];
	
	struct nv_vgpu_migration_state migration;
	
	struct work_struct init_work;
	struct work_struct cleanup_work;
	
	void *gsp_comm_ctx;
	void *bar_emul_ctx;
	void *mm_ctx;
	void *irq_ctx;
	
	u64 fb_base;
	u64 fb_size;
	/* Emulated framebuffer/vidmem size for this mdev, from its type's
	 * vram_mb (nv_vgpu_type_info); defaults to NV_VGPU_EMUL_FB_SIZE. All
	 * emulated-GSP consumers (GSCI fb_length/lb region/frtsOffset, WPR2,
	 * LOCAL_MEMORY_RANGE, vidmem allocator, BAR shadow) must agree. */
	u64 fb_emul_size;
	u64 bar0_config_base;
	
	struct list_head list;
	
	bool dry_run;

	/* PCI Config Space for VFIO mdev */
	u8 pci_config[4096];
	bool config_space_dirty;

 	/* VFIO device open/close tracking */
 	atomic_t vfio_dev_open_count;
 	wait_queue_head_t vfio_dev_close_wait;

 	/* VFIO device registration state */
 	bool vfio_dev_registered;

	/* True when this mdev holds a reference on the physical backend
	 * session (gsp_phys_backend_start() succeeded in init_work); the
	 * matching gsp_phys_backend_stop() runs in cleanup_work after the
	 * emulation thread is stopped. */
	bool phys_backend_ref;
 };

typedef struct nv_vgpu_device nv_vgpu_device;

struct nv_vgpu_driver_data {
	struct mutex lock;
	struct list_head devices;
	u32 device_count;
	struct class *class;
	struct bus_type *bus;
};

extern struct nv_vgpu_driver_data nv_vgpu_driver;

int nv_vgpu_device_create(struct mdev_device *mdev);
void nv_vgpu_device_destroy(struct mdev_device *mdev);
int nv_vgpu_device_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_device_cleanup(struct nv_vgpu_device *vgpu);

int nv_vgpu_bar_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_bar_cleanup(struct nv_vgpu_device *vgpu);
int nv_vgpu_bar_trap(struct nv_vgpu_device *vgpu, int bar_idx);
void nv_vgpu_bar_untrap(struct nv_vgpu_device *vgpu, int bar_idx);

int nv_vgpu_irq_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_irq_cleanup(struct nv_vgpu_device *vgpu);
int nv_vgpu_irq_enable(struct nv_vgpu_device *vgpu, int irq_idx);
int nv_vgpu_irq_disable(struct nv_vgpu_device *vgpu, int irq_idx);

int nv_vgpu_memory_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_memory_cleanup(struct nv_vgpu_device *vgpu);
int nv_vgpu_memory_map(struct nv_vgpu_device *vgpu, u64 gpa, u64 size, u64 *iova);
int nv_vgpu_memory_unmap(struct nv_vgpu_device *vgpu, u64 iova, u64 size);

int nv_vgpu_gsp_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_gsp_cleanup(struct nv_vgpu_device *vgpu);
int nv_vgpu_gsp_send_rpc(struct nv_vgpu_device *vgpu, void *msg, size_t size);
int nv_vgpu_gsp_recv_rpc(struct nv_vgpu_device *vgpu, void *msg, size_t size, int timeout_ms);
int nv_vgpu_gsp_load_firmware(struct nv_vgpu_device *vgpu);
void nv_vgpu_gsp_release_firmware(struct nv_vgpu_device *vgpu);
int nv_vgpu_gsp_boot(struct nv_vgpu_device *vgpu);

/* mdev driver/parent registration (main.c) */
int nv_vgpu_create_mdev_types(void);
void nv_vgpu_free_type_info(void);
int nv_vgpu_mdev_driver_register(void);
void nv_vgpu_mdev_driver_unregister(void);
int nv_vgpu_mdev_register_parent(struct pci_dev *pdev);
void nv_vgpu_mdev_unregister_parent(struct pci_dev *pdev);

int nv_vgpu_migration_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_migration_cleanup(struct nv_vgpu_device *vgpu);
int nv_vgpu_migration_save_state(struct nv_vgpu_device *vgpu, u8 **data, size_t *size);
int nv_vgpu_migration_restore_state(struct nv_vgpu_device *vgpu, u8 *data, size_t size);
int nv_vgpu_migration_prepare(struct nv_vgpu_device *vgpu);
int nv_vgpu_migration_complete(struct nv_vgpu_device *vgpu);
int nv_vgpu_migration_cancel(struct nv_vgpu_device *vgpu);
u8 *nv_vgpu_migration_get_data(struct nv_vgpu_device *vgpu, size_t *size);

static inline struct nv_vgpu_device *mdev_to_vgpu(struct mdev_device *mdev)
{
	return dev_get_drvdata(&mdev->dev);
}

/* Per-mdev emulated framebuffer size in bytes (default 512 MiB). */
static inline u64 nv_vgpu_fb_emul_size(const struct nv_vgpu_device *vgpu)
{
	return (vgpu && vgpu->fb_emul_size) ? vgpu->fb_emul_size :
					      NV_VGPU_EMUL_FB_SIZE;
}

/* NULL-tolerant logging: emulation paths may run with a NULL vgpu
 * (e.g. unit tests, early init) - fall back to plain pr_*. */
#define nv_vgpu_dbg(vgpu, fmt, ...) do { \
	if (vgpu) \
		dev_dbg(&(vgpu)->dev, KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
	else \
		pr_debug("%s: " fmt, __func__, ##__VA_ARGS__); \
} while (0)

#define nv_vgpu_info(vgpu, fmt, ...) do { \
	if (vgpu) \
		dev_info(&(vgpu)->dev, KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
	else \
		pr_info(KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
} while (0)

#define nv_vgpu_warn(vgpu, fmt, ...) do { \
	if (vgpu) \
		dev_warn(&(vgpu)->dev, KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
	else \
		pr_warn(KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
} while (0)

#define nv_vgpu_err(vgpu, fmt, ...) do { \
	if (vgpu) \
		dev_err(&(vgpu)->dev, KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
	else \
		pr_err(KBUILD_MODNAME ": %s: " fmt, __func__, ##__VA_ARGS__); \
} while (0)

#endif