/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * GSP vGPU-IO Driver - RPC Client Implementation
 */

#include "gsp_core.h"
#include <linux/slab.h>
#include <linux/wait.h>
#include <linux/delay.h>
#include <linux/kthread.h>

/* Ring Buffer Operations */
static int gsp_ring_init(struct gsp_msg_ring *ring, void *base, size_t size)
{
    ring->base = base;
    ring->size = size;
    ring->head = 0;
    ring->tail = 0;
    spin_lock_init(&ring->lock);
    return 0;
}

static bool gsp_ring_empty(struct gsp_msg_ring *ring)
{
    return ring->head == ring->tail;
}

static bool gsp_ring_full(struct gsp_msg_ring *ring)
{
    return ((ring->head + 1) % (ring->size / GSP_MSG_MAX_SIZE)) == ring->tail;
}

static int gsp_ring_write(struct gsp_msg_ring *ring, struct gsp_rpc_message *msg, size_t len)
{
    unsigned long flags;
    uint32_t next_head;
    void *dst;
    
    if (len > GSP_MSG_MAX_SIZE)
        return -EINVAL;
    
    spin_lock_irqsave(&ring->lock, flags);
    
    if (gsp_ring_full(ring)) {
        spin_unlock_irqrestore(&ring->lock, flags);
        return -EBUSY;
    }
    
    dst = ring->base + (ring->head * GSP_MSG_MAX_SIZE);
    memcpy(dst, msg, len);
    
    next_head = (ring->head + 1) % (ring->size / GSP_MSG_MAX_SIZE);
    ring->head = next_head;
    
    spin_unlock_irqrestore(&ring->lock, flags);
    return 0;
}

static int gsp_ring_read(struct gsp_msg_ring *ring, struct gsp_rpc_message *msg)
{
    unsigned long flags;
    void *src;
    
    spin_lock_irqsave(&ring->lock, flags);
    
    if (gsp_ring_empty(ring)) {
        spin_unlock_irqrestore(&ring->lock, flags);
        return 0;
    }
    
    src = ring->base + (ring->tail * GSP_MSG_MAX_SIZE);
    memcpy(msg, src, GSP_MSG_MAX_SIZE);
    
    ring->tail = (ring->tail + 1) % (ring->size / GSP_MSG_MAX_SIZE);
    
    spin_unlock_irqrestore(&ring->lock, flags);
    return 1;
}

/* RPC Client Initialization */
struct gsp_rpc_client *gsp_rpc_client_init(struct gsp_device *gsp)
{
    struct gsp_rpc_client *rpc;
    void *ring_mem;
    
    rpc = kzalloc(sizeof(*rpc), GFP_KERNEL);
    if (!rpc)
        return ERR_PTR(-ENOMEM);
    
    rpc->gsp = gsp;
    mutex_init(&rpc->lock);
    atomic_set(&rpc->seq, 0);
    init_waitqueue_head(&rpc->wq);
    
    /* Allocate shared memory for message rings (2 pages each) */
    ring_mem = dma_alloc_coherent(gsp->dev, 8 * PAGE_SIZE,
                                  &rpc->tx_ring_dma, GFP_KERNEL);
    if (!ring_mem) {
        kfree(rpc);
        return ERR_PTR(-ENOMEM);
    }
    
    rpc->tx_ring = kzalloc(sizeof(*rpc->tx_ring), GFP_KERNEL);
    rpc->rx_ring = kzalloc(sizeof(*rpc->rx_ring), GFP_KERNEL);
    if (!rpc->tx_ring || !rpc->rx_ring) {
        dma_free_coherent(gsp->dev, 8 * PAGE_SIZE, ring_mem, rpc->tx_ring_dma);
        kfree(rpc->tx_ring);
        kfree(rpc->rx_ring);
        kfree(rpc);
        return ERR_PTR(-ENOMEM);
    }
    
    gsp_ring_init(rpc->tx_ring, ring_mem, 4 * PAGE_SIZE);
    gsp_ring_init(rpc->rx_ring, ring_mem + 4 * PAGE_SIZE, 4 * PAGE_SIZE);
    
