/*
 * virtio_rng.c — Virtio 1.0 entropy source driver for StarKernel
 *
 * Modern virtio 1.0 interface only (device ID 0x1044).
 * Falls back to checking 0x1005 (legacy/transitional).
 *
 * Split virtqueue, queue depth = 1 (a single device-writable buffer is
 * the entire request shape for this device — no header, no status byte,
 * no device-specific config space at all).
 *
 * The device may fill fewer bytes than the buffer offers per request
 * (the used-ring element's len says how many); virtio_rng_get_bytes()
 * loops internally until the caller's full byte count is satisfied.
 *
 * Memory model: all allocations via kmalloc(); identity-mapped so
 * virtual address == physical address for virtqueue ring pointers.
 * Transport plumbing (capability walk, common-cfg negotiation, split
 * virtqueue layout) mirrors virtio_blk.c exactly — same device family,
 * same board, same quirks (see that file's comments for why the config
 * structs are unpacked and why notify offset math looks the way it does).
 */

#ifndef __STARKERNEL__
#error "virtio_rng.c is kernel-only"
#endif

#include <stddef.h>
#include <stdint.h>
#include <string.h>

#include "starkernel/pci.h"
#include "starkernel/virtio_rng.h"
#include "starkernel/kmalloc.h"
#include "console.h"

/* -------------------------------------------------------------------------
 * Virtio 1.0 PCI capability structures (identical to virtio_blk.c)
 * ------------------------------------------------------------------------- */

#define VIRTIO_PCI_CAP_VENDOR_ID  0x09u

#define VIRTIO_PCI_CAP_COMMON_CFG 1u
#define VIRTIO_PCI_CAP_NOTIFY_CFG 2u

#define VCAP_OFF_CAP_VNDR   0u
#define VCAP_OFF_CAP_NEXT   1u
#define VCAP_OFF_CFG_TYPE   3u
#define VCAP_OFF_BAR        4u
#define VCAP_OFF_OFFSET     8u
#define VCAP_OFF_LENGTH    12u
#define VCAP_OFF_NOTIFY_MULT 16u

typedef struct {
    volatile uint32_t device_feature_select;
    volatile uint32_t device_feature;
    volatile uint32_t driver_feature_select;
    volatile uint32_t driver_feature;
    volatile uint16_t config_msix_vector;
    volatile uint16_t num_queues;
    volatile uint8_t  device_status;
    volatile uint8_t  config_generation;
    volatile uint16_t queue_select;
    volatile uint16_t queue_size;
    volatile uint16_t queue_msix_vector;
    volatile uint16_t queue_enable;
    volatile uint16_t queue_notify_off;
    volatile uint64_t queue_desc;
    volatile uint64_t queue_driver;
    volatile uint64_t queue_device;
    volatile uint16_t queue_notify_data;
    volatile uint16_t queue_reset;
} VirtioCommonCfg;

#define VIRTIO_STATUS_ACKNOWLEDGE  0x01u
#define VIRTIO_STATUS_DRIVER       0x02u
#define VIRTIO_STATUS_DRIVER_OK    0x04u
#define VIRTIO_STATUS_FEATURES_OK  0x08u
#define VIRTIO_STATUS_FAILED       0x80u

#define VIRTIO_F_VERSION_1         (1ULL << 32)

/* -------------------------------------------------------------------------
 * Split virtqueue structures (queue depth 1 — one descriptor, no chaining)
 * ------------------------------------------------------------------------- */

#define VQUEUE_SIZE  1u

#define VRING_DESC_F_WRITE  2u

typedef struct {
    uint64_t addr;
    uint32_t len;
    uint16_t flags;
    uint16_t next;
} VirtqDesc;

typedef struct {
    uint16_t flags;
    uint16_t idx;
    uint16_t ring[VQUEUE_SIZE];
    uint16_t used_event;
} VirtqAvail;

typedef struct {
    uint32_t id;
    uint32_t len;
} VirtqUsedElem;

typedef struct {
    uint16_t      flags;
    uint16_t      idx;
    VirtqUsedElem ring[VQUEUE_SIZE];
    uint16_t      avail_event;
} VirtqUsed;

/* -------------------------------------------------------------------------
 * Driver state
 * ------------------------------------------------------------------------- */

#define VRNG_BUF_SIZE 64u /* bytes requested from the device per round */

typedef struct {
    VirtioCommonCfg   *common;
    volatile uint16_t *notify;
    uint32_t           notify_off_mult;
    uint16_t           queue_notify_off;

    VirtqDesc  *desc;
    VirtqAvail *avail;
    VirtqUsed  *used;

    uint16_t avail_idx;
    uint16_t last_used_idx;

    uint8_t *data_buf; /* VRNG_BUF_SIZE bytes, DMA-accessible */
} VirtRngState;

static VirtRngState g_vrng;
static int          g_vrng_ready = 0;

/* -------------------------------------------------------------------------
 * Capability walker (identical shape to virtio_blk.c's walk_virtio_caps)
 * ------------------------------------------------------------------------- */

