/*
 * Dynamic RPC Dispatcher - Version-aware RPC dispatch
 * Uses discovered firmware addresses for native handlers,
 * falls back to emulation for missing handlers
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/list.h>
#include <linux/errno.h>

#include "dynamic_rpc_dispatcher.h"
#include "gsp_rpc_client.h"
#include "firmware_analyzer.h"

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Open-IOV");
MODULE_DESCRIPTION("Dynamic RPC Dispatcher");

/* Dispatcher context */
struct dynamic_dispatcher {
    struct rpc_interface *current_interface;
    struct fw_analysis_result fw_analysis;
    spinlock_t lock;
    bool initialized;
};

static struct dynamic_dispatcher g_dispatcher = {
    .lock = __SPIN_LOCK_UNLOCKED(g_dispatcher.lock),
    .initialized = false,
};

/* Version 570 interface (native firmware handlers) */
static int rpc_570_alloc(struct gsp_client *client, u64 size, u64 alignment, u32 mem_type,
                         u32 flags, u64 addr_hint, u32 pid, u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    return gsp_rpc_alloc(client, size, alignment, mem_type, flags, addr_hint, pid, gpu_va, handle, actual_size);
}

static int rpc_570_free(struct gsp_client *client, u64 handle)
{
    return gsp_rpc_free(client, handle);
}

static int rpc_570_map(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size,
                       u32 prot, u32 flags, u64 offset, u64 *mapped_va)
{
    return gsp_rpc_map(client, handle, gpu_va, size, prot, flags, offset, mapped_va);
}

static int rpc_570_unmap(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size)
{
    return gsp_rpc_unmap(client, handle, gpu_va, size);
}

static int rpc_570_pin(struct gsp_client *client, u64 handle, u64 offset, u64 size,
                       u32 flags, u64 *dma_addr, u32 *num_pages)
{
    return gsp_rpc_pin(client, handle, offset, size, flags, dma_addr, num_pages);
}

static int rpc_570_unpin(struct gsp_client *client, u64 handle, u64 offset, u64 size)
{
    return gsp_rpc_unpin(client, handle, offset, size);
}

static int rpc_570_schedule(struct gsp_client *client, u64 work_handle, u64 callback, u64 arg,
                            u32 priority, u32 flags)
{
    return gsp_rpc_schedule(client, work_handle, callback, arg, priority, flags);
}

static int rpc_570_yield(struct gsp_client *client)
{
    return gsp_rpc_yield(client);
}

static int rpc_570_bar_trap(struct gsp_client *client, u64 bar_addr, u32 access_type, u32 size,
                            u64 value, u32 bar_num, u64 *ret_value)
{
    return gsp_rpc_bar_trap(client, bar_addr, access_type, size, value, bar_num, ret_value);
}

static int rpc_570_migrate_save(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                u64 memory_bitmap, u64 page_table_ptr, u64 *checkpoint)
{
    return gsp_rpc_migrate_save(client, src_gpu_id, dst_gpu_id, memory_bitmap, page_table_ptr, checkpoint);
}

static int rpc_570_migrate_restore(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                   u64 memory_bitmap, u64 page_table_ptr, u64 checkpoint)
{
    return gsp_rpc_migrate_restore(client, src_gpu_id, dst_gpu_id, memory_bitmap, page_table_ptr, checkpoint);
}

static int rpc_570_notify_register(struct gsp_client *client, u32 type,
                                   void (*callback)(struct gsp_client *, struct gsp_rpc_response *, void *),
                                   void *ctx)
{
    return gsp_rpc_notify_register(client, type, callback, ctx);
}

static void rpc_570_notify_unregister(struct gsp_client *client, u32 type)
{
    gsp_rpc_notify_unregister(client, type);
}

static struct rpc_interface interface_570 = {
    .version = "570.144W",
    .alloc = rpc_570_alloc,
    .free = rpc_570_free,
    .map = rpc_570_map,
    .unmap = rpc_570_unmap,
    .pin = rpc_570_pin,
    .unpin = rpc_570_unpin,
    .schedule = rpc_570_schedule,
    .yield = rpc_570_yield,
    .bar_trap = rpc_570_bar_trap,
    .migrate_save = rpc_570_migrate_save,
    .migrate_restore = rpc_570_migrate_restore,
    .notify_register = rpc_570_notify_register,
    .notify_unregister = rpc_570_notify_unregister,
    .use_emulation = {0},
};

/* Emulation interface (fallback) */
static int rpc_emu_alloc(struct gsp_client *client, u64 size, u64 alignment, u32 mem_type,
                         u32 flags, u64 addr_hint, u32 pid, u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    return gsp_rpc_alloc_emul(client, size, alignment, mem_type, flags, addr_hint, pid, gpu_va, handle, actual_size);
}

static int rpc_emu_free(struct gsp_client *client, u64 handle)
{
    return gsp_rpc_free_emul(client, handle);
}

