/*
 * GSP RPC Integration Layer
 * Bridges the firmware-decompiled RPC client with existing gsp_comm.c
 */

#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/list.h>
#include <linux/wait.h>
#include <linux/workqueue.h>
#include <linux/io.h>
#include <linux/pci.h>
#include <linux/interrupt.h>
#include <linux/device.h>

#include "nvidia-vgpu-vfio.h"
#include "gsp_comm.h"
#include "gsp_rpc_client.h"
#include "gsp_rpc_integration.h"
#include "bar_emulation.h"
#include "migration.h"

/* Integration context - wraps the new RPC client */
struct gsp_rpc_integration {
    struct nv_vgpu_device *vgpu;
    struct gsp_client *rpc_client;
    struct nv_vgpu_gsp_comm_ctx *comm_ctx;
    struct mutex lock;
    bool initialized;
};

/* Global integration instance (per vGPU) */
static struct gsp_rpc_integration *g_integration = NULL;

/* ============================================================
 * RPC Message Building Helpers
 * ============================================================ */

static void gsp_build_alloc_msg(struct gsp_alloc_req *req,
                                 u64 size, u64 alignment, u32 mem_type,
                                 u32 flags, u64 addr_hint, u32 pid)
{
    req->size = cpu_to_le64(size);
    req->alignment = cpu_to_le64(alignment);
    req->mem_type = cpu_to_le32(mem_type);
    req->flags = cpu_to_le32(flags);
    req->addr_hint = cpu_to_le64(addr_hint);
    req->pid = cpu_to_le32(pid);
    req->reserved = 0;
}

static void gsp_build_map_msg(struct gsp_map_req *req,
                               u64 handle, u64 gpu_va, u64 size,
                               u32 prot, u32 flags, u64 offset)
{
    req->handle = cpu_to_le64(handle);
    req->gpu_va = cpu_to_le64(gpu_va);
    req->size = cpu_to_le64(size);
    req->prot = cpu_to_le32(prot);
    req->flags = cpu_to_le32(flags);
    req->offset = cpu_to_le64(offset);
    req->reserved[0] = 0;
    req->reserved[1] = 0;
}

static void gsp_build_pin_msg(struct gsp_pin_req *req,
                               u64 handle, u64 offset, u64 size, u32 flags)
{
    req->handle = cpu_to_le64(handle);
    req->offset = cpu_to_le64(offset);
    req->size = cpu_to_le64(size);
    req->flags = cpu_to_le32(flags);
    req->reserved[0] = 0;
    req->reserved[1] = 0;
    req->reserved[2] = 0;
}

static void gsp_build_sched_msg(struct gsp_sched_req *req,
                                 u64 work_handle, u64 callback, u64 arg,
                                 u32 priority, u32 flags)
{
    req->work_handle = cpu_to_le64(work_handle);
    req->callback = cpu_to_le64(callback);
    req->arg = cpu_to_le64(arg);
    req->priority = cpu_to_le32(priority);
    req->flags = cpu_to_le32(flags);
    req->reserved[0] = 0;
    req->reserved[1] = 0;
}

static void gsp_build_bar_trap_msg(struct gsp_bar_trap_req *req,
                                    u64 bar_addr, u32 access_type, u32 size,
                                    u64 value, u32 bar_num)
{
    req->bar_addr = cpu_to_le64(bar_addr);
    req->access_type = cpu_to_le32(access_type);
    req->size = cpu_to_le32(size);
    req->value = cpu_to_le64(value);
    req->bar_num = cpu_to_le32(bar_num);
    req->reserved = 0;
}

static void gsp_build_migrate_msg(struct gsp_migrate_req *req,
                                   u64 src_gpu_id, u64 dst_gpu_id,
                                   u64 memory_bitmap, u64 page_table_ptr,
                                   u32 state, u32 flags, u64 context)
{
    req->src_gpu_id = cpu_to_le64(src_gpu_id);
    req->dst_gpu_id = cpu_to_le64(dst_gpu_id);
    req->memory_bitmap = cpu_to_le64(memory_bitmap);
    req->page_table_ptr = cpu_to_le64(page_table_ptr);
    req->state = cpu_to_le32(state);
    req->flags = cpu_to_le32(flags);
    req->context = cpu_to_le64(context);
    req->reserved = 0;
}

