/*
 * Physical GPU ownership for the real-GSP execution backend.
 *
 * The vGPU work path executes guest pushes on the physical GPU. To talk to
 * the real GSP firmware the module must own the device: the distribution
 * driver is unbound (reversibly) and this module's pci_driver claims the
 * GPU, maps its BARs, and exposes the registers the backend needs.
 *
 * Disabled by default (gsp_phys_gpu=N). Claim only when no VM holds an mdev.
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/pci.h>
#include <linux/delay.h>
#include <linux/io.h>
#include <linux/slab.h>
#include <linux/elf.h>
#include <linux/crc32.h>
#include <linux/firmware.h>
#include <linux/string.h>

#include "gsp_phys_gpu.h"
#include "gsp_phys_boot.h"

#define GSP_PHYS_VENDOR		0x10de
#define GSP_PHYS_DEVICE		0x1ed3

/*
 * GSP firmware to run on the claimed GPU (backend rung 2). Default matches
 * the distribution driver that manages the physical GPU on this host; can be
 * overridden with gsp_phys_fw=<name, without /lib/firmware prefix>.
 */
#define GSP_PHYS_FW_DEFAULT	"nvidia/570.211.01/gsp_tu10x.bin"

static char *gsp_phys_fw_name = GSP_PHYS_FW_DEFAULT;
module_param_named(gsp_phys_fw, gsp_phys_fw_name, charp, 0644);
MODULE_PARM_DESC(gsp_phys_fw, "GSP firmware image to load (default " GSP_PHYS_FW_DEFAULT ")");

/* BAR0 registers sampled after claim as proof of live access. */
#define GSP_PHYS_BOOT_0		0x00000000
#define GSP_PHYS_BOOT_1		0x00000004
#define GSP_PHYS_FALCON_OS	0x00110080
#define GSP_PHYS_GSP_MAILBOX0	0x00110040
#define GSP_PHYS_GSP_MAILBOX1	0x00110044
#define GSP_PHYS_GSP_CPUCTL	0x00110100
#define GSP_PHYS_ENGINE		0x001103c0
#define GSP_PHYS_RISCV_STATUS	0x00111240
#define GSP_PHYS_WPR2_ADDR_LO	0x001fa824
#define GSP_PHYS_WPR2_ADDR_HI	0x001fa828
#define GSP_PHYS_SEC2_CPUCTL	0x00840100
#define GSP_PHYS_SEC2_MAILBOX0	0x00840040
#define GSP_PHYS_SEC2_MAILBOX1	0x00840044

/* Falcon CPUCTL bits (same encoding on the GSP and SEC2 falcons). */
#define GSP_PHYS_CPUCTL_STARTCPU	BIT(1)
#define GSP_PHYS_CPUCTL_HALTED		BIT(4)
#define GSP_PHYS_RISCV_ACTIVE		BIT(0)

struct gsp_phys_gpu {
	struct pci_dev *pdev;
	void __iomem *bar0;
	void __iomem *bar1;
	void __iomem *bar3;
	resource_size_t bar0_len;
	resource_size_t bar1_len;
	resource_size_t bar3_len;
	const struct firmware *fw;
	u32 fw_crc32;
	u64 fw_image_size;
	bool fw_loaded;
	bool enabled;
	bool regions;
	bool drv_registered;
	bool claimed;
	struct gsp_phys_boot boot;
};

static struct gsp_phys_gpu phys;

bool gsp_phys_gpu_claimed(void)
{
	return phys.claimed;
}
EXPORT_SYMBOL_GPL(gsp_phys_gpu_claimed);

void __iomem *gsp_phys_gpu_bar0(void)
{
	return phys.bar0;
}
EXPORT_SYMBOL_GPL(gsp_phys_gpu_bar0);

static void gsp_phys_fw_unload(void)
{
	if (!phys.fw)
		return;

	release_firmware(phys.fw);
	phys.fw = NULL;
	phys.fw_loaded = false;
	phys.fw_crc32 = 0;
	phys.fw_image_size = 0;
}

static const char *gsp_phys_fw_section_name(const struct firmware *fw,
					    const Elf64_Shdr *shdr)
{
	const Elf64_Ehdr *eh = (const Elf64_Ehdr *)fw->data;
	const Elf64_Shdr *shstr;
	const char *base;

	if (eh->e_shstrndx >= eh->e_shnum)
		return "?";
	shstr = (const Elf64_Shdr *)(fw->data + eh->e_shoff +
				     (u64)eh->e_shstrndx * eh->e_shentsize);
	if (shstr->sh_offset + shstr->sh_size > fw->size ||
	    shdr->sh_name >= shstr->sh_size)
		return "?";
	base = (const char *)fw->data + shstr->sh_offset;
	return base + shdr->sh_name;
}

