/*
 * GSP Emulation Engine Implementation
 * Host-side GSP emulation thread that processes RPCs from guest via shared memory
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/mm.h>
#include <linux/dma-mapping.h>
#include <linux/mutex.h>
#include <linux/spinlock.h>
#include <linux/list.h>
#include <linux/workqueue.h>
#include <linux/delay.h>
#include <linux/kthread.h>
#include <linux/sched.h>
#include <linux/ktime.h>
#include <linux/eventfd.h>
#include <asm/io.h>

#include "nvidia-vgpu-vfio.h"
#include "gsp_emulation.h"
#include "gsp_rpc_client.h"
#include "bar_emulation.h"
#include "irq.h"
#include "memory.h"

/* ============================================================
 * Constants
 * ============================================================ */

#define GSP_EMU_DOORBELL_OFFSET     0x110C00
#define GSP_EMU_SHARED_MEM_OFFSET   0x1000
#define GSP_EMU_SHARED_MEM_SIZE     4096
#define GSP_EMU_MSIX_TABLE_OFFSET   0x2000
#define GSP_EMU_MSIX_PBA_OFFSET     0x3000
#define GSP_EMU_POLL_BUSY_US        100
#define GSP_EMU_POLL_SLEEP_MIN_US   50
#define GSP_EMU_POLL_SLEEP_MAX_US   100

/* ============================================================
 * Helper Functions
 * ============================================================ */

static inline u32 gsp_emu_read_doorbell(struct gsp_emu_context *ctx)
{
    return readl(ctx->doorbell_reg);
}

static inline void gsp_emu_write_doorbell(struct gsp_emu_context *ctx, u32 val)
{
    writel(val, ctx->doorbell_reg);
}

static inline void gsp_emu_copy_request_from_shared(struct gsp_emu_context *ctx,
                                                     struct gsp_rpc_msg *req)
{
    memcpy(req, ctx->shared_mem_virt, sizeof(*req));
}

static inline void gsp_emu_copy_response_to_shared(struct gsp_emu_context *ctx,
                                                    struct gsp_rpc_response *rsp)
{
    memcpy(ctx->shared_mem_virt + 0x80, rsp, sizeof(*rsp));
}

/* ============================================================
 * RPC Processing
 * ============================================================ */

static int gsp_emu_dispatch_rpc(struct gsp_emu_context *ctx,
                                 struct gsp_rpc_msg *req,
                                 struct gsp_rpc_response *rsp)
{
    int ret = 0;
    ktime_t start = ktime_get();

    memset(rsp, 0, sizeof(*rsp));
    rsp->type = req->type;
    rsp->seq_id = req->seq_id;

