#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/io.h>
#include <linux/vmalloc.h>
#include <linux/pci.h>
#include <linux/mm.h>
#include <linux/highmem.h>
#include <linux/eventfd.h>

#include "nvidia-vgpu-vfio.h"
#include "bar_emulation.h"
#include "gsp_comm.h"
#include "gsp_emulation.h"

#define NV_VGPU_BAR0_TRAP_REGIONS 64

/* MSI-X Table entry size (16 bytes: addr_lo, addr_hi, data, vector_control) */
#define NV_VGPU_MSIX_ENTRY_SIZE 16
#define NV_VGPU_MSIX_TABLE_ENTRIES 32
#define NV_VGPU_MSIX_TABLE_SIZE (NV_VGPU_MSIX_ENTRY_SIZE * NV_VGPU_MSIX_TABLE_ENTRIES)
#define NV_VGPU_MSIX_PBA_SIZE 8  /* 64 bits for 32 vectors */

static const struct {
	u32 offset;
	u32 size;
	const char *name;
} nv_vgpu_bar0_traps[] = {
	{ 0x000000, 0x1000, "PMC" },
	{ 0x001000, 0x1000, "PRAMIN" },
	{ 0x002000, 0x1000, "PCR" },
	{ 0x003000, 0x1000, "PRAMDAC" },
	{ 0x004000, 0x1000, "PVIO" },
	{ 0x005000, 0x1000, "PFIFO" },
	{ 0x006000, 0x1000, "PGRAPH" },
	{ 0x007000, 0x1000, "PTIMER" },
	{ 0x008000, 0x1000, "PFUSE" },
	{ 0x009000, 0x1000, "PBUS" },
	{ 0x00A000, 0x1000, "PCOP" },
	{ 0x00B000, 0x1000, "PVP" },
	{ 0x00C000, 0x1000, "PCE" },
	{ 0x00D000, 0x1000, "PDISP" },
	{ 0x00E000, 0x1000, "PNVDEC" },
	{ 0x00F000, 0x1000, "PNVENC" },
	{ 0x010000, 0x1000, "PSEC2" },
	{ 0x011000, 0x1000, "PMU" },
	{ 0x012000, 0x1000, "PGFX" },
	{ 0x013000, 0x1000, "PXVE" },
	{ 0x020000, 0x10000, "FIFO" },
	{ 0x040000, 0x10000, "GR" },
	{ 0x050000, 0x10000, "CE" },
	{ 0x060000, 0x10000, "NVDEC" },
	{ 0x070000, 0x10000, "NVENC" },
	{ 0x080000, 0x10000, "NVJPG" },
	{ 0x090000, 0x10000, "OFA" },
	{ 0x100000, 0x10000, "FB" },
	{ 0x110000, 0x10000, "LTC" },
	{ 0x120000, 0x10000, "GPC" },
	{ 0x130000, 0x10000, "TPC" },
	{ 0x140000, 0x10000, "ROP" },
	{ 0x150000, 0x10000, "HUB" },
	{ 0x160000, 0x10000, "NVLINK" },
	{ 0x170000, 0x10000, "GSP" },
	{ 0x180000, 0x10000, "RAP" },
	{ 0x190000, 0x10000, "SEC" },
	{ 0x1A0000, 0x10000, "PMU_RISCV" },
	{ 0x1B0000, 0x10000, "FALCON" },
	{ 0x1C0000, 0x10000, "XVE" },
	{ 0x1D0000, 0x10000, "DMCE" },
	{ 0x1E0000, 0x10000, "P2P" },
	{ 0x1F0000, 0x10000, "MISC" },
};

static void nv_vgpu_msix_init(struct nv_vgpu_bar_emul_ctx *ctx)
{
	memset(ctx->msix_table, 0, sizeof(ctx->msix_table));
	ctx->msix_pba = 0;
	spin_lock_init(&ctx->msix_lock);
	ctx->msix_enabled = false;
}

static void nv_vgpu_msix_write_table(struct nv_vgpu_bar_emul_ctx *ctx,
				      u32 offset, u64 val, u32 size)
{
	u32 entry_idx = offset / NV_VGPU_MSIX_ENTRY_SIZE;
	u32 entry_offset = offset % NV_VGPU_MSIX_ENTRY_SIZE;
	
	if (entry_idx >= NV_VGPU_MSIX_TABLE_ENTRIES)
		return;
	
	spin_lock(&ctx->msix_lock);
	
	switch (entry_offset) {
	case 0x0:  /* Message Address Low */
		ctx->msix_table[entry_idx].addr = (ctx->msix_table[entry_idx].addr & 0xFFFFFFFF00000000ULL) | val;
		break;
	case 0x4:  /* Message Address High */
		ctx->msix_table[entry_idx].addr = (ctx->msix_table[entry_idx].addr & 0x00000000FFFFFFFFULL) | (val << 32);
		break;
	case 0x8:  /* Message Data */
		ctx->msix_table[entry_idx].data = (u32)val;
		break;
	case 0xC:  /* Vector Control */
		ctx->msix_table[entry_idx].vector_control = (u32)val;
		break;
	}
	
	spin_unlock(&ctx->msix_lock);
}