/*
 * Rung 2: fetch and validate the real GSP firmware image while the device
 * is owned. No GPU state is touched here - the image is held for the boot
 * rung. The image is the ELF the distribution driver would load
 * (nvidia/<version>/gsp_tu10x.bin).
 */
static int gsp_phys_fw_load(void)
{
	const Elf64_Ehdr *eh;
	int ret;

	ret = request_firmware(&phys.fw, gsp_phys_fw_name,
			       phys.pdev ? &phys.pdev->dev : NULL);
	if (ret) {
		pr_err("gsp-phys-fw: request_firmware(%s) failed: %d\n",
		       gsp_phys_fw_name, ret);
		return ret;
	}

	eh = (const Elf64_Ehdr *)phys.fw->data;
	if (phys.fw->size < sizeof(*eh) ||
	    memcmp(eh->e_ident, ELFMAG, SELFMAG)) {
		pr_err("gsp-phys-fw: %s is not an ELF image (%zu bytes)\n",
		       gsp_phys_fw_name, phys.fw->size);
		goto err;
	}
	if (eh->e_ident[EI_CLASS] != ELFCLASS64 ||
	    eh->e_machine != EM_RISCV) {
		pr_err("gsp-phys-fw: unexpected ELF (class=%u machine=%u)\n",
		       eh->e_ident[EI_CLASS], eh->e_machine);
		goto err;
	}

	phys.fw_crc32 = crc32_le(0, phys.fw->data, phys.fw->size);
	phys.fw_loaded = true;

	pr_info("gsp-phys-fw: loaded %s: %zu bytes, ELF64 RISC-V, "
		"entry=0x%llx shnum=%u crc32=0x%08x\n",
		gsp_phys_fw_name, phys.fw->size, (u64)eh->e_entry,
		eh->e_shnum, phys.fw_crc32);

	if (eh->e_shoff && eh->e_shnum &&
	    eh->e_shentsize == sizeof(Elf64_Shdr) &&
	    eh->e_shoff + (u64)eh->e_shnum * eh->e_shentsize <= phys.fw->size) {
		unsigned int i;

		for (i = 0; i < eh->e_shnum; i++) {
			const Elf64_Shdr *sh = (const Elf64_Shdr *)
				(phys.fw->data + eh->e_shoff +
				 (u64)i * eh->e_shentsize);
			const char *name;

			if (sh->sh_offset > phys.fw->size ||
			    sh->sh_size > phys.fw->size - sh->sh_offset)
				continue;
			name = gsp_phys_fw_section_name(phys.fw, sh);
			if (!strcmp(name, ".fwimage"))
				phys.fw_image_size = sh->sh_size;
			pr_info("gsp-phys-fw:   [%02u] %-24s type=%u addr=0x%llx "
				"off=0x%llx size=0x%llx\n", i, name,
				sh->sh_type, (u64)sh->sh_addr,
				(u64)sh->sh_offset, (u64)sh->sh_size);
		}

		if (phys.fw_image_size)
			pr_info("gsp-phys-fw: GSP-RM image section .fwimage "
				"size=0x%llx (boot-layout input)\n",
				phys.fw_image_size);
	}

	return 0;

err:
	gsp_phys_fw_unload();
	return -EINVAL;
}

static void gsp_phys_unmap(struct pci_dev *pdev)
{
	if (phys.bar3) {
		pci_iounmap(pdev, phys.bar3);
		phys.bar3 = NULL;
	}
	if (phys.bar1) {
		pci_iounmap(pdev, phys.bar1);
		phys.bar1 = NULL;
	}
	if (phys.bar0) {
		pci_iounmap(pdev, phys.bar0);
		phys.bar0 = NULL;
	}
	if (phys.regions) {
		pci_release_regions(pdev);
		phys.regions = false;
	}
	if (phys.enabled) {
		pci_disable_device(pdev);
		phys.enabled = false;
	}
}

/*
 * Rung 3 stage 3a: read-only falcon state probe. Logs the boot-relevant
 * register state before any boot attempt: both falcons' run/halt bits, the
 * RISC-V active bit, WPR2 up/down and the security-falcon mailboxes (the
 * booter reports its result there).
 */