    switch (req->type) {
    case GSP_RPC_ALLOC: {
        struct gsp_alloc_req *areq = (struct gsp_alloc_req *)req->payload;
        struct gsp_alloc_rsp arsp = {0};
        u64 gpu_va = 0, handle = 0, actual_size = 0;

        ret = gsp_rpc_alloc(ctx->rpc_client,
                            areq->size, areq->alignment, areq->mem_type,
                            areq->flags, areq->addr_hint, areq->pid,
                            &gpu_va, &handle, &actual_size);
        if (!ret) {
            arsp.gpu_va = gpu_va;
            arsp.handle = handle;
            arsp.size = actual_size;
            memcpy(rsp->result, &arsp, sizeof(arsp));
        }
        break;
    }

    case GSP_RPC_FREE: {
        u64 handle = req->payload[0];
        ret = gsp_rpc_free(ctx->rpc_client, handle);
        break;
    }

    case GSP_RPC_MAP: {
        struct gsp_map_req *mreq = (struct gsp_map_req *)req->payload;
        struct gsp_map_rsp mrsp = {0};
        u64 mapped_va = 0;

        ret = gsp_rpc_map(ctx->rpc_client,
                          mreq->handle, mreq->gpu_va, mreq->size,
                          mreq->prot, mreq->flags, mreq->offset,
                          &mapped_va);
        if (!ret) {
            mrsp.gpu_va = mapped_va;
            memcpy(rsp->result, &mrsp, sizeof(mrsp));
        }
        break;
    }

    case GSP_RPC_UNMAP: {
        u64 handle = req->payload[0];
        u64 gpu_va = req->payload[1];
        u64 size = req->payload[2];
        ret = gsp_rpc_unmap(ctx->rpc_client, handle, gpu_va, size);
        break;
    }

    case GSP_RPC_PIN: {
        struct gsp_pin_req *preq = (struct gsp_pin_req *)req->payload;
        struct gsp_pin_rsp prsp = {0};
        u64 dma_addr = 0;
        u32 num_pages = 0;

        ret = gsp_rpc_pin(ctx->rpc_client,
                          preq->handle, preq->offset, preq->size,
                          preq->flags, &dma_addr, &num_pages);
        if (!ret) {
            prsp.dma_addr = dma_addr;
            prsp.num_pages = num_pages;
            memcpy(rsp->result, &prsp, sizeof(prsp));
        }
        break;
    }

    case GSP_RPC_UNPIN: {
        u64 handle = req->payload[0];
        u64 offset = req->payload[1];
        u64 size = req->payload[2];
        ret = gsp_rpc_unpin(ctx->rpc_client, handle, offset, size);
        break;
    }

    case GSP_RPC_SCHEDULE: {
        struct gsp_sched_req *sreq = (struct gsp_sched_req *)req->payload;
        ret = gsp_rpc_schedule(ctx->rpc_client,
                               sreq->work_handle, sreq->callback, sreq->arg,
                               sreq->priority, sreq->flags);
        break;
    }

    case GSP_RPC_YIELD: {
        ret = gsp_rpc_yield(ctx->rpc_client);
        break;
    }

    case GSP_RPC_NOTIFY: {
        /* NOTIFY is handled via registration, not direct RPC */
        nv_vgpu_info(ctx->vgpu, "NOTIFY RPC received (type=0x%x)", req->subcmd);
        ret = 0;
        break;
    }

    case GSP_RPC_BAR_TRAP: {
        struct gsp_bar_trap_req *breq = (struct gsp_bar_trap_req *)req->payload;
        struct gsp_bar_trap_rsp brsp = {0};
        u64 ret_value = 0;

        ret = gsp_rpc_bar_trap(ctx->rpc_client,
                               breq->bar_addr, breq->access_type, breq->size,
                               breq->value, breq->bar_num, &ret_value);
        if (!ret) {
            brsp.value = ret_value;
            brsp.handled = 1;
            memcpy(rsp->result, &brsp, sizeof(brsp));
        }
        break;
    }

    case GSP_RPC_MIGRATE: {
        struct gsp_migrate_req *mreq = (struct gsp_migrate_req *)req->payload;
        struct gsp_migrate_rsp mrsp = {0};
        u64 checkpoint = 0;

        if (mreq->flags == GSP_MIGRATE_SUBCMD_SAVE) {
            ret = gsp_rpc_migrate_save(ctx->rpc_client,
                                       mreq->src_gpu_id, mreq->dst_gpu_id,
                                       mreq->memory_bitmap, mreq->page_table_ptr,
                                       &checkpoint);
            if (!ret) {
                mrsp.checkpoint = checkpoint;
                memcpy(rsp->result, &mrsp, sizeof(mrsp));
            }
        } else if (mreq->flags == GSP_MIGRATE_SUBCMD_RESTORE) {
            ret = gsp_rpc_migrate_restore(ctx->rpc_client,
                                          mreq->src_gpu_id, mreq->dst_gpu_id,
                                          mreq->memory_bitmap, mreq->page_table_ptr,
                                          mreq->context);
        }
        break;
    }

    case GSP_RPC_DEBUG: {
        nv_vgpu_info(ctx->vgpu, "DEBUG RPC received (seq=%u)", req->seq_id);
        ret = 0;
        break;
    }

    default:
        nv_vgpu_warn(ctx->vgpu, "Unknown RPC type: 0x%x", req->type);
        rsp->status = -ENOTSUPP;
        ret = -ENOTSUPP;
        break;
    }

    if (ret)
        rsp->status = ret;

    ctx->total_latency_ns += ktime_to_ns(ktime_sub(ktime_get(), start));
    return ret;
}

/* ============================================================
 * MSI-X Interrupt Injection
 * ============================================================ */

void gsp_emu_inject_completion_interrupt(struct gsp_emu_context *ctx, u32 seq_id)
{
    struct eventfd_ctx *trigger = NULL;

    /* Write to MSI-X PBA (Pending Bit Array) at BAR0+0x3000 */
    if (ctx->bar0_vaddr) {
        u32 *pba = (u32 *)(ctx->bar0_vaddr + GSP_EMU_MSIX_PBA_OFFSET);
        writel(1, pba); /* Set pending bit for vector 0 */
    }

    /* Look up the vector-0 trigger LIVE: QEMU wires eventfds via
     * VFIO_DEVICE_SET_IRQS after gsp_emu_init runs, so an init-time copy
     * would be stale (NULL). */
    if (ctx->vgpu)
        trigger = ctx->vgpu->irqs[0].trigger;

    /* Signal eventfd for VFIO interrupt delivery */
    if (trigger)
        eventfd_signal(trigger);
}

