/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * GSP vGPU-IO Driver - VF Management
 */

#include "gsp_core.h"
#include <linux/slab.h>
#include <linux/idr.h>
#include <linux/list.h>

/* VF Reference Counting */
static void gsp_vf_release(struct kref *ref)
{
    struct gsp_vf *vf = container_of(ref, struct gsp_vf, ref);
    kfree(vf);
}

struct gsp_vf *gsp_vf_get(struct gsp_device *gsp, uint32_t id)
{
    struct gsp_vf *vf;
    
    mutex_lock(&gsp->vf_lock);
    vf = idr_find(&gsp->vf_idr, id);
    if (vf && vf->active)
        kref_get(&vf->ref);
    else
        vf = NULL;
    mutex_unlock(&gsp->vf_lock);
    
    return vf;
}

void gsp_vf_put(struct gsp_vf *vf)
{
    if (vf)
        kref_put(&vf->ref, gsp_vf_release);
}

/* VF Creation */
int gsp_vf_create(struct gsp_device *gsp, const struct gsp_vf_config *config)
{
    struct gsp_vf *vf;
    int ret, id;
    
    if (!config || config->id == 0 || config->id > gsp->max_vfs)
        return -EINVAL;
    
    vf = kzalloc(sizeof(*vf), GFP_KERNEL);
    if (!vf)
        return -ENOMEM;
    
    vf->gsp = gsp;
    vf->id = config->id;
    kref_init(&vf->ref);
    mutex_init(&vf->lock);
    mutex_init(&vf->wpr_lock);
    INIT_LIST_HEAD(&vf->wpr_list);
    INIT_LIST_HEAD(&vf->list);
    atomic_set(&vf->refcount, 1);
    
    /* Copy configuration */
    vf->bar1_offset = config->bar1_offset;
    vf->bar1_size = config->bar1_size;
    vf->pt_base = config->pt_base;
    vf->pt_size = config->pt_size;
    vf->cmd_base = config->cmd_base;
    vf->cmd_size = config->cmd_size;
    
    /* Allocate VF ID */
    mutex_lock(&gsp->vf_lock);
    id = idr_alloc(&gsp->vf_idr, vf, config->id, config->id + 1, GFP_KERNEL);
    if (id < 0) {
        ret = id;
        mutex_unlock(&gsp->vf_lock);
        goto err_free;
    }
    
    list_add(&vf->list, &gsp->vf_list);
    vf->active = true;
    vf->state = GSP_VF_ACTIVE;
    mutex_unlock(&gsp->vf_lock);
    
    /* Initialize subsystems */
    ret = gsp_vf_wpr_init(vf);
    if (ret) {
        dev_err(gsp->dev, "VF %u: WPR init failed: %d\n", vf->id, ret);
        goto err_wpr;
    }
    
    ret = gsp_vf_bar_init(vf);
    if (ret) {
        dev_err(gsp->dev, "VF %u: BAR init failed: %d\n", vf->id, ret);
        goto err_bar;
    }
    
    dev_info(gsp->dev, "VF %u created (bar1=0x%llx-0x%llx, pt=0x%llx-0x%llx)\n",
             vf->id, vf->bar1_offset, vf->bar1_offset + vf->bar1_size,
             vf->pt_base, vf->pt_base + vf->pt_size);
    
    return 0;

err_bar:
    gsp_vf_wpr_fini(vf);
err_wpr:
    mutex_lock(&gsp->vf_lock);
    list_del(&vf->list);
    idr_remove(&gsp->vf_idr, vf->id);
    vf->active = false;
    mutex_unlock(&gsp->vf_lock);
err_free:
    kfree(vf);
    return ret;
}

/* VF Destruction */
void gsp_vf_destroy(struct gsp_vf *vf)
{
    struct gsp_device *gsp = vf->gsp;
    uint32_t id = vf->id;
    
    if (!vf->active)
        return;
    
    /* Cancel any ongoing migration */
    mutex_lock(&vf->lock);
    if (vf->migration) {
        gsp_migration_cancel(vf->migration);
        vf->migration = NULL;
    }
    vf->active = false;
    vf->state = GSP_VF_DESTROYING;
    mutex_unlock(&vf->lock);
    
    /* Remove from device lists */
    mutex_lock(&gsp->vf_lock);
    list_del(&vf->list);
    idr_remove(&gsp->vf_idr, id);
    mutex_unlock(&gsp->vf_lock);
    
    /* Cleanup subsystems */
    gsp_vf_bar_fini(vf);
    gsp_vf_wpr_fini(vf);
    
    dev_info(gsp->dev, "VF %u destroyed\n", id);
    
    gsp_vf_put(vf);
}