static int rpc_emu_map(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size,
                       u32 prot, u32 flags, u64 offset, u64 *mapped_va)
{
    return gsp_rpc_map_emul(client, handle, gpu_va, size, prot, flags, offset, mapped_va);
}

static int rpc_emu_unmap(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size)
{
    return gsp_rpc_unmap_emul(client, handle, gpu_va, size);
}

static int rpc_emu_pin(struct gsp_client *client, u64 handle, u64 offset, u64 size,
                       u32 flags, u64 *dma_addr, u32 *num_pages)
{
    return gsp_rpc_pin_emul(client, handle, offset, size, flags, dma_addr, num_pages);
}

static int rpc_emu_unpin(struct gsp_client *client, u64 handle, u64 offset, u64 size)
{
    return gsp_rpc_unpin_emul(client, handle, offset, size);
}

static int rpc_emu_schedule(struct gsp_client *client, u64 work_handle, u64 callback, u64 arg,
                            u32 priority, u32 flags)
{
    return 0; /* Emulation always succeeds */
}

static int rpc_emu_yield(struct gsp_client *client)
{
    return 0;
}

static int rpc_emu_bar_trap(struct gsp_client *client, u64 bar_addr, u32 access_type, u32 size,
                            u64 value, u32 bar_num, u64 *ret_value)
{
    return gsp_rpc_bar_trap(client, bar_addr, access_type, size, value, bar_num, ret_value);
}

static int rpc_emu_migrate_save(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                u64 memory_bitmap, u64 page_table_ptr, u64 *checkpoint)
{
    return 0;
}

static int rpc_emu_migrate_restore(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                   u64 memory_bitmap, u64 page_table_ptr, u64 checkpoint)
{
    return 0;
}

static int rpc_emu_notify_register(struct gsp_client *client, u32 type,
                                   void (*callback)(struct gsp_client *, struct gsp_rpc_response *, void *),
                                   void *ctx)
{
    return gsp_rpc_notify_register(client, type, callback, ctx);
}

static void rpc_emu_notify_unregister(struct gsp_client *client, u32 type)
{
    gsp_rpc_notify_unregister(client, type);
}

static struct rpc_interface interface_emu = {
    .version = "emulation",
    .alloc = rpc_emu_alloc,
    .free = rpc_emu_free,
    .map = rpc_emu_map,
    .unmap = rpc_emu_unmap,
    .pin = rpc_emu_pin,
    .unpin = rpc_emu_unpin,
    .schedule = rpc_emu_schedule,
    .yield = rpc_emu_yield,
    .bar_trap = rpc_emu_bar_trap,
    .migrate_save = rpc_emu_migrate_save,
    .migrate_restore = rpc_emu_migrate_restore,
    .notify_register = rpc_emu_notify_register,
    .notify_unregister = rpc_emu_notify_unregister,
    .use_emulation = {1},
};

int dynamic_dispatcher_init(const char *fw_path)
{
    struct fw_analysis_result result;
    struct rpc_interface *iface;
    bool use_native = false;
    int ret;

    spin_lock(&g_dispatcher.lock);
    if (g_dispatcher.initialized) {
        spin_unlock(&g_dispatcher.lock);
        return 0;
    }
    spin_unlock(&g_dispatcher.lock);

    /*
     * Firmware analysis must run outside the lock: request_firmware()
     * waits for userspace and sleeps, so holding g_dispatcher.lock
     * (a spinlock) across it is a scheduling-while-atomic bug.
     */
    ret = fw_analyzer_load_and_analyze(fw_path, &result);
    if (ret) {
        pr_warn("vgsp-dyn-dispatch: Firmware analysis failed, using emulation: %d\n", ret);
        iface = &interface_emu;
    } else if (result.matched && strcmp(result.version, "570.144W") == 0) {
        iface = &interface_570;
        use_native = true;
        pr_info("vgsp-dyn-dispatch: Using native 570.144W interface\n");
    } else {
        iface = &interface_emu;
        pr_info("vgsp-dyn-dispatch: Using emulation interface (version: %s)\n", result.version);
    }

    spin_lock(&g_dispatcher.lock);
    if (!g_dispatcher.initialized) {
        if (!ret) {
            g_dispatcher.fw_analysis = result;
            if (use_native)
                interface_570.fw_addrs = result.addresses;
        }
        g_dispatcher.current_interface = iface;
        g_dispatcher.initialized = true;
    }
    spin_unlock(&g_dispatcher.lock);

    if (!ret)
        fw_analyzer_print_result(&result);
    return 0;
}

void dynamic_dispatcher_fini(void)
{
    spin_lock(&g_dispatcher.lock);
    g_dispatcher.initialized = false;
    g_dispatcher.current_interface = NULL;
    spin_unlock(&g_dispatcher.lock);
}

struct rpc_interface *dynamic_dispatcher_get_interface(void)
{
    struct rpc_interface *iface;