/* ============================================================
 * Emulation Thread
 * ============================================================ */

static int gsp_emu_thread_fn(void *data)
{
    struct gsp_emu_context *ctx = data;
    struct gsp_rpc_msg req;
    struct gsp_rpc_response rsp;
    u32 doorbell_val;
    int ret;

    nv_vgpu_info(ctx->vgpu, "GSP emulation thread started");

    while (!kthread_should_stop()) {
        /* Read doorbell */
        doorbell_val = gsp_emu_read_doorbell(ctx);

        if (doorbell_val != 0) {
            /* ACK doorbell immediately */
            gsp_emu_write_doorbell(ctx, 0);

            /* Copy request from shared memory */
            gsp_emu_copy_request_from_shared(ctx, &req);

            /* Validate sequence ID */
            if (req.seq_id == 0) {
                nv_vgpu_warn(ctx->vgpu, "Invalid seq_id=0, ignoring RPC");
                continue;
            }

            /* Process RPC */
            ret = gsp_emu_dispatch_rpc(ctx, &req, &rsp);

            /* Write response to shared memory */
            gsp_emu_copy_response_to_shared(ctx, &rsp);

            /* Inject completion interrupt */
            gsp_emu_inject_completion_interrupt(ctx, req.seq_id);

            if (ret)
                ctx->failed_rpcs++;
            else
                ctx->processed_rpcs++;
        } else {
            /* Adaptive polling: busy-wait first, then sleep (interruptible for kthread_stop) */
            static int poll_count = 0;
            if (poll_count < GSP_EMU_POLL_BUSY_US) {
                udelay(1);
                poll_count++;
            } else {
                /* Use interruptible sleep so kthread_stop can wake us */
                schedule_timeout_interruptible(usecs_to_jiffies(GSP_EMU_POLL_SLEEP_MAX_US));
                poll_count = 0;
            }
        }
    }

    nv_vgpu_info(ctx->vgpu, "GSP emulation thread stopped (processed=%llu, failed=%llu)",
                 ctx->processed_rpcs, ctx->failed_rpcs);
    complete(&ctx->thread_stopped);
    return 0;
}

/* ============================================================
 * Initialization & Cleanup
 * ============================================================ */

int gsp_emu_init(struct nv_vgpu_device *vgpu)
{
    struct nv_vgpu_bar_emul_ctx *bar_ctx;
    struct gsp_emu_context *ctx;
    int ret = 0;

    nv_vgpu_info(vgpu, "Initializing GSP emulation engine");

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

    ctx->vgpu = vgpu;
    ctx->running = false;
    ctx->initialized = false;
    ctx->stopping = false;
    init_completion(&ctx->thread_stopped);
    spin_lock_init(&ctx->lock);

    /* The guest's BAR0 is the emulation shadow (vzalloc'd by
     * nv_vgpu_bar_emul_init, mapped to the guest via VFIO mmap) - NOT the
     * host physical BAR0 (vgpu->bars[0].vaddr, which belongs to the real GPU
     * and the host nvidia driver). All guest doorbell writes, RPC messages,
     * MSI-X table writes and PBA reads land in the shadow, so the emulation
     * engine must operate on the shadow exclusively. Polling the physical
     * BAR0 would never see guest traffic and would corrupt host GPU state. */
    bar_ctx = (struct nv_vgpu_bar_emul_ctx *)vgpu->bar_emul_ctx;
    if (!bar_ctx || !bar_ctx->bar0_shadow) {
        nv_vgpu_err(vgpu, "BAR0 emulation shadow not available");
        ret = -ENODEV;
        goto err_free_ctx;
    }

    ctx->bar0_vaddr = (void __iomem *)bar_ctx->bar0_shadow;

    /* Doorbell register at BAR0+0x110C00 (guest-visible shadow) */
    ctx->doorbell_reg = ctx->bar0_vaddr + GSP_EMU_DOORBELL_OFFSET;

    /* RPC shared memory at BAR0+0x1000. This ALIASES the guest-visible
     * shadow - it is NOT a separate allocation. shared_mem_dma stays 0 to
     * mark it as aliased so cleanup does not dma_free it. */
    ctx->shared_mem_size = GSP_EMU_SHARED_MEM_SIZE;
    ctx->shared_mem = ctx->bar0_vaddr + GSP_EMU_SHARED_MEM_OFFSET;
    ctx->shared_mem_virt = (void *)ctx->shared_mem;
    ctx->shared_mem_dma = 0;

    /* Initialize RPC client with the shadow BAR0 base and doorbell */
    ctx->rpc_client = gsp_client_init_emulation(&vgpu->dev, ctx->bar0_vaddr, ctx->doorbell_reg);
    if (IS_ERR(ctx->rpc_client)) {
        ret = PTR_ERR(ctx->rpc_client);
        nv_vgpu_err(vgpu, "Failed to init GSP RPC client: %d", ret);
        goto err_free_ctx;
    }

    /* Enable emulation mode in RPC client */
    gsp_rpc_emulation_set(true);

    /* NOTE: MSI-X eventfd triggers are NOT copied here - QEMU wires them
     * via VFIO_DEVICE_SET_IRQS later, so gsp_emu_inject_completion_interrupt
     * looks them up live from vgpu->irqs[]. */

    /* Start emulation thread */
    ctx->thread = kthread_run(gsp_emu_thread_fn, ctx, "gsp_emu_%u", vgpu->vgpu_id);
    if (IS_ERR(ctx->thread)) {
        ret = PTR_ERR(ctx->thread);
        nv_vgpu_err(vgpu, "Failed to start GSP emulation thread: %d", ret);
        goto err_fini_client;
    }

    ctx->running = true;
    ctx->initialized = true;
    vgpu->gsp_comm_ctx = ctx; /* Reuse existing field for emulation context */

    nv_vgpu_info(vgpu, "GSP emulation engine initialized successfully");
    return 0;

err_fini_client:
    gsp_client_fini(ctx->rpc_client);
err_free_ctx:
    /* shared_mem_virt aliases the BAR0 shadow (owned by bar_emulation),
     * nothing to free here */
    kfree(ctx);
    return ret;
}