    /* Map rings to GSP via queue head registers (BAR0) */
    writel(rpc->tx_ring_dma, gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0));
    writel(rpc->tx_ring_dma >> 32, gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0) + 4);
    writel(rpc->rx_ring_dma, gsp->bar0 + GSP_BAR0_QUEUE_HEAD(1));
    writel(rpc->rx_ring_dma >> 32, gsp->bar0 + GSP_BAR0_QUEUE_HEAD(1) + 4);
    
    /* Allocate response buffers */
    for (int i = 0; i < GSP_MAX_PENDING; i++) {
        rpc->resp_buf[i] = kzalloc(GSP_MSG_MAX_SIZE, GFP_KERNEL);
        if (!rpc->resp_buf[i]) {
            /* Cleanup on failure */
            for (int j = 0; j < i; j++)
                kfree(rpc->resp_buf[j]);
            gsp_rpc_client_fini(rpc);
            return ERR_PTR(-ENOMEM);
        }
    }
    
    dev_info(gsp->dev, "RPC client initialized\n");
    return rpc;
}

void gsp_rpc_client_fini(struct gsp_rpc_client *rpc)
{
    if (!rpc)
        return;
    
    for (int i = 0; i < GSP_MAX_PENDING; i++)
        kfree(rpc->resp_buf[i]);
    
    if (rpc->tx_ring) {
        dma_free_coherent(rpc->gsp->dev, 8 * PAGE_SIZE,
                          rpc->tx_ring->base, rpc->tx_ring_dma);
        kfree(rpc->tx_ring);
    }
    kfree(rpc->rx_ring);
    kfree(rpc);
}

/* Synchronous RPC Send */
int gsp_rpc_send_sync(struct gsp_rpc_client *rpc,
                      enum gsp_rpc_type type,
                      void *payload, size_t payload_size,
                      void *response, size_t *response_size)
{
    struct gsp_rpc_message msg = {0};
    struct completion done;
    uint32_t seq;
    int ret;
    
    if (payload_size > GSP_MSG_MAX_SIZE - sizeof(struct gsp_rpc_header))
        return -EINVAL;
    
    seq = atomic_inc_return(&rpc->seq);
    init_completion(&done);
    
    msg.hdr.magic = 0x47535052;  /* "GSPR" */
    msg.hdr.version = 1;
    msg.hdr.rpc_type = type;
    msg.hdr.sequence = seq;
    msg.hdr.timestamp_ns = ktime_get_ns();
    msg.hdr.payload_size = payload_size;
    msg.hdr.flags = GSP_MSG_SYNC;
    msg.hdr.correlation_id = seq;
    
    if (payload && payload_size > 0)
        memcpy(msg.payload, payload, payload_size);
    
    mutex_lock(&rpc->lock);
    
    rpc->pending[seq % GSP_MAX_PENDING] = &done;
    
    ret = gsp_ring_write(rpc->tx_ring, &msg, GSP_MSG_MAX_SIZE);
    if (ret) {
        rpc->pending[seq % GSP_MAX_PENDING] = NULL;
        mutex_unlock(&rpc->lock);
        return ret;
    }
    
    /* Ring doorbell - write to queue head to signal GSP */
    writel(1, rpc->gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0));
    
    mutex_unlock(&rpc->lock);
    
    /* Wait for completion (5 second timeout) */
    if (!wait_for_completion_timeout(&done, msecs_to_jiffies(5000))) {
        mutex_lock(&rpc->lock);
        rpc->pending[seq % GSP_MAX_PENDING] = NULL;
        mutex_unlock(&rpc->lock);
        atomic64_inc(&rpc->gsp->rpc_errors);
        return -ETIMEDOUT;
    }
    
    /* Get response */
    struct gsp_rpc_response *resp = rpc->resp_buf[seq % GSP_MAX_PENDING];
    
    atomic64_inc(&rpc->gsp->rpc_total);
    
    if (resp->status != 0) {
        atomic64_inc(&rpc->gsp->rpc_errors);
        return gsp_errno_to_errno(resp->status);
    }
    
    if (response && response_size) {
        size_t copy_size = min(*response_size, (size_t)resp->payload_size);
        memcpy(response, resp->payload, copy_size);
        *response_size = copy_size;
    }
    
    return 0;
}