static void gsp_phys_dump_falcon(void)
{
	u32 gsp_cpuctl = readl(phys.bar0 + GSP_PHYS_GSP_CPUCTL);
	u32 sec2_cpuctl = readl(phys.bar0 + GSP_PHYS_SEC2_CPUCTL);
	u32 riscv = readl(phys.bar0 + GSP_PHYS_RISCV_STATUS);
	u32 wpr2_lo = readl(phys.bar0 + GSP_PHYS_WPR2_ADDR_LO);
	u32 wpr2_hi = readl(phys.bar0 + GSP_PHYS_WPR2_ADDR_HI);

	pr_info("gsp-phys: GSP falcon CPUCTL=0x%08x [%s%s]\n", gsp_cpuctl,
		gsp_cpuctl & GSP_PHYS_CPUCTL_STARTCPU ? "started " : "",
		gsp_cpuctl & GSP_PHYS_CPUCTL_HALTED ? "halted" : "running");
	pr_info("gsp-phys: SEC2 falcon CPUCTL=0x%08x [%s%s] MAILBOX0=0x%08x "
		"MAILBOX1=0x%08x\n", sec2_cpuctl,
		sec2_cpuctl & GSP_PHYS_CPUCTL_STARTCPU ? "started " : "",
		sec2_cpuctl & GSP_PHYS_CPUCTL_HALTED ? "halted" : "running",
		readl(phys.bar0 + GSP_PHYS_SEC2_MAILBOX0),
		readl(phys.bar0 + GSP_PHYS_SEC2_MAILBOX1));
	pr_info("gsp-phys: RISCV_STATUS=0x%08x [%sactive] WPR2_LO=0x%08x "
		"WPR2_HI=0x%08x [%s]\n", riscv,
		riscv & GSP_PHYS_RISCV_ACTIVE ? "" : "not ",
		wpr2_lo, wpr2_hi, wpr2_hi ? "up" : "down");
}

static int gsp_phys_probe(struct pci_dev *pdev, const struct pci_device_id *id)
{
	u32 boot0, boot1;
	int ret;

	pr_info("gsp-phys: probe %s\n", pci_name(pdev));

	ret = pci_enable_device(pdev);
	if (ret) {
		pr_err("gsp-phys: pci_enable_device failed: %d\n", ret);
		return ret;
	}
	phys.enabled = true;

	ret = pci_request_regions(pdev, "open-iov-gsp-phys");
	if (ret) {
		pr_err("gsp-phys: pci_request_regions failed: %d\n", ret);
		goto err_unmap;
	}
	phys.regions = true;

	phys.bar0 = pci_iomap(pdev, 0, 0);
	phys.bar0_len = pci_resource_len(pdev, 0);
	if (!phys.bar0) {
		pr_err("gsp-phys: BAR0 map failed\n");
		ret = -ENOMEM;
		goto err_unmap;
	}
	phys.bar1 = pci_iomap(pdev, 1, 0);
	phys.bar1_len = pci_resource_len(pdev, 1);
	if (!phys.bar1) {
		pr_err("gsp-phys: BAR1 map failed\n");
		ret = -ENOMEM;
		goto err_unmap;
	}
	phys.bar3 = pci_iomap(pdev, 3, 0);
	phys.bar3_len = pci_resource_len(pdev, 3);
	if (!phys.bar3) {
		pr_err("gsp-phys: BAR3 map failed\n");
		ret = -ENOMEM;
		goto err_unmap;
	}

	boot0 = readl(phys.bar0 + GSP_PHYS_BOOT_0);
	boot1 = readl(phys.bar0 + GSP_PHYS_BOOT_1);
	pr_info("gsp-phys: BAR0 len=0x%llx BAR1 len=0x%llx BAR3 len=0x%llx\n",
		(u64)phys.bar0_len, (u64)phys.bar1_len, (u64)phys.bar3_len);
	pr_info("gsp-phys: BOOT0=0x%08x BOOT1=0x%08x (chip family in BOOT0[15:4])\n",
		boot0, boot1);
	pr_info("gsp-phys: GSP_CPUCTL=0x%08x MAILBOX0=0x%08x MAILBOX1=0x%08x "
		"ENGINE=0x%08x FALCON_OS=0x%08x\n",
		readl(phys.bar0 + GSP_PHYS_GSP_CPUCTL),
		readl(phys.bar0 + GSP_PHYS_GSP_MAILBOX0),
		readl(phys.bar0 + GSP_PHYS_GSP_MAILBOX1),
		readl(phys.bar0 + GSP_PHYS_ENGINE),
		readl(phys.bar0 + GSP_PHYS_FALCON_OS));

	gsp_phys_dump_falcon();

	phys.pdev = pdev;
	phys.claimed = true;
	pr_info("gsp-phys: claimed %s\n", pci_name(pdev));
	return 0;

err_unmap:
	gsp_phys_unmap(pdev);
	return ret;
}

