/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * GSP vGPU-IO Driver - Core Implementation
 * Guest vGPU driver - communicates with GSP already booted by host GRID driver
 */

#include "gsp_core.h"
#include <linux/module.h>
#include <linux/pci.h>
#include <linux/firmware.h>
#include <linux/interrupt.h>
#include <linux/dma-mapping.h>
#include <linux/debugfs.h>
#include <linux/seq_file.h>
#include <linux/delay.h>
#include <linux/msi.h>
#include <linux/miscdevice.h>
#include <linux/uaccess.h>

#define DRIVER_NAME "gsp-vgpu"
#define DRIVER_VERSION "0.1.0"

/* Emulation mode: skip GSP RPC for BAR/WPR when GSP lacks vGPU handlers */
bool emulation_mode = true;  /* ON by default */
static bool emulation_override = false;  /* true = force emulation, false = force native */
module_param(emulation_mode, bool, 0444);
MODULE_PARM_DESC(emulation_mode, "Enable emulation mode for BAR/WPR (skip GSP RPC) - default ON, auto-detects if GSP has vGPU handlers");
module_param(emulation_override, bool, 0444);
MODULE_PARM_DESC(emulation_override, "Override auto-detection: 1=force emulation, 0=force native (default: auto-detect)");

/* Forward declarations */
int gsp_firmware_load(struct gsp_device *gsp);
int gsp_setup_irq(struct gsp_device *gsp);
int gsp_rpc_init(struct gsp_device *gsp);
void gsp_shutdown(struct gsp_device *gsp);
static void gsp_detect_emulation_mode(struct gsp_device *gsp);
static long gsp_ioctl(struct file *file, unsigned int cmd, unsigned long arg);
static int gsp_open(struct inode *inode, struct file *file);
static int gsp_release(struct inode *inode, struct file *file);

/* Char device file operations */
static const struct file_operations gsp_fops = {
    .owner = THIS_MODULE,
    .open = gsp_open,
    .release = gsp_release,
    .unlocked_ioctl = gsp_ioctl,
    .compat_ioctl = gsp_ioctl,
};

/* Char device open */
static int gsp_open(struct inode *inode, struct file *file)
{
    struct gsp_device *gsp = container_of(file->private_data, struct gsp_device, miscdev);
    file->private_data = gsp;
    return 0;
}

/* Char device release */
static int gsp_release(struct inode *inode, struct file *file)
{
    return 0;
}