/* Async RPC Send */
int gsp_rpc_send_async(struct gsp_rpc_client *rpc,
                       enum gsp_rpc_type type,
                       void *payload, size_t payload_size,
                       void (*callback)(struct gsp_rpc_client *,
                                        struct gsp_rpc_response *,
                                        void *),
                       void *ctx)
{
    struct gsp_rpc_message msg = {0};
    uint32_t seq;
    int ret;
    
    if (payload_size > GSP_MSG_MAX_SIZE - sizeof(struct gsp_rpc_header))
        return -EINVAL;
    
    seq = atomic_inc_return(&rpc->seq);
    
    msg.hdr.magic = 0x47535052;
    msg.hdr.version = 1;
    msg.hdr.rpc_type = type;
    msg.hdr.sequence = seq;
    msg.hdr.timestamp_ns = ktime_get_ns();
    msg.hdr.payload_size = payload_size;
    msg.hdr.flags = GSP_MSG_ASYNC;
    msg.hdr.correlation_id = (uint64_t)callback | ((uint64_t)ctx << 32);
    
    if (payload && payload_size > 0)
        memcpy(msg.payload, payload, payload_size);
    
    mutex_lock(&rpc->lock);
    ret = gsp_ring_write(rpc->tx_ring, &msg, GSP_MSG_MAX_SIZE);
    if (!ret)
        writel(1, rpc->gsp->bar0 + GSP_BAR0_QUEUE_HEAD(0));
    mutex_unlock(&rpc->lock);
    
    return ret;
}

/* ISR - Handle Completions */
irqreturn_t gsp_isr(int irq, void *data)
{
    struct gsp_device *gsp = data;
    struct gsp_rpc_message msg;
    int handled = 0;
    
    while (gsp_ring_read(gsp->rpc->rx_ring, &msg)) {
        handled = 1;
        
        if (msg.hdr.flags & GSP_MSG_SYNC) {
            uint32_t seq = msg.hdr.sequence;
            struct completion *done = gsp->rpc->pending[seq % GSP_MAX_PENDING];
            struct gsp_rpc_response *resp = gsp->rpc->resp_buf[seq % GSP_MAX_PENDING];
            
            if (done && resp) {
                resp->status = ((struct gsp_rpc_response *)msg.payload)->status;
                resp->payload_size = msg.hdr.payload_size;
                memcpy(resp->payload, ((struct gsp_rpc_response *)msg.payload)->payload,
                       msg.hdr.payload_size);
                complete(done);
            }
            gsp->rpc->pending[seq % GSP_MAX_PENDING] = NULL;
        } else {
            /* Async notification - handle via callback */
            gsp_handle_notification(gsp, &msg);
        }
    }
    
    return handled ? IRQ_HANDLED : IRQ_NONE;
}

/* Notification Handler */
void gsp_handle_notification(struct gsp_device *gsp, struct gsp_rpc_message *msg)
{
    switch (msg->hdr.rpc_type) {
    case GSP_RPC_NOTIFY:
        /* Handle VF notifications (WPR, migration, etc.) */
        break;
    case GSP_RPC_BAR_TRAP:
        atomic64_inc(&gsp->bar_traps);
        break;
    default:
        dev_dbg(gsp->dev, "Unhandled notification type 0x%02x\n", msg->hdr.rpc_type);
    }
}

/* Error Code Translation */
int gsp_errno_to_errno(int gsp_err)
{
    switch (gsp_err) {
    case GSP_OK: return 0;
    case GSP_EPERM: return -EPERM;
    case GSP_ENOENT: return -ENOENT;
    case GSP_EINVAL: return -EINVAL;
    case GSP_ENODEV: return -ENODEV;
    case GSP_ENOMEM: return -ENOMEM;
    case GSP_EBUSY: return -EBUSY;
    case GSP_ENOSPC: return -ENOSPC;
    case GSP_ETIMEDOUT: return -ETIMEDOUT;
    case GSP_EOVERFLOW: return -EOVERFLOW;
    case GSP_EDEV: return -EIO;
    default: return -EIO;
    }
}