    spin_lock(&g_dispatcher.lock);
    iface = g_dispatcher.current_interface;
    spin_unlock(&g_dispatcher.lock);

    return iface;
}

bool dynamic_dispatcher_use_emulation(enum gsp_rpc_type type)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface)
        return true;

    if (type >= GSP_RPC_MAX_TYPE)
        return true;

    return iface->use_emulation[type];
}

int dynamic_dispatcher_alloc(struct gsp_client *client, u64 size, u64 alignment, u32 mem_type,
                             u32 flags, u64 addr_hint, u32 pid, u64 *gpu_va, u64 *handle, u64 *actual_size)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->alloc)
        return -ENODEV;
    return iface->alloc(client, size, alignment, mem_type, flags, addr_hint, pid, gpu_va, handle, actual_size);
}

int dynamic_dispatcher_free(struct gsp_client *client, u64 handle)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->free)
        return -ENODEV;
    return iface->free(client, handle);
}

int dynamic_dispatcher_map(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size,
                           u32 prot, u32 flags, u64 offset, u64 *mapped_va)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->map)
        return -ENODEV;
    return iface->map(client, handle, gpu_va, size, prot, flags, offset, mapped_va);
}

int dynamic_dispatcher_unmap(struct gsp_client *client, u64 handle, u64 gpu_va, u64 size)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->unmap)
        return -ENODEV;
    return iface->unmap(client, handle, gpu_va, size);
}

int dynamic_dispatcher_pin(struct gsp_client *client, u64 handle, u64 offset, u64 size,
                           u32 flags, u64 *dma_addr, u32 *num_pages)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->pin)
        return -ENODEV;
    return iface->pin(client, handle, offset, size, flags, dma_addr, num_pages);
}

int dynamic_dispatcher_unpin(struct gsp_client *client, u64 handle, u64 offset, u64 size)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->unpin)
        return -ENODEV;
    return iface->unpin(client, handle, offset, size);
}

int dynamic_dispatcher_schedule(struct gsp_client *client, u64 work_handle, u64 callback, u64 arg,
                                u32 priority, u32 flags)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->schedule)
        return -ENODEV;
    return iface->schedule(client, work_handle, callback, arg, priority, flags);
}

int dynamic_dispatcher_yield(struct gsp_client *client)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->yield)
        return -ENODEV;
    return iface->yield(client);
}

int dynamic_dispatcher_bar_trap(struct gsp_client *client, u64 bar_addr, u32 access_type, u32 size,
                                u64 value, u32 bar_num, u64 *ret_value)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->bar_trap)
        return -ENODEV;
    return iface->bar_trap(client, bar_addr, access_type, size, value, bar_num, ret_value);
}

int dynamic_dispatcher_migrate_save(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                    u64 memory_bitmap, u64 page_table_ptr, u64 *checkpoint)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->migrate_save)
        return -ENODEV;
    return iface->migrate_save(client, src_gpu_id, dst_gpu_id, memory_bitmap, page_table_ptr, checkpoint);
}

int dynamic_dispatcher_migrate_restore(struct gsp_client *client, u64 src_gpu_id, u64 dst_gpu_id,
                                       u64 memory_bitmap, u64 page_table_ptr, u64 checkpoint)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->migrate_restore)
        return -ENODEV;
    return iface->migrate_restore(client, src_gpu_id, dst_gpu_id, memory_bitmap, page_table_ptr, checkpoint);
}

int dynamic_dispatcher_notify_register(struct gsp_client *client, u32 type,
                                       void (*callback)(struct gsp_client *, struct gsp_rpc_response *, void *),
                                       void *ctx)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (!iface || !iface->notify_register)
        return -ENODEV;
    return iface->notify_register(client, type, callback, ctx);
}

void dynamic_dispatcher_notify_unregister(struct gsp_client *client, u32 type)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    if (iface && iface->notify_unregister)
        iface->notify_unregister(client, type);
}

const char *dynamic_dispatcher_get_version(void)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    return iface ? iface->version : "unknown";
}

const struct fw_addresses *dynamic_dispatcher_get_fw_addresses(void)
{
    struct rpc_interface *iface = dynamic_dispatcher_get_interface();
    return iface ? &iface->fw_addrs : NULL;
}

EXPORT_SYMBOL_GPL(dynamic_dispatcher_init);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_fini);
// EXPORT_SYMBOL_GPL(dynamic_dispatcher_get_interface); // Internal function, not exported
EXPORT_SYMBOL_GPL(dynamic_dispatcher_use_emulation);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_alloc);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_free);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_map);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_unmap);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_pin);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_unpin);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_schedule);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_yield);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_bar_trap);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_migrate_save);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_migrate_restore);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_notify_register);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_notify_unregister);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_get_version);
EXPORT_SYMBOL_GPL(dynamic_dispatcher_get_fw_addresses);

MODULE_DESCRIPTION("Dynamic RPC Dispatcher");
MODULE_LICENSE("GPL");