#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 "nvidia-vgpu-vfio.h"
#include "bar_emulation.h"
#include "gsp_comm.h"

#define NV_VGPU_BAR0_TRAP_REGIONS 64

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

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

	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;

	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;

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