//! virtio-rng: random bytes from the host (QEMU reads the host's
//! /dev/urandom). `-device virtio-rng-pci`.
//!
//! The simplest virtio device there is: one queue, and a request is just an
//! empty buffer the device fills with random bytes. We poll for the answer
//! (it takes microseconds) instead of waiting for an interrupt.
//!
//! Why bother? A computer is deterministic; good randomness has to come
//! from outside (hardware noise -- or here, the host). TCP needs
//! unpredictable sequence numbers and ports, or an attacker can guess them
//! and inject packets; programs need random keys and seeds. Without the
//! device we fall back to a generator seeded from the CPU's time stamp
//! counter -- fine for games, *not* for secrets, and we say so in the log.
//! (Real kernels also mix in the CPU's RDRAND and interrupt timings.)

use super::{Transport, Virtqueue};
use crate::memory;
use crate::pci::PciDevice;
use alloc::{format, string::String};
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use spin::Mutex;
use x86_64::{PhysAddr, instructions::interrupts::without_interrupts};

struct Rng {
    transport: Transport,
    queue: Virtqueue,
    buf_phys: u64, // one page the device writes random bytes into
}

static RNG: Mutex<Option<Rng>> = Mutex::new(None);
static HAVE_DEVICE: AtomicBool = AtomicBool::new(false);
/// State of the fallback generator (xorshift64*).
static FALLBACK: AtomicU64 = AtomicU64::new(0);

/// Is this the virtio entropy device? (transitional id 0x1005, modern 0x1044)
pub fn is_virtio_rng(d: &PciDevice) -> bool {
    d.vendor_id == super::VENDOR && matches!(d.device_id, 0x1005 | 0x1044)
}

pub fn init(devices: &[PciDevice]) -> String {
    let Some(pci) = devices.iter().find(|d| is_virtio_rng(d)) else {
        return String::from("rng: no virtio-rng, using a TSC-seeded generator (weak!)");
    };
    let setup = || -> Option<Rng> {
        let transport = Transport::new(pci)?;
        transport.begin_init(0)?;
        let queue = transport.setup_queue(0, 8, false)?;
        transport.driver_ok();
        let buf_phys = memory::alloc_contiguous_frames(1)?.as_u64();
        Some(Rng { transport, queue, buf_phys })
    };
    match setup() {
        Some(rng) => {
            without_interrupts(|| *RNG.lock() = Some(rng));
            HAVE_DEVICE.store(true, Ordering::Relaxed);
            format!("rng: virtio-rng at {:02x}:{:02x}.{}", pci.bus, pci.device, pci.function)
        }
        None => String::from("rng: virtio-rng found but setup failed, using a TSC-seeded generator (weak!)"),
    }
}

/// Is the host's entropy device in use (vs the weak fallback)?
pub fn available() -> bool {
    HAVE_DEVICE.load(Ordering::Relaxed)
}

/// Fill `buf` with random bytes. Returns false if the bytes came from the
/// weak fallback generator.
pub fn fill(buf: &mut [u8]) -> bool {
    let mut done = 0;
    while done < buf.len() && available() {
        match device_read(&mut buf[done..]) {
            Some(n) => done += n,
            None => {
                // No answer within a second: stop trusting the device.
                HAVE_DEVICE.store(false, Ordering::Relaxed);
                crate::serial_println!("[rng] virtio-rng stopped answering, using the fallback");
            }
        }
    }
    if done < buf.len() {
        fallback_fill(&mut buf[done..]);
        return false;
    }
    true
}

/// One request: up to a page of random bytes. Returns how many we got.
fn device_read(buf: &mut [u8]) -> Option<usize> {
    without_interrupts(|| {
        let mut guard = RNG.lock();
        let rng = guard.as_mut()?;
        let want = buf.len().min(4096);
        rng.queue.push(&[(rng.buf_phys, want as u32, true)])?;
        rng.queue.notify();
        let (_, len) = rng.queue.wait_used(1_000_000)?;
        let n = (len as usize).min(want);
        let src: *const u8 = memory::phys_to_virt(PhysAddr::new(rng.buf_phys)).as_ptr();
        buf[..n].copy_from_slice(unsafe { core::slice::from_raw_parts(src, n) });
        Some(n)
    })
}

/// A random u64 (for seeds and port numbers).
pub fn u64() -> u64 {
    let mut bytes = [0u8; 8];
    fill(&mut bytes);
    u64::from_le_bytes(bytes)
}

/// xorshift64* (Marsaglia / Vigna), seeded and stirred with the TSC.
fn fallback_fill(buf: &mut [u8]) {
    let tsc = unsafe { core::arch::x86_64::_rdtsc() };
    let mut x = FALLBACK.load(Ordering::Relaxed) ^ tsc.rotate_left(17) | 1;
    for chunk in buf.chunks_mut(8) {
        x ^= x >> 12;
        x ^= x << 25;
        x ^= x >> 27;
        let out = x.wrapping_mul(0x2545_f491_4f6c_dd1d).to_le_bytes();
        chunk.copy_from_slice(&out[..chunk.len()]);
    }
    FALLBACK.store(x, Ordering::Relaxed);
}

/// Acknowledge a (shared-line) interrupt; we poll, so nothing else to do.
pub fn handle_irq() {
    if let Some(guard) = RNG.try_lock() {
        if let Some(rng) = guard.as_ref() {
            rng.transport.read_isr();
        }
    }
}