//! Host-side KVM binary stats reader via KVM_GET_STATS_FD.
//!
//! The ioctl returns an fd whose contents are a self-describing binary
//! format: header -> descriptors -> data. Descriptors are static
//! (parsed once at open time). Data values are live u64s re-read via
//! pread at the header's data_offset.
//!
//! Stats fds have `noop_llseek` but `FMODE_PREAD`. The initial
//! sequential read works (position starts at 0); subsequent reads
//! use pread.
//!
//! `KVM_GET_STATS_FD` and `struct kvm_stats_desc` live in the generic
//! `include/uapi/linux/kvm.h` and are available on every KVM-supported
//! architecture. The descriptor schema (per-stat names, sizes, offsets)
//! is published by each arch's `kvm_vcpu_stats_desc[]` (e.g.
//! `arch/x86/kvm/x86.c`, `arch/arm64/kvm/guest.c`). The reader below
//! parses whatever descriptors the kernel publishes; arch-specific
//! stat names that don't appear in a given kernel's descriptor list
//! are silently absent from the resulting map and consumers that look
//! them up via [`crate::vmm::KvmStatsTotals::sum`] /
//! [`crate::vmm::KvmStatsTotals::avg`] return 0.

use std::collections::HashMap;
use std::os::fd::OwnedFd;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use vmm_sys_util::ioctl_io_nr;

ioctl_io_nr!(KVM_GET_STATS_FD, kvm_bindings::KVMIO, 0xce);

/// Size of `kvm_stats_desc` fixed fields before the variable-length name.
///
/// Mirrors `struct kvm_stats_desc` from `include/uapi/linux/kvm.h`:
///   __u32 flags;        //  0..4
///   __s16 exponent;     //  4..6
///   __u16 size;         //  6..8
///   __u32 offset;       //  8..12
///   __u32 bucket_size;  // 12..16
///   char  name[];       // 16..
/// The reads below at offsets 6 (size) and 8 (offset) and the
/// `DESC_FIXED_SIZE..DESC_FIXED_SIZE + name_size` slice for the
/// trailing name string match this layout exactly. `flags`, `exponent`,
/// and `bucket_size` are unused (only scalar value reads are emitted,
/// and bucket_size only applies to histogram stats).
const DESC_FIXED_SIZE: usize = 16;

/// Issue KVM_GET_STATS_FD on an fd (VmFd or VcpuFd).
fn get_stats_fd<F: AsRawFd>(fd: &F) -> Option<OwnedFd> {
    let ret = unsafe { libc::ioctl(fd.as_raw_fd(), KVM_GET_STATS_FD()) };
    if ret < 0 {
        let err = std::io::Error::last_os_error();
        tracing::warn!(
            errno = err.raw_os_error().unwrap_or(0),
            error = %err,
            "KVM_GET_STATS_FD ioctl failed"
        );
        None
    } else {
        Some(unsafe { OwnedFd::from_raw_fd(ret) })
    }
}

/// Parsed descriptor for a single stat.
#[derive(Debug, Clone)]
struct StatDesc {
    name: String,
    /// Number of u64 values (1 for scalars).
    size: usize,
    /// Byte offset from the start of the data section.
    offset: usize,
}

/// Metadata parsed from the stats fd. Static for the fd's lifetime.
#[derive(Debug)]
struct StatsMeta {
    /// File offset where data values begin.
    data_offset: usize,
    /// Byte size of the data section.
    data_size: usize,
    /// Parsed descriptors.
    descs: Vec<StatDesc>,
}

/// Read the entire stats fd content via sequential read. Called
/// once from open(); position tracking relies on initial sequential
/// read (position starts at 0; noop_llseek prevents resetting).
fn read_initial(fd: RawFd) -> Option<Vec<u8>> {
    let mut buf = Vec::with_capacity(8192);
    let mut chunk = [0u8; 4096];
    loop {
        let n = loop {
            let r = unsafe { libc::read(fd, chunk.as_mut_ptr() as *mut _, chunk.len()) };
            if r >= 0 {
                break r;
            }
            let err = std::io::Error::last_os_error();
            if err.raw_os_error() == Some(libc::EINTR) {
                continue;
            }
            tracing::warn!(
                errno = err.raw_os_error().unwrap_or(0),
                error = %err,
                "KVM stats fd read failed"
            );
            return None;
        };
        if n == 0 {
            break;
        }
        buf.extend_from_slice(&chunk[..n as usize]);
    }
    if buf.len() < 24 {
        return None;
    }
    Some(buf)
}