static void *walk_virtio_caps(const PciDevice *d, uint8_t cap_type,
                               uint32_t *extra_out) {
    uint8_t cap_ptr = pci_read8(d, (uint16_t)PCI_CFG_CAP_PTR) & 0xFCu;
    if (!cap_ptr) return NULL;

    int limit = 48;
    while (cap_ptr && limit--) {
        uint8_t vndr  = pci_read8(d, cap_ptr + (uint16_t)VCAP_OFF_CAP_VNDR);
        uint8_t next  = pci_read8(d, cap_ptr + (uint16_t)VCAP_OFF_CAP_NEXT);
        uint8_t ctype = pci_read8(d, cap_ptr + (uint16_t)VCAP_OFF_CFG_TYPE);

        if (vndr == (uint8_t)VIRTIO_PCI_CAP_VENDOR_ID && ctype == cap_type) {
            uint8_t  bar    = pci_read8 (d, cap_ptr + (uint16_t)VCAP_OFF_BAR);
            uint32_t offset = pci_read32(d, cap_ptr + (uint16_t)VCAP_OFF_OFFSET);
            uint32_t length = pci_read32(d, cap_ptr + (uint16_t)VCAP_OFF_LENGTH);

            if (bar > 5u) { cap_ptr = next & 0xFCu; continue; }

            uint64_t bar_base = pci_bar(d, (int)bar);
            if (!bar_base) { cap_ptr = next & 0xFCu; continue; }

            if (pci_map_bar(bar_base, (uint64_t)length + offset) != 0) {
                cap_ptr = next & 0xFCu; continue;
            }

            if (extra_out && cap_type == VIRTIO_PCI_CAP_NOTIFY_CFG) {
                *extra_out = pci_read32(d, cap_ptr + (uint16_t)VCAP_OFF_NOTIFY_MULT);
            }

            return (void *)(uintptr_t)(bar_base + offset);
        }

        cap_ptr = next & 0xFCu;
    }
    return NULL;
}

static inline void wmb(void) { __asm__ volatile("" : : : "memory"); }
static inline void rmb(void) { __asm__ volatile("" : : : "memory"); }

/* -------------------------------------------------------------------------
 * One round: ask the device to fill data_buf, return bytes actually
 * written (may be less than VRNG_BUF_SIZE).
 * ------------------------------------------------------------------------- */

static int vrng_request(uint32_t *bytes_out) {
    VirtRngState *s = &g_vrng;
    VirtqDesc    *d = s->desc;

    d[0].addr  = (uint64_t)(uintptr_t)s->data_buf;
    d[0].len   = (uint32_t)VRNG_BUF_SIZE;
    d[0].flags = (uint16_t)VRING_DESC_F_WRITE;
    d[0].next  = 0;

    uint16_t avail_idx = s->avail_idx & (uint16_t)(VQUEUE_SIZE - 1u);
    s->avail->ring[avail_idx] = 0;
    wmb();
    s->avail->idx = (uint16_t)(s->avail->idx + 1u);
    s->avail_idx  = s->avail->idx;
    wmb();

    uint16_t notify_idx = (uint16_t)(s->queue_notify_off *
                                      (s->notify_off_mult & 0xFFFFu));
    volatile uint16_t *doorbell = s->notify + notify_idx;
    *doorbell = 0;
    wmb();

    uint32_t spin = 0x2000000u;
    while (s->used->idx == s->last_used_idx) {
        rmb();
        if (!--spin) return -2;
    }

    uint32_t used_slot = (uint32_t)s->last_used_idx & (VQUEUE_SIZE - 1u);
    /* The device-controlled used length is trusted after this point:
     * data_buf is only VRNG_BUF_SIZE bytes, so clamp anything larger to
     * the buffer size to keep the caller's later data_buf copy in-bounds
     * against a buggy or malicious device. */
    uint32_t used_len = s->used->ring[used_slot].len;
    if (used_len > VRNG_BUF_SIZE) used_len = VRNG_BUF_SIZE;
    *bytes_out = used_len;
    s->last_used_idx = s->used->idx;

    return 0;
}

/* -------------------------------------------------------------------------
 * Device initialisation
 * ------------------------------------------------------------------------- */

