#include <linux/slab.h>
#include <linux/pci.h>
#include <linux/iommu.h>
#include <linux/vmalloc.h>

#include "nvidia-vgpu-vfio.h"
#include "bar_emulation.h"
#include "gsp_phys_gpu.h"
#include "irq.h"

int nv_vgpu_bar_init(struct nv_vgpu_device *vgpu)
{
	struct pci_dev *pdev = vgpu->pdev;
	struct nv_vgpu_bar *bar;
	u64 bar_addr;
	u64 bar_size;
	int i, ret;

	if (vgpu->dry_run) {
		nv_vgpu_info(vgpu, "DRY-RUN: BAR init validation");
		for (i = 0; i < 4; i++) {
			bar_addr = pci_resource_start(pdev, i);
			bar_size = pci_resource_len(pdev, i);
			if (bar_addr && bar_size) {
				nv_vgpu_info(vgpu, "  BAR%d: gpa=0x%llx, size=0x%llx (%s)", 
					     i, bar_addr, bar_size,
					     i == 0 ? "config" : (i == 1 ? "FB" : (i == 2 ? "I/O" : "other")));
			}
		}
		nv_vgpu_info(vgpu, "DRY-RUN: BAR validation passed");
		return 0;
	}

	for (i = 0; i < NV_VGPU_MAX_REGIONS; i++) {
		bar = &vgpu->bars[i];
		mutex_init(&bar->lock);
		bar->mapped = false;
		bar->trapped = false;
	}

	bar_addr = pci_resource_start(pdev, 0);
	bar_size = pci_resource_len(pdev, 0);
	if (bar_addr && bar_size) {
		vgpu->bars[0].gpa = bar_addr;
		vgpu->bars[0].size = bar_size;
		vgpu->bars[0].vaddr = ioremap_wc(bar_addr, bar_size);
		if (!vgpu->bars[0].vaddr) {
			nv_vgpu_err(vgpu, "Failed to map BAR0");
			return -ENOMEM;
		}
		vgpu->bar0_config_base = bar_addr;
		nv_vgpu_info(vgpu, "BAR0: gpa=0x%llx, size=0x%llx, vaddr=%p", bar_addr, bar_size, vgpu->bars[0].vaddr);
	}

	bar_addr = pci_resource_start(pdev, 1);
	bar_size = pci_resource_len(pdev, 1);
	if (bar_addr && bar_size) {
		vgpu->bars[1].gpa = bar_addr;
		vgpu->bars[1].size = bar_size;
		/* Emulated vidmem: a dedicated host-RAM buffer backing the guest's
		 * VRAM aperture. Do NOT map the host GPU's real BAR1 - guest VRAM
		 * accesses must hit our own clean buffer, not the (broken) host
		 * GPU's VRAM. Allocated lazily on first use to keep mdev create
		 * cheap; 256MB to match the GSCI fb size. */
		vgpu->bars[1].vaddr = vzalloc(nv_vgpu_fb_emul_size(vgpu));
		if (!vgpu->bars[1].vaddr) {
			nv_vgpu_err(vgpu, "Failed to allocate emulated vidmem (%llu MiB)",
				    (unsigned long long)(nv_vgpu_fb_emul_size(vgpu) >> 20));
			return -ENOMEM;
		}
		nv_vgpu_info(vgpu, "BAR1: gpa=0x%llx, size=0x%llx, emulated vidmem=%p",
			     bar_addr, bar_size, vgpu->bars[1].vaddr);
	}

	bar_addr = pci_resource_start(pdev, 2);
	bar_size = pci_resource_len(pdev, 2);
	if (bar_addr && bar_size) {
		vgpu->bars[2].gpa = bar_addr;
		vgpu->bars[2].size = bar_size;
		/*
		 * Emulation-only mode (gsp_phys_gpu=0): the host nvidia driver
		 * owns the physical GPU. Never ioremap the real BAR2 nor expose
		 * its PFNs to the guest - guest BAR2 MMIO against a device the
		 * distribution driver owns is the suspected trigger of the Xid
		 * 140 uncorrectable-ECC incident (2026-10-03). Back the region
		 * with a private host-RAM shadow instead, exactly like BAR1/BAR3.
		 * Claim/backend mode keeps the real mapping.
		 */
		if (gsp_phys_gpu_enabled()) {
			vgpu->bars[2].vaddr = ioremap_wc(bar_addr, bar_size);
			if (!vgpu->bars[2].vaddr) {
				nv_vgpu_err(vgpu, "Failed to map BAR2");
				return -ENOMEM;
			}
			nv_vgpu_info(vgpu, "BAR2: gpa=0x%llx, size=0x%llx, vaddr=%p (real, claim mode)",
				     bar_addr, bar_size, vgpu->bars[2].vaddr);
		} else {
			vgpu->bars[2].vaddr = vzalloc(bar_size);
			vgpu->bars[2].emulated = true;
			if (!vgpu->bars[2].vaddr) {
				nv_vgpu_err(vgpu, "Failed to allocate emulated BAR2 (%llu bytes)",
					    (unsigned long long)bar_size);
				return -ENOMEM;
			}
			nv_vgpu_info(vgpu, "BAR2: gpa=0x%llx, size=0x%llx, emulated shadow=%p",
				     bar_addr, bar_size, vgpu->bars[2].vaddr);
		}
	}

	bar_addr = pci_resource_start(pdev, 3);
	bar_size = pci_resource_len(pdev, 3);
	if (bar_addr && bar_size) {
		vgpu->bars[3].gpa = bar_addr;
		vgpu->bars[3].size = bar_size;
		/* Emulated BAR3: dedicated host-RAM shadow. Do NOT ioremap the
		 * host GPU's real BAR3 - guest writes through this region must
		 * hit our own clean buffer, not the (running) host GPU's
		 * instance memory window. 32MB to match the GSCI-reported BAR3. */
		vgpu->bars[3].vaddr = vzalloc(bar_size);
		if (!vgpu->bars[3].vaddr) {
			nv_vgpu_err(vgpu, "Failed to allocate emulated BAR3 (%llu bytes)",
				    (unsigned long long)bar_size);
			return -ENOMEM;
		}
		nv_vgpu_info(vgpu, "BAR3: gpa=0x%llx, size=0x%llx, emulated shadow=%p",
			     bar_addr, bar_size, vgpu->bars[3].vaddr);
	}

	ret = nv_vgpu_bar_emul_init(vgpu);
	if (ret) {
		nv_vgpu_err(vgpu, "BAR emulation init failed: %d", ret);
		return ret;
	}

	/* MSI/MSI-X is managed by nvidia.ko - don't try to set up our own */
	nv_vgpu_info(vgpu, "MSI/MSI-X managed by nvidia.ko");

	vgpu->fb_base = pci_resource_start(pdev, 1);
	vgpu->fb_size = pci_resource_len(pdev, 1);

	return 0;
}