/* Memory Allocation */
int gsp_vf_mem_alloc(struct gsp_vf *vf, uint64_t size, uint32_t flags,
                     uint64_t *handle, uint64_t *va)
{
    struct gsp_alloc_payload payload = {
        .size = size,
        .type = flags & 0x3,  /* Low bits: mem type */
        .alignment = 0x1000,
        .va_hint = 0,
        .flags = flags >> 2,
        .pid = current->pid,
        .vf_id = vf->id,
    };
    struct gsp_alloc_response resp;
    size_t resp_size = sizeof(resp);
    int ret;
    
    if (!handle || !va)
        return -EINVAL;
    
    /* In emulation mode, handle locally */
    if (emulation_mode) {
        static uint64_t next_handle = 0x1000;
        static uint64_t next_va = 0x100000000ULL;  /* 4GB base */
        
        *handle = next_handle++;
        *va = next_va;
        next_va += (size + 0xFFF) & ~0xFFF;  /* Align to 4KB */
        
        dev_dbg(vf->gsp->dev, "VF %u: alloc handle=0x%llx va=0x%llx size=0x%llx (emulation)\n",
                vf->id, *handle, *va, size);
        
        return 0;
    }
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_ALLOC,
                            &payload, sizeof(payload),
                            &resp, &resp_size);
    if (ret)
        return ret;
    
    *handle = resp.handle;
    *va = resp.va;
    
    dev_dbg(vf->gsp->dev, "VF %u: alloc handle=0x%llx va=0x%llx size=0x%llx\n",
            vf->id, *handle, *va, size);
    
    return 0;
}

int gsp_vf_mem_free(struct gsp_vf *vf, uint64_t handle)
{
    /* FREE uses primary dispatcher sub-command */
    struct gsp_subcmd_desc desc;
    size_t desc_size = sizeof(desc);
    int ret;
    
    /* Subcmd for FREE: need to determine which subcmd corresponds to FREE */
    /* For now, use a generic FREE RPC if implemented */
    
    /* Placeholder - actual implementation depends on FREE subcmd mapping */
    dev_dbg(vf->gsp->dev, "VF %u: free handle=0x%llx\n", vf->id, handle);
    
    return 0;
}

/* Memory Pinning */
int gsp_vf_mem_pin(struct gsp_vf *vf, uint64_t handle,
                   uint64_t offset, uint64_t size,
                   uint64_t *dma_addr)
{
    struct gsp_pin_payload payload = {
        .handle = handle,
        .offset = offset,
        .size = size,
        .vf_id = vf->id,
        .flags = GSP_PIN_READ | GSP_PIN_WRITE,
    };
    struct gsp_pin_response resp;
    size_t resp_size = sizeof(resp);
    int ret;
    
    if (!dma_addr)
        return -EINVAL;
    
    /* In emulation mode, handle locally */
    if (emulation_mode) {
        static uint64_t next_dma = 0x200000000ULL;  /* 8GB base */
        *dma_addr = next_dma;
        next_dma += (size + 0xFFF) & ~0xFFF;
        
        dev_dbg(vf->gsp->dev, "VF %u: pin handle=0x%llx offset=0x%llx size=0x%llx -> dma=0x%llx (emulation)\n",
                vf->id, handle, offset, size, *dma_addr);
        
        return 0;
    }
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_PIN,
                            &payload, sizeof(payload),
                            &resp, &resp_size);
    if (ret)
        return ret;
    
    *dma_addr = resp.dma_addr;
    
    dev_dbg(vf->gsp->dev, "VF %u: pin handle=0x%llx offset=0x%llx size=0x%llx -> dma=0x%llx\n",
            vf->id, handle, offset, size, *dma_addr);
    
    return 0;
}

int gsp_vf_mem_unpin(struct gsp_vf *vf, uint64_t handle)
{
    dev_dbg(vf->gsp->dev, "VF %u: unpin handle=0x%llx\n", vf->id, handle);
    return 0;
}

/* Memory Mapping */
int gsp_vf_mem_map(struct gsp_vf *vf, uint64_t handle, uint64_t gpu_va,
                   uint64_t size, uint32_t prot, uint32_t flags,
                   uint64_t offset, uint64_t *mapped_va)
{
    /* In emulation mode, just return the requested VA */
    if (emulation_mode) {
        if (mapped_va)
            *mapped_va = gpu_va;
        
        dev_dbg(vf->gsp->dev, "VF %u: map handle=0x%llx gpu_va=0x%llx size=0x%llx -> mapped_va=0x%llx (emulation)\n",
                vf->id, handle, gpu_va, size, gpu_va);
        
        return 0;
    }
    
    /* TODO: Implement GSP RPC for MAP */
    dev_dbg(vf->gsp->dev, "VF %u: map handle=0x%llx gpu_va=0x%llx size=0x%llx\n",
            vf->id, handle, gpu_va, size);
    
    if (mapped_va)
        *mapped_va = gpu_va;
    
    return 0;
}

int gsp_vf_mem_unmap(struct gsp_vf *vf, uint64_t handle, uint64_t gpu_va, uint64_t size)
{
    dev_dbg(vf->gsp->dev, "VF %u: unmap handle=0x%llx gpu_va=0x%llx size=0x%llx\n",
            vf->id, handle, gpu_va, size);
    return 0;
}

/* Scheduling */
int gsp_vf_schedule(struct gsp_vf *vf, uint32_t priority, uint64_t time_slice_ns)
{
    /* Use primary dispatcher SCHEDULE subcmd or secondary dispatcher */
    dev_dbg(vf->gsp->dev, "VF %u: schedule prio=%u slice=%llu ns\n",
            vf->id, priority, time_slice_ns);
    return 0;
}

int gsp_vf_yield(struct gsp_vf *vf)
{
    dev_dbg(vf->gsp->dev, "VF %u: yield\n", vf->id);
    return 0;
}