void gsp_emu_cleanup(struct nv_vgpu_device *vgpu)
{
    struct gsp_emu_context *ctx = vgpu->gsp_comm_ctx;

    if (!ctx)
        return;

    nv_vgpu_info(vgpu, "Cleaning up GSP emulation engine");

    /* Stop emulation thread with timeout */
    if (ctx->thread) {
        ctx->stopping = true;
        kthread_stop(ctx->thread);
        /* Wait for thread to actually stop with timeout (5 seconds) */
        if (!wait_for_completion_timeout(&ctx->thread_stopped, msecs_to_jiffies(5000))) {
            nv_vgpu_warn(vgpu, "GSP emulation thread did not stop in time, forcing cleanup");
        }
        ctx->thread = NULL;
    }
    ctx->running = false;

    /* Disable emulation mode */
    gsp_rpc_emulation_set(false);

    /* Cleanup RPC client */
    if (ctx->rpc_client) {
        gsp_client_fini(ctx->rpc_client);
        ctx->rpc_client = NULL;
    }

    /* Shared memory aliases the BAR0 shadow (owned by bar_emulation) -
     * only free it if it was a standalone DMA allocation (dma != 0) */
    if (ctx->shared_mem_virt && ctx->shared_mem_dma) {
        dma_free_wc(&vgpu->pdev->dev, ctx->shared_mem_size,
                    ctx->shared_mem_virt, ctx->shared_mem_dma);
    }
    ctx->shared_mem_virt = NULL;

    /* BAR0 shadow is managed by nv_vgpu_bar_emul_init/cleanup */
    ctx->bar0_vaddr = NULL;

    vgpu->gsp_comm_ctx = NULL;
    kfree(ctx);

    nv_vgpu_info(vgpu, "GSP emulation engine cleaned up");
}

/* ============================================================
 * Public RPC Processing (for external callers)
 * ============================================================ */

int gsp_emu_process_rpc(struct gsp_emu_context *ctx)
{
    struct gsp_rpc_msg req;
    struct gsp_rpc_response rsp;
    int ret;

    if (!ctx || !ctx->initialized)
        return -ENODEV;

    gsp_emu_copy_request_from_shared(ctx, &req);

    if (req.seq_id == 0)
        return -EINVAL;

    ret = gsp_emu_dispatch_rpc(ctx, &req, &rsp);

    gsp_emu_copy_response_to_shared(ctx, &rsp);
    gsp_emu_inject_completion_interrupt(ctx, req.seq_id);

    return ret;
}

/* ============================================================
 * Statistics
 * ============================================================ */

void gsp_emu_get_stats(struct gsp_emu_context *ctx,
                       u64 *processed, u64 *failed, u64 *total_latency_ns)
{
    if (!ctx)
        return;

    if (processed)
        *processed = ctx->processed_rpcs;
    if (failed)
        *failed = ctx->failed_rpcs;
    if (total_latency_ns)
        *total_latency_ns = ctx->total_latency_ns;
}

EXPORT_SYMBOL(gsp_emu_process_rpc);
EXPORT_SYMBOL(gsp_emu_get_stats);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Open-IOV Project");
MODULE_DESCRIPTION("GSP Emulation Engine for Host-Side vGPU");