void nv_vgpu_bar_cleanup(struct nv_vgpu_device *vgpu)
{
	int i;

	if (vgpu->dry_run) {
		nv_vgpu_info(vgpu, "DRY-RUN: BAR cleanup (no hardware resources to release)");
		return;
	}

	nv_vgpu_irq_teardown_msix(vgpu);
	nv_vgpu_irq_teardown_msi(vgpu);
	nv_vgpu_bar_emul_cleanup(vgpu);

	/* Free the emulated vidmem buffer (vzalloc'd, not ioremap'd) */
	if (vgpu->bars[1].vaddr) {
		vfree(vgpu->bars[1].vaddr);
		vgpu->bars[1].vaddr = NULL;
	}

	/* Free the emulated BAR3 shadow (vzalloc'd, not ioremap'd) */
	if (vgpu->bars[3].vaddr) {
		vfree(vgpu->bars[3].vaddr);
		vgpu->bars[3].vaddr = NULL;
	}

	/* Free the emulated BAR2 shadow (emulation-only mode) */
	if (vgpu->bars[2].emulated && vgpu->bars[2].vaddr) {
		vfree(vgpu->bars[2].vaddr);
		vgpu->bars[2].vaddr = NULL;
	}

	for (i = 0; i < NV_VGPU_MAX_REGIONS; i++) {
		if (i == 1 || i == 3)
			continue; /* shadow buffers already freed above */
		if (i == 2 && vgpu->bars[2].emulated)
			continue; /* shadow buffer already freed above */
		if (vgpu->bars[i].vaddr) {
			iounmap(vgpu->bars[i].vaddr);
			vgpu->bars[i].vaddr = NULL;
		}
		vgpu->bars[i].mapped = false;
		vgpu->bars[i].trapped = false;
		vgpu->bars[i].iova = 0;
	}
}

int nv_vgpu_bar_trap(struct nv_vgpu_device *vgpu, int bar_idx)
{
	if (bar_idx < 0 || bar_idx >= NV_VGPU_MAX_REGIONS)
		return -EINVAL;

	if (!vgpu->bars[bar_idx].vaddr)
		return -ENODEV;

	return nv_vgpu_bar_trap_register(vgpu, bar_idx, 0, vgpu->bars[bar_idx].size);
}

void nv_vgpu_bar_untrap(struct nv_vgpu_device *vgpu, int bar_idx)
{
	if (bar_idx < 0 || bar_idx >= NV_VGPU_MAX_REGIONS)
		return;

	nv_vgpu_bar_untrap_register(vgpu, bar_idx, 0);
}