static u64 nv_vgpu_msix_read_table(struct nv_vgpu_bar_emul_ctx *ctx,
				    u32 offset, u32 size)
{
	u32 entry_idx = offset / NV_VGPU_MSIX_ENTRY_SIZE;
	u32 entry_offset = offset % NV_VGPU_MSIX_ENTRY_SIZE;
	u64 val = 0;
	
	if (entry_idx >= NV_VGPU_MSIX_TABLE_ENTRIES)
		return 0;
	
	spin_lock(&ctx->msix_lock);
	
	switch (entry_offset) {
	case 0x0:  /* Message Address Low */
		val = ctx->msix_table[entry_idx].addr & 0xFFFFFFFF;
		break;
	case 0x4:  /* Message Address High */
		val = (ctx->msix_table[entry_idx].addr >> 32) & 0xFFFFFFFF;
		break;
	case 0x8:  /* Message Data */
		val = ctx->msix_table[entry_idx].data;
		break;
	case 0xC:  /* Vector Control */
		val = ctx->msix_table[entry_idx].vector_control;
		break;
	}
	
	spin_unlock(&ctx->msix_lock);
	return val;
}

static u64 nv_vgpu_msix_read_pba(struct nv_vgpu_bar_emul_ctx *ctx, u32 size)
{
	u64 val;
	spin_lock(&ctx->msix_lock);
	val = ctx->msix_pba;
	spin_unlock(&ctx->msix_lock);
	return val;
}

void nv_vgpu_msix_set_pending(struct nv_vgpu_device *vgpu, int vector)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	
	if (!ctx || vector < 0 || vector >= NV_VGPU_MSIX_TABLE_ENTRIES)
		return;
	
	spin_lock(&ctx->msix_lock);
	ctx->msix_pba |= (1ULL << vector);
	spin_unlock(&ctx->msix_lock);
}

void nv_vgpu_msix_clear_pending(struct nv_vgpu_device *vgpu, int vector)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	
	if (!ctx || vector < 0 || vector >= NV_VGPU_MSIX_TABLE_ENTRIES)
		return;
	
	spin_lock(&ctx->msix_lock);
	ctx->msix_pba &= ~(1ULL << vector);
	spin_unlock(&ctx->msix_lock);
}

int nv_vgpu_bar_emul_init(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_bar_emul_ctx *ctx;
	int i, ret;

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

	ctx->vgpu = vgpu;
	ctx->bar0_shadow_size = NV_VGPU_BAR0_SIZE;
	ctx->bar0_shadow = vzalloc(ctx->bar0_shadow_size);
	if (!ctx->bar0_shadow) {
		ret = -ENOMEM;
		goto err_free;
	}

	spin_lock_init(&ctx->shadow_lock);
	ctx->trap_count = 0;

	for (i = 0; i < ARRAY_SIZE(nv_vgpu_bar0_traps); i++) {
		if (ctx->trap_count >= 256)
			break;
		ctx->trap_regs[ctx->trap_count].offset = nv_vgpu_bar0_traps[i].offset;
		ctx->trap_regs[ctx->trap_count].trapped = true;
		ctx->trap_count++;
	}

	/* Initialize MSI-X emulation */
	nv_vgpu_msix_init(ctx);

	vgpu->bar_emul_ctx = ctx;
	return 0;

err_free:
	kfree(ctx);
	return ret;
}

void nv_vgpu_bar_emul_cleanup(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;

	if (!ctx)
		return;

	if (ctx->bar0_shadow)
		vfree(ctx->bar0_shadow);

	kfree(ctx);
	vgpu->bar_emul_ctx = NULL;
}

int nv_vgpu_bar_trap_register(struct nv_vgpu_device *vgpu, int bar_idx,
			       u32 offset, u32 size)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	unsigned long flags;
	int i;

	if (!ctx || bar_idx != 0)
		return -EINVAL;

	spin_lock_irqsave(&ctx->shadow_lock, flags);

	for (i = 0; i < ctx->trap_count; i++) {
		if (ctx->trap_regs[i].offset == offset) {
			ctx->trap_regs[i].trapped = true;
			spin_unlock_irqrestore(&ctx->shadow_lock, flags);
			return 0;
		}
	}

	if (ctx->trap_count >= 256) {
		spin_unlock_irqrestore(&ctx->shadow_lock, flags);
		return -ENOSPC;
	}

	ctx->trap_regs[ctx->trap_count].offset = offset;
	ctx->trap_regs[ctx->trap_count].trapped = true;
	ctx->trap_count++;

	spin_unlock_irqrestore(&ctx->shadow_lock, flags);
	return 0;
}