static void gsp_phys_remove(struct pci_dev *pdev)
{
	pr_info("gsp-phys: remove %s\n", pci_name(pdev));

	phys.claimed = false;
	gsp_phys_unmap(pdev);
	phys.pdev = NULL;
}

static const struct pci_device_id gsp_phys_ids[] = {
	{ PCI_DEVICE(GSP_PHYS_VENDOR, GSP_PHYS_DEVICE) },
	{ }
};
MODULE_DEVICE_TABLE(pci, gsp_phys_ids);

static struct pci_driver gsp_phys_driver = {
	.name = "open-iov-gsp-phys",
	.id_table = gsp_phys_ids,
	.probe = gsp_phys_probe,
	.remove = gsp_phys_remove,
};

int gsp_phys_gpu_claim(void)
{
	struct pci_dev *pdev;
	int ret;

	if (phys.claimed)
		return 0;

	pdev = pci_get_device(GSP_PHYS_VENDOR, GSP_PHYS_DEVICE, NULL);
	if (!pdev) {
		pr_err("gsp-phys: device %04x:%04x not found\n",
		       GSP_PHYS_VENDOR, GSP_PHYS_DEVICE);
		return -ENODEV;
	}

	if (pdev->dev.driver) {
		pr_info("gsp-phys: unbinding %s from %s\n",
			pdev->dev.driver->name, pci_name(pdev));
		device_release_driver(&pdev->dev);
	}
	if (pdev->dev.driver) {
		pr_err("gsp-phys: %s still bound after unbind (usage count?)\n",
		       pdev->dev.driver->name);
		pci_dev_put(pdev);
		return -EBUSY;
	}

	ret = pci_register_driver(&gsp_phys_driver);
	if (ret) {
		pr_err("gsp-phys: register driver failed: %d\n", ret);
		pci_dev_put(pdev);
		return ret;
	}
	phys.drv_registered = true;

	if (!phys.claimed) {
		pr_err("gsp-phys: our driver did not bind %s\n", pci_name(pdev));
		pci_unregister_driver(&gsp_phys_driver);
		phys.drv_registered = false;
		pci_dev_put(pdev);
		return -ENODEV;
	}

	pci_dev_put(pdev);
	pr_info("gsp-phys: physical GPU claimed (backend rung 1)\n");

	if (gsp_phys_fw_load()) {
		pr_warn("gsp-phys: firmware load failed; backend stays at rung 1\n");
	} else {
		pr_info("gsp-phys: real GSP firmware staged (backend rung 2)\n");

		if (gsp_phys_layout_enabled() &&
		    gsp_phys_boot_layout_init(&phys.boot, phys.pdev,
					      phys.fw_image_size, phys.bar1_len))
			pr_warn("gsp-phys: boot layout/allocation failed; "
				"backend stays at rung 2\n");
	}
	return 0;
}
EXPORT_SYMBOL_GPL(gsp_phys_gpu_claim);

void gsp_phys_gpu_release(void)
{
	struct pci_dev *pdev;

	if (!phys.drv_registered && !phys.claimed)
		return;

	if (phys.pdev)
		gsp_phys_boot_layout_fini(&phys.boot, phys.pdev);

	gsp_phys_fw_unload();

	pdev = pci_get_device(GSP_PHYS_VENDOR, GSP_PHYS_DEVICE, NULL);

	if (phys.drv_registered) {
		pci_unregister_driver(&gsp_phys_driver);
		phys.drv_registered = false;
	}

	if (pdev) {
		int ret = device_attach(&pdev->dev);
		struct device_driver *drv = pdev->dev.driver;

		if (drv)
			pr_info("gsp-phys: released %s; rebound to %s (rc=%d)\n",
				pci_name(pdev), drv->name, ret);
		else
			pr_warn("gsp-phys: released %s; no driver rebound "
				"(rc=%d), run `modprobe nvidia`\n",
				pci_name(pdev), ret);
		pci_dev_put(pdev);
	}
	phys.claimed = false;
}