/* ============================================================
 * Integration API
 * ============================================================ */

int nv_vgpu_gsp_rpc_integration_init(struct nv_vgpu_device *vgpu,
                                      struct nv_vgpu_gsp_comm_ctx *comm_ctx)
{
    struct gsp_rpc_integration *integ;
    phys_addr_t fw_base, doorbell;
    int ret;

    if (!vgpu || !comm_ctx)
        return -EINVAL;

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

    integ->vgpu = vgpu;
    integ->comm_ctx = comm_ctx;
    mutex_init(&integ->lock);

    /* Calculate firmware base and doorbell from BAR0 */
    if (vgpu->bars[0].vaddr) {
        fw_base = vgpu->bars[0].gpa;
        doorbell = vgpu->bars[0].gpa + NV_VGPU_GSP_DOORBELL_OFFSET;
    } else {
        /* Default for TU102/TU104 */
        fw_base = GSP_FW_BASE;
        doorbell = GSP_FW_BASE + NV_VGPU_GSP_DOORBELL_OFFSET;
    }

    /* Initialize the new RPC client */
    integ->rpc_client = gsp_client_init(&vgpu->pdev->dev, fw_base, doorbell);
    if (IS_ERR(integ->rpc_client)) {
        ret = PTR_ERR(integ->rpc_client);
        integ->rpc_client = NULL;
        goto err_free;
    }

    integ->initialized = true;
    g_integration = integ;

    nv_vgpu_info(vgpu, "GSP RPC integration initialized (fw_base=0x%llx, doorbell=0x%llx)",
                 (unsigned long long)fw_base, (unsigned long long)doorbell);

    return 0;

err_free:
    kfree(integ);
    return ret;
}

void nv_vgpu_gsp_rpc_integration_cleanup(struct nv_vgpu_device *vgpu)
{
    struct gsp_rpc_integration *integ = g_integration;

    if (!integ)
        return;

    mutex_lock(&integ->lock);
    
    if (integ->rpc_client) {
        gsp_client_fini(integ->rpc_client);
        integ->rpc_client = NULL;
    }

    integ->initialized = false;
    mutex_unlock(&integ->lock);

    kfree(integ);
    g_integration = NULL;

    if (vgpu)
        nv_vgpu_info(vgpu, "GSP RPC integration cleaned up");
}

/* ============================================================
 * Enhanced RPC Operations (using firmware-decompiled addresses)
 * ============================================================ */