/* Char device ioctl */
static long gsp_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
    struct gsp_device *gsp = file->private_data;
    void __user *argp = (void __user *)arg;
    int ret = 0;
    
    mutex_lock(&gsp->chardev_lock);
    
    switch (cmd) {
    case GSP_IOCTL_CREATE_VF: {
        struct gsp_vf_config config;
        if (copy_from_user(&config, argp, sizeof(config))) {
            ret = -EFAULT;
            break;
        }
        ret = gsp_vf_create(gsp, &config);
        break;
    }
    case GSP_IOCTL_DESTROY_VF: {
        uint32_t vf_id;
        if (copy_from_user(&vf_id, argp, sizeof(vf_id))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        gsp_vf_destroy(vf);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_ALLOC: {
        struct gsp_alloc_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_alloc(vf, args.size, args.flags, &args.handle, &args.va);
        if (!ret)
            ret = copy_to_user(argp, &args, sizeof(args)) ? -EFAULT : 0;
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_FREE: {
        struct gsp_free_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_free(vf, args.handle);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_PIN: {
        struct gsp_pin_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_pin(vf, args.handle, args.offset, args.size, &args.dma_addr);
        if (!ret)
            ret = copy_to_user(argp, &args, sizeof(args)) ? -EFAULT : 0;
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_UNPIN: {
        struct gsp_unpin_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_unpin(vf, args.handle);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_MAP: {
        struct gsp_map_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_map(vf, args.handle, args.gpu_va, args.size,
                             args.prot, args.flags, args.offset, &args.mapped_va);
        if (!ret)
            ret = copy_to_user(argp, &args, sizeof(args)) ? -EFAULT : 0;
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_UNMAP: {
        struct gsp_unmap_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_mem_unmap(vf, args.handle, args.gpu_va, args.size);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_BAR_CONFIG: {
        struct gsp_bar_config_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_bar_init(vf);  /* Re-init with new aperture */
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_BAR_FORCE: {
        struct gsp_bar_force_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_bar_force_access(vf, args.enable);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_BAR_SAVE: {
        struct gsp_bar_save_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_migration_bar_save(vf, args.buf, args.size);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_BAR_RESTORE: {
        struct gsp_bar_restore_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_migration_bar_restore(vf, args.buf, args.size);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_WPR_PROTECT: {
        struct gsp_wpr_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_wpr_protect(vf, args.base, args.size, args.flags);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_WPR_UNPROTECT: {
        struct gsp_wpr_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_wpr_unprotect(vf, args.base);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_WPR_TIMEOUT: {
        struct gsp_wpr_timeout_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_wpr_set_timeout(vf, args.timeout_ns);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_SCHEDULE: {
        struct gsp_schedule_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_schedule(vf, args.priority, args.time_slice_ns);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_YIELD: {
        uint32_t vf_id;
        if (copy_from_user(&vf_id, argp, sizeof(vf_id))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_vf_yield(vf);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_MIGRATE: {
        struct gsp_migrate_args args;
        if (copy_from_user(&args, argp, sizeof(args))) {
            ret = -EFAULT;
            break;
        }
        struct gsp_vf *vf = gsp_vf_get(gsp, args.vf_id);
        if (!vf) {
            ret = -ENOENT;
            break;
        }
        ret = gsp_migration_start(vf, args.dst_vf);
        gsp_vf_put(vf);
        break;
    }
    case GSP_IOCTL_GET_INFO: {
        char info[256];
        snprintf(info, sizeof(info), "GSP vGPU-IO Driver v%s (FW: 570.144W)", DRIVER_VERSION);
        ret = copy_to_user(argp, info, sizeof(info)) ? -EFAULT : 0;
        break;
    }
    default:
        ret = -ENOTTY;
    }
    
    mutex_unlock(&gsp->chardev_lock);
    return ret;
}

/* PCI Device IDs */
static const struct pci_device_id gsp_pci_ids[] = {
    { PCI_DEVICE(0x10DE, 0x1E07), .driver_data = 0 },  /* TU102 */
    { PCI_DEVICE(0x10DE, 0x1E08), .driver_data = 0 },  /* TU104 */
    { PCI_DEVICE(0x10DE, 0x1ED3), .driver_data = 0 },  /* TU104 RTX 2080 Super Mobile */
    { 0 }
};
MODULE_DEVICE_TABLE(pci, gsp_pci_ids);

/* Module Parameters */
static int max_vfs = 7;
module_param(max_vfs, int, 0444);
MODULE_PARM_DESC(max_vfs, "Maximum number of VFs (default: 7)");

static char *firmware_name = "gsp-570.144.fwimage.bin";
module_param(firmware_name, charp, 0444);
MODULE_PARM_DESC(firmware_name, "GSP firmware filename (for version verification only)");

/* PCI Probe */
int gsp_probe(struct pci_dev *pdev, const struct pci_device_id *id)
{
    struct gsp_device *gsp;
    int ret;
    
    gsp = devm_kzalloc(&pdev->dev, sizeof(*gsp), GFP_KERNEL);
    if (!gsp)
        return -ENOMEM;
    
    gsp->pdev = pdev;
    gsp->dev = &pdev->dev;
    gsp->max_vfs = max_vfs;
    pci_set_drvdata(pdev, gsp);
    
    mutex_init(&gsp->vf_lock);
    idr_init(&gsp->vf_idr);
    INIT_LIST_HEAD(&gsp->vf_list);
    
    /* Enable PCI device */
    ret = pcim_enable_device(pdev);
    if (ret) {
        dev_err(&pdev->dev, "Failed to enable PCI device: %d\n", ret);
        return ret;
    }
    
    pci_set_master(pdev);
    
    /* Map BARs */
    gsp->bar0 = pcim_iomap(pdev, 0, 0);
    gsp->bar1 = pcim_iomap(pdev, 1, 0);
    if (!gsp->bar0 || !gsp->bar1) {
        dev_err(&pdev->dev, "Failed to map BARs\n");
        return -ENOMEM;
    }
    
    dev_info(&pdev->dev, "BAR0: %p, BAR1: %p\n", gsp->bar0, gsp->bar1);
    
    /* Load Firmware (for version verification only - GSP already booted by host) */
    ret = gsp_firmware_load(gsp);
    if (ret) {
        dev_err(&pdev->dev, "Firmware load failed: %d\n", ret);
        return ret;
    }
    
    /* Initialize RPC client (sets up rings, doorbells) */
    gsp->rpc = gsp_rpc_client_init(gsp);
    if (IS_ERR(gsp->rpc)) {
        ret = PTR_ERR(gsp->rpc);
        dev_err(&pdev->dev, "RPC client init failed: %d\n", ret);
        return ret;
    }
    
    /* Initialize RPC communication with GSP */
    ret = gsp_rpc_init(gsp);
    if (ret) {
        dev_err(&pdev->dev, "RPC init failed: %d\n", ret);
        return ret;
    }
    
    /* Setup IRQ */
    ret = gsp_setup_irq(gsp);
    if (ret) {
        dev_err(&pdev->dev, "IRQ setup failed: %d\n", ret);
        return ret;
    }
    
    /* Initialize char device */
    mutex_init(&gsp->chardev_lock);
    gsp->miscdev.minor = MISC_DYNAMIC_MINOR;
    gsp->miscdev.name = "gsp-vgpu";
    gsp->miscdev.fops = &gsp_fops;
    gsp->miscdev.parent = &pdev->dev;
    gsp->miscdev.mode = 0660;
    
    ret = misc_register(&gsp->miscdev);
    if (ret) {
        dev_err(&pdev->dev, "Failed to register misc device: %d\n", ret);
        return ret;
    }
    
    /* Debugfs */
    gsp_debugfs_init(gsp);
    
    dev_info(&pdev->dev, "GSP vGPU-IO driver loaded (FW: 570.144W)\n");
    return 0;
}

/* PCI Remove */
void gsp_remove(struct pci_dev *pdev)
{
    struct gsp_device *gsp = pci_get_drvdata(pdev);
    struct gsp_vf *vf, *tmp;
    
    if (!gsp)
        return;
    
    /* Destroy all VFs */
    list_for_each_entry_safe(vf, tmp, &gsp->vf_list, list)
        gsp_vf_destroy(vf);
    
    /* Unregister char device */
    misc_deregister(&gsp->miscdev);
    
    gsp_debugfs_fini(gsp);
    gsp_shutdown(gsp);
    gsp_rpc_client_fini(gsp->rpc);
    
    if (gsp->fw)
        release_firmware(gsp->fw);
    
    idr_destroy(&gsp->vf_idr);
    
    dev_info(&pdev->dev, "GSP vGPU-IO driver unloaded\n");
}

/* Firmware Load - only for version verification */
int gsp_firmware_load(struct gsp_device *gsp)
{
    const struct firmware *fw;
    int ret;
    
    ret = request_firmware(&fw, firmware_name, gsp->dev);
    if (ret) {
        dev_err(gsp->dev, "Failed to load firmware '%s': %d\n", firmware_name, ret);
        return ret;
    }
    
    /* Verify MD5: 13ce12a392a31eeb79a2e6ebd3a4a93a */
    if (fw->size != 0x1C50000) {  /* ~28.5 MB */
        dev_warn(gsp->dev, "Firmware size unexpected: %zu (expected ~28.5MB)\n", fw->size);
    }
    
    dev_info(gsp->dev, "Firmware verified: %zu bytes\n", fw->size);
    gsp->fw = fw;
    return 0;
}

/* Initialize RPC communication with already-running GSP */
int gsp_rpc_init(struct gsp_device *gsp)
{
    /* 
     * The GSP is already running (booted by host GRID driver).
     * We just need to verify communication works by reading mailbox registers
     * and optionally sending a test RPC.
     */
    
    /* Read mailbox registers to verify GSP is responsive */
    uint32_t mbox0 = readl(gsp->bar0 + GSP_BAR0_FALCON_MAILBOX0);
    uint32_t mbox1 = readl(gsp->bar0 + GSP_BAR0_FALCON_MAILBOX1);
    
    dev_info(gsp->dev, "GSP mailbox0: 0x%08x, mailbox1: 0x%08x\n", mbox0, mbox1);
    
    /* Check if GSP appears to be running (non-zero mailbox or queue heads) */
    uint32_t qhead0 = readl(gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0));
    dev_info(gsp->dev, "GSP queue head 0: 0x%08x\n", qhead0);
    
    /* Ring doorbell to signal we're ready - write to queue head */
    writel(1, gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0));
    
    gsp->fw_loaded = true;  /* Indicates GSP communication established */
    dev_info(gsp->dev, "GSP RPC communication initialized\n");
    
    /* Auto-detect emulation mode if not overridden */
    if (!emulation_override) {
        gsp_detect_emulation_mode(gsp);
    } else if (emulation_mode) {
        dev_info(gsp->dev, "Emulation mode FORCED ON by module parameter\n");
    } else {
        dev_info(gsp->dev, "Emulation mode FORCED OFF by module parameter\n");
    }
    
    return 0;
}

/* Auto-detect if GSP has vGPU RPC handlers */
static void gsp_detect_emulation_mode(struct gsp_device *gsp)
{
    struct gsp_bar_trap_payload test_payload = {0};
    struct gsp_bar_trap_response test_resp = {0};
    size_t resp_size = sizeof(test_resp);
    int ret;
    
    dev_info(gsp->dev, "Auto-detecting GSP vGPU RPC handler support...\n");
    
    /* Send a test BAR_TRAP RPC to check if GSP has vGPU handlers */
    test_payload.bar_index = 1;
    test_payload.trap_type = GSP_BAR_READ;  /* Non-destructive read */
    test_payload.offset = 0;
    test_payload.value = 0;
    test_payload.size = 8;
    test_payload.vf_id = 0;
    
    /* Try synchronous RPC with short timeout (1 second) */
    ret = gsp_rpc_send_sync(gsp->rpc, GSP_RPC_BAR_TRAP,
                            &test_payload, sizeof(test_payload),
                            &test_resp, &resp_size);
    
    if (ret == -ETIMEDOUT || ret == -ENOTTY || ret == -EIO || ret == -ENODEV) {
        /* GSP doesn't have vGPU handlers - enable emulation */
        emulation_mode = true;
        dev_info(gsp->dev, "GSP lacks vGPU RPC handlers (ret=%d) - ENABLING emulation mode\n", ret);
    } else if (ret == 0) {
        /* GSP responded - has vGPU handlers */
        emulation_mode = false;
        dev_info(gsp->dev, "GSP has vGPU RPC handlers - DISABLING emulation mode\n");
    } else {
        /* Other error - default to emulation for safety */
        emulation_mode = true;
        dev_info(gsp->dev, "GSP RPC test returned %d - ENABLING emulation mode (safe default)\n", ret);
    }
}

/* GSP Shutdown - just clean up our side */
void gsp_shutdown(struct gsp_device *gsp)
{
    if (gsp->fw_loaded) {
        /* Don't reset GSP - host driver owns it */
        gsp->fw_loaded = false;
    }
}

/* IRQ Setup */
int gsp_setup_irq(struct gsp_device *gsp)
{
    int ret;
    struct msix_entry entries[1];
    
    /* Try MSI-X first */
    ret = pci_enable_msix_range(gsp->pdev, entries, 1, 1);
    if (ret > 0) {
        gsp->msix_enabled = true;
        gsp->irq = entries[0].vector;
        gsp->msix_entries[0] = entries[0];
    } else {
        gsp->irq = gsp->pdev->irq;
    }
    
    ret = devm_request_irq(gsp->dev, gsp->irq, gsp_isr,
                           gsp->msix_enabled ? 0 : IRQF_SHARED,
                           DRIVER_NAME, gsp);
    if (ret) {
        dev_err(gsp->dev, "Failed to request IRQ %d: %d\n", gsp->irq, ret);
        return ret;
    }
    
    dev_info(gsp->dev, "IRQ %d registered (MSI-X: %s)\n", 
             gsp->irq, gsp->msix_enabled ? "yes" : "no");
    return 0;
}

/* PCI Driver Structure */
static struct pci_driver gsp_pci_driver = {
    .name = DRIVER_NAME,
    .id_table = gsp_pci_ids,
    .probe = gsp_probe,
    .remove = gsp_remove,
};

/* Module Init/Exit */
static int __init gsp_module_init(void)
{
    return pci_register_driver(&gsp_pci_driver);
}

static void __exit gsp_module_exit(void)
{
    pci_unregister_driver(&gsp_pci_driver);
}

module_init(gsp_module_init);
module_exit(gsp_module_exit);

MODULE_LICENSE("GPL v2");
MODULE_AUTHOR("NVIDIA vGPU-IO Team");
MODULE_DESCRIPTION("NVIDIA GSP 570.144W vGPU-IO Driver");
MODULE_VERSION(DRIVER_VERSION);
MODULE_FIRMWARE("gsp-570.144.fwimage.bin");