/// Parse header + descriptors from the initial read buffer.
fn parse_meta_from_buf(buf: &[u8]) -> Option<StatsMeta> {
    let name_size = u32::from_ne_bytes(buf[4..8].try_into().ok()?) as usize;
    let num_desc = u32::from_ne_bytes(buf[8..12].try_into().ok()?) as usize;
    let desc_offset = u32::from_ne_bytes(buf[16..20].try_into().ok()?) as usize;
    let data_offset = u32::from_ne_bytes(buf[20..24].try_into().ok()?) as usize;

    let desc_stride = DESC_FIXED_SIZE + name_size;

    let mut descs = Vec::with_capacity(num_desc);
    for i in 0..num_desc {
        let d_start = desc_offset + i * desc_stride;
        if d_start + desc_stride > buf.len() {
            break;
        }
        let d = &buf[d_start..d_start + desc_stride];

        let size = u16::from_ne_bytes(d[6..8].try_into().ok()?) as usize;
        let offset = u32::from_ne_bytes(d[8..12].try_into().ok()?) as usize;

        let name_bytes = &d[DESC_FIXED_SIZE..DESC_FIXED_SIZE + name_size];
        let name_end = name_bytes.iter().position(|&b| b == 0).unwrap_or(name_size);
        let name = String::from_utf8_lossy(&name_bytes[..name_end]).to_string();

        descs.push(StatDesc { name, size, offset });
    }

    let data_size = buf.len().saturating_sub(data_offset);

    Some(StatsMeta {
        data_offset,
        data_size,
        descs,
    })
}

/// Read live data section via pread.
fn pread_data(fd: RawFd, meta: &StatsMeta) -> Option<Vec<u8>> {
    let mut buf = vec![0u8; meta.data_size];
    let n = unsafe {
        libc::pread(
            fd,
            buf.as_mut_ptr() as *mut _,
            meta.data_size,
            meta.data_offset as libc::off_t,
        )
    };
    if n < 0 {
        return None;
    }
    buf.truncate(n as usize);
    Some(buf)
}

/// Extract scalar stat values from a data buffer.
/// Descriptor offsets are byte offsets into the data section.
fn extract_stats(meta: &StatsMeta, data: &[u8]) -> HashMap<String, u64> {
    let mut map = HashMap::with_capacity(meta.descs.len());
    for desc in &meta.descs {
        if desc.size != 1 {
            continue;
        }
        let off = desc.offset;
        if off + 8 > data.len() {
            continue;
        }
        let val = u64::from_ne_bytes(data[off..off + 8].try_into().unwrap());
        map.insert(desc.name.clone(), val);
    }
    map
}

/// Per-vCPU stats reader. Holds the stats fd and pre-parsed metadata.
/// Opened before vCPUs move to threads; the stats fd holds a kernel
/// reference independent of VcpuFd ownership.
struct VcpuStatsReader {
    fd: OwnedFd,
    meta: StatsMeta,
}

impl VcpuStatsReader {
    /// Open a stats fd and parse metadata from the initial read.
    fn open<F: AsRawFd>(vcpu: &F) -> Option<Self> {
        let fd = get_stats_fd(vcpu)?;
        let buf = read_initial(fd.as_raw_fd())?;
        let meta = parse_meta_from_buf(&buf)?;
        Some(VcpuStatsReader { fd, meta })
    }

    /// Read current absolute stat values via pread.
    fn read_snapshot(&self) -> HashMap<String, u64> {
        let data = match pread_data(self.fd.as_raw_fd(), &self.meta) {
            Some(d) => d,
            None => return HashMap::new(),
        };
        extract_stats(&self.meta, &data)
    }
}