int nv_vgpu_gsp_rpc_alloc_fw(struct nv_vgpu_device *vgpu,
                              u64 size, u64 alignment, u32 mem_type,
                              u32 flags, u64 addr_hint, u32 pid,
                              u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_alloc_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_alloc_msg(&req, size, alignment, mem_type, flags, addr_hint, pid);

    ret = gsp_rpc_alloc(integ->rpc_client,
                         size, alignment, mem_type, flags, addr_hint, pid,
                         gpu_va, handle, actual_size);

    if (ret == 0) {
        nv_vgpu_dbg(vgpu, "ALLOC: size=%llu align=%llu -> gpu_va=0x%llx handle=0x%llx",
                    size, alignment, *gpu_va, *handle);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_free_fw(struct nv_vgpu_device *vgpu, u64 handle)
{
    struct gsp_rpc_integration *integ = g_integration;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_rpc_free(integ->rpc_client, handle);
}

int nv_vgpu_gsp_rpc_map_fw(struct nv_vgpu_device *vgpu,
                            u64 handle, u64 gpu_va, u64 size,
                            u32 prot, u32 flags, u64 offset,
                            u64 *mapped_va)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_map_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_map_msg(&req, handle, gpu_va, size, prot, flags, offset);

    ret = gsp_rpc_map(integ->rpc_client, handle, gpu_va, size,
                       prot, flags, offset, mapped_va);

    if (ret == 0) {
        nv_vgpu_dbg(vgpu, "MAP: handle=0x%llx gpu_va=0x%llx size=%llu -> mapped_va=0x%llx",
                    handle, gpu_va, size, *mapped_va);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_unmap_fw(struct nv_vgpu_device *vgpu,
                              u64 handle, u64 gpu_va, u64 size)
{
    struct gsp_rpc_integration *integ = g_integration;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_rpc_unmap(integ->rpc_client, handle, gpu_va, size);
}

int nv_vgpu_gsp_rpc_pin_fw(struct nv_vgpu_device *vgpu,
                            u64 handle, u64 offset, u64 size,
                            u32 flags, u64 *dma_addr, u32 *num_pages)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_pin_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_pin_msg(&req, handle, offset, size, flags);

    ret = gsp_rpc_pin(integ->rpc_client, handle, offset, size,
                       flags, dma_addr, num_pages);

    if (ret == 0) {
        nv_vgpu_dbg(vgpu, "PIN: handle=0x%llx offset=%llu size=%llu -> dma_addr=0x%llx pages=%u",
                    handle, offset, size, *dma_addr, *num_pages);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_unpin_fw(struct nv_vgpu_device *vgpu,
                              u64 handle, u64 offset, u64 size)
{
    struct gsp_rpc_integration *integ = g_integration;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_rpc_unpin(integ->rpc_client, handle, offset, size);
}

int nv_vgpu_gsp_rpc_schedule_fw(struct nv_vgpu_device *vgpu,
                                 u64 work_handle, u64 callback, u64 arg,
                                 u32 priority, u32 flags)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_sched_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_sched_msg(&req, work_handle, callback, arg, priority, flags);

    ret = gsp_rpc_schedule(integ->rpc_client, work_handle, callback, arg,
                            priority, flags);

    if (ret == 0) {
        nv_vgpu_dbg(vgpu, "SCHEDULE: work_handle=0x%llx priority=%u", work_handle, priority);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_yield_fw(struct nv_vgpu_device *vgpu)
{
    struct gsp_rpc_integration *integ = g_integration;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_rpc_yield(integ->rpc_client);
}

int nv_vgpu_gsp_rpc_bar_trap_fw(struct nv_vgpu_device *vgpu,
                                 u64 bar_addr, u32 access_type, u32 size,
                                 u64 value, u32 bar_num, u64 *ret_value)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_bar_trap_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_bar_trap_msg(&req, bar_addr, access_type, size, value, bar_num);

    ret = gsp_rpc_bar_trap(integ->rpc_client, bar_addr, access_type, size,
                            value, bar_num, ret_value);

    if (ret == 0) {
        nv_vgpu_dbg(vgpu, "BAR_TRAP: bar=%u addr=0x%llx type=%u -> value=0x%llx",
                    bar_num, bar_addr, access_type, *ret_value);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_migrate_save_fw(struct nv_vgpu_device *vgpu,
                                     u64 src_gpu_id, u64 dst_gpu_id,
                                     u64 memory_bitmap, u64 page_table_ptr,
                                     u64 *checkpoint)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_migrate_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_migrate_msg(&req, src_gpu_id, dst_gpu_id,
                           memory_bitmap, page_table_ptr,
                           GSP_MIGRATE_SUBCMD_SAVE, 0, 0);

    ret = gsp_rpc_migrate_save(integ->rpc_client,
                                src_gpu_id, dst_gpu_id,
                                memory_bitmap, page_table_ptr,
                                checkpoint);

    if (ret == 0) {
        nv_vgpu_info(vgpu, "MIGRATE SAVE: src=%llu dst=%llu -> checkpoint=0x%llx",
                     src_gpu_id, dst_gpu_id, *checkpoint);
    }

    return ret;
}

int nv_vgpu_gsp_rpc_migrate_restore_fw(struct nv_vgpu_device *vgpu,
                                        u64 src_gpu_id, u64 dst_gpu_id,
                                        u64 memory_bitmap, u64 page_table_ptr,
                                        u64 checkpoint)
{
    struct gsp_rpc_integration *integ = g_integration;
    struct gsp_migrate_req req;
    int ret;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    gsp_build_migrate_msg(&req, src_gpu_id, dst_gpu_id,
                           memory_bitmap, page_table_ptr,
                           GSP_MIGRATE_SUBCMD_RESTORE, 0, checkpoint);

    ret = gsp_rpc_migrate_restore(integ->rpc_client,
                                   src_gpu_id, dst_gpu_id,
                                   memory_bitmap, page_table_ptr,
                                   checkpoint);

    if (ret == 0) {
        nv_vgpu_info(vgpu, "MIGRATE RESTORE: src=%llu dst=%llu checkpoint=0x%llx",
                     src_gpu_id, dst_gpu_id, checkpoint);
    }

    return ret;
}

/* ============================================================
 * WPR Management (from firmware strings)
 * ============================================================ */

int nv_vgpu_gsp_wpr_write_mutex_acquire(struct nv_vgpu_device *vgpu, u32 wpr_id)
{
    struct gsp_rpc_integration *integ = g_integration;

    (void)vgpu; /* integ is global, vgpu not directly used */

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_wpr_write_mutex_acquire(integ->rpc_client, wpr_id);
}

int nv_vgpu_gsp_wpr_write_mutex_release(struct nv_vgpu_device *vgpu, u32 wpr_id)
{
    struct gsp_rpc_integration *integ = g_integration;

    (void)vgpu;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_wpr_write_mutex_release(integ->rpc_client, wpr_id);
}

int nv_vgpu_gsp_wpr_scratch_mutex_acquire(struct nv_vgpu_device *vgpu, u32 wpr_id)
{
    struct gsp_rpc_integration *integ = g_integration;

    (void)vgpu;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_wpr_scratch_mutex_acquire(integ->rpc_client, wpr_id);
}

int nv_vgpu_gsp_wpr_scratch_mutex_release(struct nv_vgpu_device *vgpu, u32 wpr_id)
{
    struct gsp_rpc_integration *integ = g_integration;

    (void)vgpu;

    if (!integ || !integ->initialized || !integ->rpc_client)
        return -ENODEV;

    return gsp_wpr_scratch_mutex_release(integ->rpc_client, wpr_id);
}

/* ============================================================
 * Firmware Address Accessors (for debugging/validation)
 * ============================================================ */

u64 nv_vgpu_gsp_get_primary_dispatcher(void)
{
    return GSP_RPC_DISPATCH_PRIMARY;
}

u64 nv_vgpu_gsp_get_secondary_dispatcher(void)
{
    return GSP_RPC_DISPATCH_SECONDARY;
}

u64 nv_vgpu_gsp_get_alloc_main(void)
{
    return GSP_ALLOC_MAIN;
}

u64 nv_vgpu_gsp_get_schedule_main(void)
{
    return GSP_SCHEDULE_MAIN;
}

u64 nv_vgpu_gsp_get_bar_trap_handler_jt(void)
{
    return GSP_BAR_TRAP_HANDLER_JT;
}

u64 nv_vgpu_gsp_get_bar_trap_handler_primary(void)
{
    return GSP_BAR_TRAP_HANDLER_PRIM;
}

u64 nv_vgpu_gsp_get_entry_point(void)
{
    return GSP_ENTRY_POINT;
}

EXPORT_SYMBOL(nv_vgpu_gsp_rpc_integration_init);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_integration_cleanup);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_alloc_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_free_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_map_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_unmap_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_pin_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_unpin_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_schedule_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_yield_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_bar_trap_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_migrate_save_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_rpc_migrate_restore_fw);
EXPORT_SYMBOL(nv_vgpu_gsp_wpr_write_mutex_acquire);
EXPORT_SYMBOL(nv_vgpu_gsp_wpr_write_mutex_release);
EXPORT_SYMBOL(nv_vgpu_gsp_wpr_scratch_mutex_acquire);
EXPORT_SYMBOL(nv_vgpu_gsp_wpr_scratch_mutex_release);
EXPORT_SYMBOL(nv_vgpu_gsp_get_primary_dispatcher);
EXPORT_SYMBOL(nv_vgpu_gsp_get_secondary_dispatcher);
EXPORT_SYMBOL(nv_vgpu_gsp_get_alloc_main);
EXPORT_SYMBOL(nv_vgpu_gsp_get_schedule_main);
EXPORT_SYMBOL(nv_vgpu_gsp_get_bar_trap_handler_jt);
EXPORT_SYMBOL(nv_vgpu_gsp_get_bar_trap_handler_primary);
EXPORT_SYMBOL(nv_vgpu_gsp_get_entry_point);