/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * GSP vGPU-IO Driver - WPR (Write/Read Protect Region) Management
 */

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

/* WPR Payload (for RPC) */
struct gsp_wpr_payload {
    uint32_t vf_id;
    uint32_t op;          /* 0=protect_write, 1=protect_read, 2=unprotect */
    uint64_t base;
    uint64_t size;
    uint32_t flags;
    uint64_t timeout_ns;
} __attribute__((packed));

struct gsp_wpr_response {
    int64_t status;
} __attribute__((packed));

/* WPR Initialization */
int gsp_vf_wpr_init(struct gsp_vf *vf)
{
    INIT_LIST_HEAD(&vf->wpr_list);
    mutex_init(&vf->wpr_lock);
    return 0;
}

void gsp_vf_wpr_fini(struct gsp_vf *vf)
{
    struct gsp_wpr_region *rgn, *tmp;
    
    mutex_lock(&vf->wpr_lock);
    list_for_each_entry_safe(rgn, tmp, &vf->wpr_list, list) {
        list_del(&rgn->list);
        kfree(rgn);
    }
    mutex_unlock(&vf->wpr_lock);
}

/* Protect Region (Write) */
int gsp_vf_wpr_protect_write(struct gsp_vf *vf, uint64_t base, uint64_t size)
{
    return gsp_vf_wpr_protect(vf, base, size, GSP_WPR_WRITE);
}

/* Protect Region (Read) */
int gsp_vf_wpr_protect_read(struct gsp_vf *vf, uint64_t base, uint64_t size)
{
    return gsp_vf_wpr_protect(vf, base, size, GSP_WPR_READ);
}

/* Generic Protect */
int gsp_vf_wpr_protect(struct gsp_vf *vf, uint64_t base, uint64_t size, uint32_t flags)
{
    struct gsp_wpr_payload payload = {
        .vf_id = vf->id,
        .op = (flags & GSP_WPR_READ) ? 1 : 0,
        .base = base,
        .size = size,
        .flags = flags,
        .timeout_ns = 10000000,  /* 10ms default */
    };
    struct gsp_wpr_response resp;
    size_t resp_size = sizeof(resp);
    struct gsp_wpr_region *rgn;
    int ret;
    
    /* Validate alignment */
    if (base & 0xFFF || size & 0xFFF)
        return -EINVAL;
    
    rgn = kzalloc(sizeof(*rgn), GFP_KERNEL);
    if (!rgn)
        return -ENOMEM;
    
    rgn->base = base;
    rgn->size = size;
    rgn->flags = flags;
    rgn->vf_id = vf->id;
    rgn->timeout_ns = payload.timeout_ns;
    
    /* In emulation mode, handle locally */
    if (emulation_mode) {
        mutex_lock(&vf->wpr_lock);
        list_add(&rgn->list, &vf->wpr_list);
        mutex_unlock(&vf->wpr_lock);
        
        dev_dbg(vf->gsp->dev, "VF %u: WPR protect [0x%llx, 0x%llx) flags=0x%x (emulation)\n",
                vf->id, base, base + size, flags);
        
        return 0;
    }
    
    /* Send to GSP WPR handler (0x04F00B0C/0x04F00C78) */
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MAP,  /* Uses MAP RPC for WPR */
                            &payload, sizeof(payload), &resp, &resp_size);
    if (ret) {
        kfree(rgn);
        return ret;
    }
    
    if (resp.status != 0) {
        kfree(rgn);
        return gsp_errno_to_errno(resp.status);
    }
    
    mutex_lock(&vf->wpr_lock);
    list_add(&rgn->list, &vf->wpr_list);
    mutex_unlock(&vf->wpr_lock);
    
    dev_dbg(vf->gsp->dev, "VF %u: WPR protect [0x%llx, 0x%llx) flags=0x%x\n",
            vf->id, base, base + size, flags);
    
    return 0;
}

/* Unprotect Region */
int gsp_vf_wpr_unprotect(struct gsp_vf *vf, uint64_t base)
{
    struct gsp_wpr_payload payload = {
        .vf_id = vf->id,
        .op = 2,  /* unprotect */
        .base = base,
        .size = 0,
        .flags = 0,
    };
    struct gsp_wpr_response resp;
    size_t resp_size = sizeof(resp);
    struct gsp_wpr_region *rgn, *tmp;
    bool found = false;
    int ret;
    
    /* In emulation mode, handle locally */
    if (emulation_mode) {
        mutex_lock(&vf->wpr_lock);
        list_for_each_entry_safe(rgn, tmp, &vf->wpr_list, list) {
            if (rgn->base == base) {
                list_del(&rgn->list);
                kfree(rgn);
                found = true;
                break;
            }
        }
        mutex_unlock(&vf->wpr_lock);
        
        dev_dbg(vf->gsp->dev, "VF %u: WPR unprotect 0x%llx (emulation)\n", vf->id, base);
        
        return found ? 0 : -ENOENT;
    }
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_UNMAP,  /* Uses UNMAP for WPR */
                            &payload, sizeof(payload), &resp, &resp_size);
    if (ret)
        return ret;
    
    if (resp.status != 0)
        return gsp_errno_to_errno(resp.status);
    
    mutex_lock(&vf->wpr_lock);
    list_for_each_entry_safe(rgn, tmp, &vf->wpr_list, list) {
        if (rgn->base == base) {
            list_del(&rgn->list);
            kfree(rgn);
            found = true;
            break;
        }
    }
    mutex_unlock(&vf->wpr_lock);
    
    if (!found)
        return -ENOENT;
    
    dev_dbg(vf->gsp->dev, "VF %u: WPR unprotect 0x%llx\n", vf->id, base);
    
    return 0;
}