int nv_vgpu_bar_untrap_register(struct nv_vgpu_device *vgpu, int bar_idx,
				 u32 offset)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	unsigned long flags;
	int i;

	if (!ctx || bar_idx != 0)
		return -EINVAL;

	spin_lock_irqsave(&ctx->shadow_lock, flags);

	for (i = 0; i < ctx->trap_count; i++) {
		if (ctx->trap_regs[i].offset == offset) {
			ctx->trap_regs[i].trapped = false;
			spin_unlock_irqrestore(&ctx->shadow_lock, flags);
			return 0;
		}
	}

	spin_unlock_irqrestore(&ctx->shadow_lock, flags);
	return -ENOENT;
}

int nv_vgpu_bar_emul_read(struct nv_vgpu_device *vgpu, int bar_idx,
			   u64 offset, void *val, size_t size)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	unsigned long flags;
	u32 *val32 = val;
	u64 *val64 = val;

	if (!ctx || bar_idx != 0 || offset >= ctx->bar0_shadow_size)
		return -EINVAL;

	/* Handle MSI-X Table reads (at BAR0+0x2000) */
	if (offset >= GSP_EMU_MSIX_TABLE_OFFSET && 
	    offset < GSP_EMU_MSIX_TABLE_OFFSET + NV_VGPU_MSIX_TABLE_SIZE) {
		u32 msix_offset = offset - GSP_EMU_MSIX_TABLE_OFFSET;
		*val64 = nv_vgpu_msix_read_table(ctx, msix_offset, size);
		return 0;
	}

	/* Handle MSI-X PBA reads (at BAR0+0x3000) */
	if (offset >= GSP_EMU_MSIX_PBA_OFFSET && 
	    offset < GSP_EMU_MSIX_PBA_OFFSET + NV_VGPU_MSIX_PBA_SIZE) {
		*val64 = nv_vgpu_msix_read_pba(ctx, size);
		return 0;
	}

	spin_lock_irqsave(&ctx->shadow_lock, flags);

	switch (size) {
	case 4:
		*val32 = *(u32 *)((u8 *)ctx->bar0_shadow + offset);
		break;
	case 8:
		*val64 = *(u64 *)((u8 *)ctx->bar0_shadow + offset);
		break;
	default:
		spin_unlock_irqrestore(&ctx->shadow_lock, flags);
		return -EINVAL;
	}

	spin_unlock_irqrestore(&ctx->shadow_lock, flags);
	return 0;
}

int nv_vgpu_bar_emul_write(struct nv_vgpu_device *vgpu, int bar_idx,
			    u64 offset, void *val, size_t size)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	unsigned long flags;
	u32 *val32 = val;
	u64 *val64 = val;
	bool trapped = false;
	int i;

	if (!ctx || bar_idx != 0 || offset >= ctx->bar0_shadow_size)
		return -EINVAL;

	/* Handle MSI-X Table writes (at BAR0+0x2000) */
	if (offset >= GSP_EMU_MSIX_TABLE_OFFSET && 
	    offset < GSP_EMU_MSIX_TABLE_OFFSET + NV_VGPU_MSIX_TABLE_SIZE) {
		u32 msix_offset = offset - GSP_EMU_MSIX_TABLE_OFFSET;
		nv_vgpu_msix_write_table(ctx, msix_offset, *val64, size);
		return 0;
	}

	spin_lock_irqsave(&ctx->shadow_lock, flags);

	for (i = 0; i < ctx->trap_count; i++) {
		if (ctx->trap_regs[i].offset == offset && ctx->trap_regs[i].trapped) {
			trapped = true;
			ctx->trap_regs[i].value = (size == 4) ? *val32 : (u32)*val64;
			ctx->trap_regs[i].is_write = true;
			break;
		}
	}

	switch (size) {
	case 4:
		*(u32 *)((u8 *)ctx->bar0_shadow + offset) = *val32;
		break;
	case 8:
		*(u64 *)((u8 *)ctx->bar0_shadow + offset) = *val64;
		break;
	default:
		spin_unlock_irqrestore(&ctx->shadow_lock, flags);
		return -EINVAL;
	}

	spin_unlock_irqrestore(&ctx->shadow_lock, flags);

	if (trapped) {
		nv_vgpu_dbg(vgpu, "BAR0 trap: offset=0x%llx, size=%zu, write=0x%x",
			    offset, size, (size == 4) ? *val32 : (u32)*val64);
		nv_vgpu_gsp_rpc_bar_trap(vgpu, offset, 1, size, *(u64 *)val, 0, NULL);
	}

	return 0;
}

int nv_vgpu_bar_emul_handle_fault(struct nv_vgpu_device *vgpu,
				   unsigned long addr, bool write)
{
	struct nv_vgpu_bar_emul_ctx *ctx = vgpu->bar_emul_ctx;
	struct nv_vgpu_bar *bar = &vgpu->bars[0];
	u64 offset = addr - (unsigned long)bar->vaddr;
	u32 val = 0;
	int ret;

	if (!ctx || !bar->mapped || !bar->trapped)
		return -EINVAL;

	if (write) {
		ret = nv_vgpu_bar_emul_write(vgpu, 0, offset, &val, 4);
	} else {
		ret = nv_vgpu_bar_emul_read(vgpu, 0, offset, &val, 4);
	}

	return ret;
}