static int vrng_init_device(const PciDevice *pci) {
    VirtRngState *s = &g_vrng;

    pci_enable(pci);

    uint32_t notify_mult = 0;
    VirtioCommonCfg *common = (VirtioCommonCfg *)
                              walk_virtio_caps(pci, VIRTIO_PCI_CAP_COMMON_CFG, NULL);
    volatile uint16_t *notify = (volatile uint16_t *)
                              walk_virtio_caps(pci, VIRTIO_PCI_CAP_NOTIFY_CFG, &notify_mult);

    if (!common || !notify) {
        console_println("virtio-rng: cap walk failed");
        return -2;
    }

    s->common          = common;
    s->notify          = notify;
    s->notify_off_mult = notify_mult;

    common->device_status = 0;
    wmb();
    common->device_status = (uint8_t)VIRTIO_STATUS_ACKNOWLEDGE;
    wmb();
    common->device_status = (uint8_t)(VIRTIO_STATUS_ACKNOWLEDGE | VIRTIO_STATUS_DRIVER);
    wmb();

    common->driver_feature_select = 1;
    wmb();
    common->driver_feature = (uint32_t)(VIRTIO_F_VERSION_1 >> 32);
    common->driver_feature_select = 0;
    wmb();
    common->driver_feature = 0;
    wmb();

    common->device_status = (uint8_t)(VIRTIO_STATUS_ACKNOWLEDGE |
                                      VIRTIO_STATUS_DRIVER |
                                      VIRTIO_STATUS_FEATURES_OK);
    wmb();
    rmb();
    if (!(common->device_status & (uint8_t)VIRTIO_STATUS_FEATURES_OK)) {
        console_println("virtio-rng: FEATURES_OK rejected");
        common->device_status = (uint8_t)VIRTIO_STATUS_FAILED;
        return -2;
    }

    common->queue_select = 0;
    wmb();
    uint16_t max_size = common->queue_size;
    if (max_size == 0) {
        console_println("virtio-rng: bad queue size");
        return -2;
    }
    uint16_t qsize = (uint16_t)VQUEUE_SIZE;
    common->queue_size = qsize;
    s->queue_notify_off = common->queue_notify_off;
    wmb();

    common->config_msix_vector = 0xFFFFu;
    common->queue_msix_vector  = 0xFFFFu;
    wmb();

    size_t desc_bytes  = (size_t)qsize * sizeof(VirtqDesc);
    size_t avail_bytes = sizeof(uint16_t) * 2u +
                         (size_t)qsize * sizeof(uint16_t) +
                         sizeof(uint16_t);
    size_t used_bytes  = sizeof(uint16_t) * 2u +
                         (size_t)qsize * sizeof(VirtqUsedElem) +
                         sizeof(uint16_t);

    s->desc  = (VirtqDesc  *)kmalloc_aligned(desc_bytes, 64);
    s->avail = (VirtqAvail *)kmalloc_aligned(avail_bytes, 2);
    s->used  = (VirtqUsed  *)kmalloc_aligned(used_bytes, 4);
    if (!s->desc || !s->avail || !s->used) {
        console_println("virtio-rng: queue alloc failed");
        return -2;
    }
    memset(s->desc,  0, desc_bytes);
    memset(s->avail, 0, avail_bytes);
    memset(s->used,  0, used_bytes);

    s->data_buf = (uint8_t *)kmalloc_aligned(VRNG_BUF_SIZE, 16);
    if (!s->data_buf) {
        console_println("virtio-rng: buf alloc failed");
        return -2;
    }

    s->avail_idx     = 0;
    s->last_used_idx = 0;

    common->queue_desc   = (uint64_t)(uintptr_t)s->desc;
    common->queue_driver = (uint64_t)(uintptr_t)s->avail;
    common->queue_device = (uint64_t)(uintptr_t)s->used;
    wmb();

    common->queue_enable = 1;
    wmb();

    common->device_status = (uint8_t)(VIRTIO_STATUS_ACKNOWLEDGE |
                                      VIRTIO_STATUS_DRIVER |
                                      VIRTIO_STATUS_FEATURES_OK |
                                      VIRTIO_STATUS_DRIVER_OK);
    wmb();

    return 0;
}

/* -------------------------------------------------------------------------
 * Public entry points
 * ------------------------------------------------------------------------- */

int virtio_rng_init(void) {
    PciDevice pci;
    int found = pci_find_first(VIRTIO_RNG_PCI_VENDOR_ID, VIRTIO_RNG_DEVICE_MODERN, &pci);
    if (found != 0)
        found = pci_find_first(VIRTIO_RNG_PCI_VENDOR_ID, VIRTIO_RNG_DEVICE_LEGACY, &pci);
    if (found != 0) {
        console_println("virtio-rng: no device on PCI bus 0");
        return -1;
    }

    console_println("virtio-rng: found device");

    int rc = vrng_init_device(&pci);
    if (rc != 0) return rc;

    g_vrng_ready = 1;
    return 0;
}

int virtio_rng_ready(void) {
    return g_vrng_ready;
}

int virtio_rng_get_bytes(uint8_t *buf, size_t n) {
    if (!g_vrng_ready) return -1;
    if (!buf) return -2;

    size_t filled = 0;
    uint32_t attempts_left = 64u; /* generous bound: real entropy always yields > 0 bytes */

    while (filled < n) {
        uint32_t got = 0;
        int rc = vrng_request(&got);
        if (rc != 0) return -2;

        if (got == 0) {
            if (!--attempts_left) return -2;
            continue;
        }

        size_t take = (size_t)got;
        if (take > n - filled) take = n - filled;
        memcpy(buf + filled, g_vrng.data_buf, take);
        filled += take;
    }

    return 0;
}