/* Set Timeout */
int gsp_vf_wpr_set_timeout(struct gsp_vf *vf, uint64_t timeout_ns)
{
    /* Update timeout for all WPR regions */
    struct gsp_wpr_region *rgn;
    
    mutex_lock(&vf->wpr_lock);
    list_for_each_entry(rgn, &vf->wpr_list, list) {
        rgn->timeout_ns = timeout_ns;
    }
    mutex_unlock(&vf->wpr_lock);
    
    return 0;
}

/* Register Completion Callback */
int gsp_vf_wpr_register_callback(struct gsp_vf *vf, enum gsp_wpr_event type,
                                 void (*cb)(struct gsp_vf *, uint64_t, uint32_t),
                                 void *ctx)
{
    struct gsp_wpr_region *rgn;
    
    mutex_lock(&vf->wpr_lock);
    list_for_each_entry(rgn, &vf->wpr_list, list) {
        switch (type) {
        case GSP_WPR_WRITE_COMPLETE:
        case GSP_WPR_WRITE_TIMEOUT:
            rgn->completion_cb = cb;
            break;
        case GSP_WPR_READ_COMPLETE:
        case GSP_WPR_READ_TIMEOUT:
            rgn->timeout_cb = cb;
            break;
        }
    }
    mutex_unlock(&vf->wpr_lock);
    
    return 0;
}

/* WPR Event Handler (called from notification) */
void gsp_wpr_handle_event(struct gsp_device *gsp, struct gsp_rpc_message *msg)
{
    struct gsp_rpc_response *resp = (struct gsp_rpc_response *)msg->payload;
    uint32_t event_type = resp->payload[0];
    uint32_t vf_id = resp->payload[1];
    uint64_t address = *(uint64_t *)&resp->payload[2];
    struct gsp_vf *vf;
    struct gsp_wpr_region *rgn;
    
    vf = gsp_vf_get(gsp, vf_id);
    if (!vf)
        return;
    
    atomic64_inc(&gsp->wpr_events);
    
    mutex_lock(&vf->wpr_lock);
    list_for_each_entry(rgn, &vf->wpr_list, list) {
        if (address >= rgn->base && address < rgn->base + rgn->size) {
            switch (event_type) {
            case GSP_WPR_WRITE_COMPLETE:
            case GSP_WPR_READ_COMPLETE:
                if (rgn->completion_cb)
                    rgn->completion_cb(vf, address, event_type);
                break;
            case GSP_WPR_WRITE_TIMEOUT:
            case GSP_WPR_READ_TIMEOUT:
                if (rgn->timeout_cb)
                    rgn->timeout_cb(vf, address, event_type);
                break;
            }
            break;
        }
    }
    mutex_unlock(&vf->wpr_lock);
    
    gsp_vf_put(vf);
}

/* Default Completion Callback */
void gsp_wpr_default_complete_cb(struct gsp_vf *vf, uint64_t addr, uint32_t type)
{
    dev_dbg(vf->gsp->dev, "VF %u: WPR %s complete at 0x%llx\n",
            vf->id, type == GSP_WPR_WRITE_COMPLETE ? "write" : "read", addr);
    
    /* If this was a page table write, invalidate TLB */
    if (addr >= vf->pt_base && addr < vf->pt_base + vf->pt_size) {
        /* TODO: TLB invalidate */
    }
}

/* Default Timeout Callback */
void gsp_wpr_default_timeout_cb(struct gsp_vf *vf, uint64_t addr, uint32_t type)
{
    dev_warn(vf->gsp->dev, "VF %u: WPR %s timeout at 0x%llx\n",
             vf->id, type == GSP_WPR_WRITE_TIMEOUT ? "write" : "read", addr);
    
    /* Throttle or reset VF */
    if (type == GSP_WPR_WRITE_TIMEOUT) {
        /* TODO: VF throttle */
    } else {
        /* TODO: VF reset */
    }
}