/// Holds pre-opened stats readers for all vCPUs. Read after VM exit
/// to capture cumulative totals. Same lifetime constraint as
/// [`VcpuStatsReader`] — opened before vCPU threads start.
pub(crate) struct StatsContext {
    readers: Vec<VcpuStatsReader>,
}

impl StatsContext {
    /// Read absolute cumulative stats from each vCPU after VM exit.
    pub(crate) fn read_stats(&self) -> crate::vmm::KvmStatsTotals {
        let per_vcpu = self.readers.iter().map(|r| r.read_snapshot()).collect();
        crate::vmm::KvmStatsTotals { per_vcpu }
    }
}

/// Open stats fds for all vCPUs. Called before vCPUs move to threads.
/// Returns None if KVM_GET_STATS_FD is not supported.
pub(crate) fn open_stats_context(vcpus: &[kvm_ioctls::VcpuFd]) -> Option<StatsContext> {
    let mut readers = Vec::with_capacity(vcpus.len());
    for vcpu in vcpus {
        readers.push(VcpuStatsReader::open(vcpu)?);
    }
    Some(StatsContext { readers })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::vmm::kvm::KtstrKvm;
    use crate::vmm::topology::Topology;

    /// Single-LLC, single-NUMA topology for stats-fd unit tests.
    /// `cores_per_llc` is the only dimension the call sites vary;
    /// every other field stays at the minimum value KtstrKvm accepts.
    fn test_topo(cores_per_llc: u32) -> Topology {
        Topology {
            llcs: 1,
            cores_per_llc,
            threads_per_core: 1,
            numa_nodes: 1,
            nodes: None,
            distances: None,
        }
    }

    #[test]
    fn stats_fd_returns_some_on_modern_kernel() {
        // KVM_GET_STATS_FD is generic-uapi
        // (KVM_CAP_BINARY_STATS_FD in
        // kvm_vm_ioctl_check_extension_generic), present on every
        // KVM since Linux 5.17. A host that can create a VM
        // (KtstrKvm::new succeeds below) is such a kernel, so the
        // ioctl on the VM fd MUST yield a stats fd — not the
        // vacuous "if Some" the previous version allowed.
        let vm = KtstrKvm::new(test_topo(1), 64, false).unwrap();
        let fd = get_stats_fd(&*vm.vm_fd).expect("KVM_GET_STATS_FD on VM fd must return Some");
        // The fd must be usable, not merely a valid number: the
        // initial read returns the self-describing header+desc
        // blob and parse_meta yields a non-empty descriptor list.
        let buf = read_initial(fd.as_raw_fd()).expect("initial read of VM stats fd");
        let meta = parse_meta_from_buf(&buf).expect("parse VM stats metadata");
        assert!(
            !meta.descs.is_empty(),
            "VM stats fd must publish at least one descriptor",
        );
    }

    #[test]
    fn vcpu_stats_fd_returns_some() {
        // Same generic-uapi guarantee as the VM-level test:
        // KVM_GET_STATS_FD on a vCPU fd MUST return Some on any
        // KVM that could create the vCPU. Exercise the fd through
        // parse and assert the per-vCPU "exits" stat is present so
        // the test proves usable stats, not just a non-negative
        // fd id.
        let vm = KtstrKvm::new(test_topo(1), 64, false).unwrap();
        let fd = get_stats_fd(&vm.vcpus[0]).expect("KVM_GET_STATS_FD on vCPU fd must return Some");
        let buf = read_initial(fd.as_raw_fd()).expect("initial read of vCPU stats fd");
        let meta = parse_meta_from_buf(&buf).expect("parse vCPU stats metadata");
        assert!(
            meta.descs.iter().any(|d| d.name == "exits"),
            "vCPU stats descriptors must include the 'exits' counter",
        );
    }

    #[test]
    fn parse_meta_from_vcpu() {
        let vm = KtstrKvm::new(test_topo(1), 64, false).unwrap();
        if let Some(fd) = get_stats_fd(&vm.vcpus[0]) {
            let buf = read_initial(fd.as_raw_fd()).unwrap();
            let meta = parse_meta_from_buf(&buf).unwrap();
            assert!(!meta.descs.is_empty());
            assert!(meta.descs.iter().any(|d| d.name == "exits"));
        }
    }

    #[test]
    fn pread_data_after_initial_read() {
        let vm = KtstrKvm::new(test_topo(1), 64, false).unwrap();
        if let Some(fd) = get_stats_fd(&vm.vcpus[0]) {
            let buf = read_initial(fd.as_raw_fd()).unwrap();
            let meta = parse_meta_from_buf(&buf).unwrap();
            let data = pread_data(fd.as_raw_fd(), &meta);
            assert!(data.is_some(), "pread should work after initial read");
            assert!(!data.unwrap().is_empty());
        }
    }

    #[test]
    fn vcpu_reader_opens_and_reads() {
        let vm = KtstrKvm::new(test_topo(2), 64, false).unwrap();
        if let Some(reader) = VcpuStatsReader::open(&vm.vcpus[0]) {
            let snap = reader.read_snapshot();
            assert!(snap.contains_key("exits"), "should contain 'exits'");
        }
    }

    #[test]
    fn stats_context_opens_and_reads() {
        let vm = KtstrKvm::new(test_topo(2), 64, false).unwrap();
        if let Some(ctx) = open_stats_context(&vm.vcpus) {
            assert_eq!(ctx.readers.len(), 2);
            let stats = ctx.read_stats();
            assert_eq!(stats.per_vcpu.len(), 2);
        }
    }

    #[test]
    fn interesting_stats_are_trust_relevant() {
        use crate::vmm::KVM_INTERESTING_STATS;
        assert!(KVM_INTERESTING_STATS.contains(&"exits"));
        assert!(KVM_INTERESTING_STATS.contains(&"halt_exits"));
        assert!(KVM_INTERESTING_STATS.contains(&"halt_successful_poll"));
        assert!(KVM_INTERESTING_STATS.contains(&"halt_attempted_poll"));
        assert!(KVM_INTERESTING_STATS.contains(&"halt_wait_ns"));
        assert!(KVM_INTERESTING_STATS.contains(&"signal_exits"));
        assert!(KVM_INTERESTING_STATS.contains(&"hypercalls"));
        assert!(KVM_INTERESTING_STATS.contains(&"preemption_reported"));
    }

    #[test]
    fn kvm_stats_totals_serde_roundtrip() {
        use crate::vmm::KvmStatsTotals;
        let mut totals = KvmStatsTotals::default();
        let mut m = HashMap::new();
        m.insert("exits".to_string(), 1000u64);
        m.insert("halt_exits".to_string(), 200u64);
        totals.per_vcpu = vec![m];
        let json = serde_json::to_string(&totals).unwrap();
        let loaded: KvmStatsTotals = serde_json::from_str(&json).unwrap();
        assert_eq!(loaded.per_vcpu.len(), 1);
        assert_eq!(loaded.avg("exits"), 1000);
        assert_eq!(loaded.avg("halt_exits"), 200);
    }

    #[test]
    fn sum_and_avg() {
        use crate::vmm::KvmStatsTotals;
        let mut d = KvmStatsTotals::default();
        let mut m1 = HashMap::new();
        m1.insert("exits".to_string(), 100u64);
        let mut m2 = HashMap::new();
        m2.insert("exits".to_string(), 200u64);
        d.per_vcpu = vec![m1, m2];
        assert_eq!(d.sum("exits"), 300);
        assert_eq!(d.avg("exits"), 150);
        assert_eq!(d.sum("nonexistent"), 0);
        assert_eq!(d.avg("nonexistent"), 0);
    }

    #[test]
    fn avg_empty() {
        use crate::vmm::KvmStatsTotals;
        let d = KvmStatsTotals::default();
        assert_eq!(d.avg("exits"), 0);
    }
}
