// DAX session — orchestrates SETUPMAPPING/REMOVEMAPPING by combining a
// FUSE filesystem backend (which knows how to mmap a file region into
// the host process) with an HVF mapper (which knows how to plumb that
// host VA into the guest's DAX window via `hv_vm_map`).
//
// This module is deliberately VM-substrate-agnostic: it takes an
// `Arc<dyn HvfMapper>` so tests can drive the full session machinery
// without a real VM. The real impl lives in `crate::hvf::dax_mapper`
// (added in a subsequent slice).

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;

use super::backend::{Errno, FsBackend};
use super::protocol::{
    RemoveMappingIn, RemoveMappingOne, SetupMappingIn, FUSE_SETUPMAPPING_FLAG_READ,
    FUSE_SETUPMAPPING_FLAG_WRITE,
};

/// Permission bits we pass through to the underlying HVF mapping.
/// We deliberately keep the bit shape stable across HVF / KVM so the
/// trait is portable.
pub const DAX_PROT_READ: u32 = 1 << 0;
pub const DAX_PROT_WRITE: u32 = 1 << 1;

/// Abstracts `hv_vm_map` / `hv_vm_unmap` so the DAX session can be
/// tested without a real HVF VM. The real impl wraps applevisor-sys
/// in `crate::hvf` and is wired into the worker once we plumb the
/// device into the VMM device list.
pub trait HvfMapper: Send + Sync {
    /// Map `len` bytes from `host_va` into guest physical at
    /// `gpa` with `prot` (DAX_PROT_*).
    fn map(&self, host_va: *mut u8, gpa: u64, len: u64, prot: u32) -> Result<(), Errno>;

    /// Unmap a previously-established region. `gpa` + `len` must be
    /// 16 KiB-aligned and match an existing mapping; HVF tolerates
    /// unmapping the exact range we mapped.
    fn unmap(&self, gpa: u64, len: u64) -> Result<(), Errno>;
}

/// A no-op mapper used by tests + the not-yet-wired-up production
/// path. Records every map/unmap call so tests can inspect them.
pub struct MockHvfMapper {
    log: Mutex<Vec<MockCall>>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MockCall {
    Map { host_va: usize, gpa: u64, len: u64, prot: u32 },
    Unmap { gpa: u64, len: u64 },
}

impl MockHvfMapper {
    pub fn new() -> Self {
        Self {
            log: Mutex::new(Vec::new()),
        }
    }
    pub fn calls(&self) -> Vec<MockCall> {
        self.log.lock().unwrap().clone()
    }
}

impl Default for MockHvfMapper {
    fn default() -> Self {
        Self::new()
    }
}

impl HvfMapper for MockHvfMapper {
    fn map(&self, host_va: *mut u8, gpa: u64, len: u64, prot: u32) -> Result<(), Errno> {
        self.log.lock().unwrap().push(MockCall::Map {
            host_va: host_va as usize,
            gpa,
            len,
            prot,
        });
        Ok(())
    }
    fn unmap(&self, gpa: u64, len: u64) -> Result<(), Errno> {
        self.log.lock().unwrap().push(MockCall::Unmap { gpa, len });
        Ok(())
    }
}

/// One active DAX mapping. Recorded by `DaxSession` so REMOVEMAPPING
/// can find the right host VA + len to tear down.
#[derive(Clone, Debug)]
struct Slot {
    host_va: *mut u8,
    len: u64,
}

// SAFETY: host_va is a process-local pointer used only by the
// owning DaxSession (which itself is Send+Sync). We don't dereference
// it across threads — only HvfMapper does the actual map/unmap.
unsafe impl Send for Slot {}
unsafe impl Sync for Slot {}

/// Per-mount DAX state. Owned by the FuseServer when a DAX mount is
/// configured.
pub struct DaxSession {
    /// Guest physical base of the DAX window. All FUSE moffset values
    /// are relative to this.
    dax_base_gpa: u64,
    /// Window length. SETUPMAPPING with moffset+len > window_len is rejected.
    dax_window_len: u64,
    /// Aggregate cap on the sum of currently-mapped byte counts.
    /// SETUPMAPPING with `currently_mapped + req.len > cap` returns
    /// ENOSPC. Defends against a hostile guest churning
    /// map/remap/leak cycles to consume host stage-2 page-table
    /// memory. Default = `dax_window_len` (the natural upper bound).
    cap_bytes: u64,
    /// Running sum of `Slot.len` across all active entries. Atomic so
    /// concurrent SETUPMAPPING calls don't race.
    current_mapped_bytes: AtomicU64,
    /// Backend that knows how to mmap file regions into host VA space.
    backend: std::sync::Arc<dyn FsBackend>,
    /// HVF mapper.
    mapper: std::sync::Arc<dyn HvfMapper>,
    /// Active mappings, keyed by moffset (window-relative).
    active: Mutex<BTreeMap<u64, Slot>>,
}

impl DaxSession {
    pub fn new(
        dax_base_gpa: u64,
        dax_window_len: u64,
        backend: std::sync::Arc<dyn FsBackend>,
        mapper: std::sync::Arc<dyn HvfMapper>,
    ) -> Self {
        Self::with_cap(dax_base_gpa, dax_window_len, dax_window_len, backend, mapper)
    }

    /// Construct with an explicit aggregate cap. The cap bounds the
    /// SUM of currently-active mapping lengths. Pass a smaller value
    /// than the window to tighten the limit for hostile-tenant
    /// deployments — e.g. cap at 2 GiB while the window is 8 GiB.
    pub fn with_cap(
        dax_base_gpa: u64,
        dax_window_len: u64,
        cap_bytes: u64,
        backend: std::sync::Arc<dyn FsBackend>,
        mapper: std::sync::Arc<dyn HvfMapper>,
    ) -> Self {
        Self {
            dax_base_gpa,
            dax_window_len,
            cap_bytes,
            current_mapped_bytes: AtomicU64::new(0),
            backend,
            mapper,
            active: Mutex::new(BTreeMap::new()),
        }
    }

    /// Handle FUSE_SETUPMAPPING. Returns Ok(()) on success, negated
    /// errno on failure.
    pub fn setup(&self, nodeid: u64, req: &SetupMappingIn) -> Result<(), Errno> {
        log_dax(|| format!(
            "SETUPMAPPING: nodeid={nodeid} fh={} foffset={:#x} len={:#x} flags={:#x} moffset={:#x}",
            req.fh, req.foffset, req.len, req.flags, req.moffset
        ));
        // Bounds + alignment checks. Apple Silicon page granule is
        // 16 KiB; FUSE_INIT advertised map_alignment=14 to the guest.
        if (req.moffset & 0x3FFF) != 0 || (req.len & 0x3FFF) != 0 {
            return Err(super::backend::EINVAL);
        }
        if req.moffset.checked_add(req.len).map_or(true, |end| end > self.dax_window_len) {
            return Err(super::backend::EINVAL);
        }
        // Aggregate-bytes cap. Pre-reserve the increment with
        // compare-exchange so concurrent SETUPMAPPINGs can't both
        // observe room then race past the cap. On reject the cap
        // returns ENOSPC; the guest's virtio-fs driver retries
        // after dropping references (REMOVEMAPPING).
        let mut current = self.current_mapped_bytes.load(Ordering::Acquire);
        loop {
            let proposed = match current.checked_add(req.len) {
                Some(v) if v <= self.cap_bytes => v,
                _ => return Err(super::backend::ENOSPC),
            };
            match self.current_mapped_bytes.compare_exchange_weak(
                current,
                proposed,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => break,
                Err(observed) => current = observed,
            }
        }
        // Permissions: virtio-fs sets at least one of READ / WRITE.
        let mut prot = 0u32;
        if req.flags & FUSE_SETUPMAPPING_FLAG_READ != 0 {
            prot |= DAX_PROT_READ;
        }
        if req.flags & FUSE_SETUPMAPPING_FLAG_WRITE != 0 {
            prot |= DAX_PROT_WRITE;
        }
        if prot == 0 {
            // Release the cap reservation we just took.
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            return Err(super::backend::EINVAL);
        }

        // Ask the backend for a host VA covering the file region.
        let host_va = match self
            .backend
            .dax_map(nodeid, req.fh, req.foffset, req.len, prot)
        {
            Ok(va) => va,
            Err(e) => {
                log_dax(|| format!("backend.dax_map FAILED: errno={e}"));
                self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
                return Err(e);
            }
        };

        // Plumb into the guest.
        let gpa = self.dax_base_gpa + req.moffset;
        log_dax(|| format!(
            "hv_vm_map: host_va={:#x} gpa={gpa:#x} len={:#x} prot={prot:#x}",
            host_va as usize, req.len
        ));
        self.mapper.map(host_va, gpa, req.len, prot).map_err(|e| {
            log_dax(|| format!("mapper.map FAILED: errno={e}"));
            // Failed to map into the guest — release the backend's
            // mapping so we don't leak.
            let _ = self
                .backend
                .dax_unmap(nodeid, host_va, req.len);
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            e
        })?;

        // Record.
        let mut active = self.active.lock().unwrap();
        if active.contains_key(&req.moffset) {
            // Overlap with an existing slot. Refuse — the guest is
            // expected to REMOVEMAPPING before re-mapping the same
            // moffset.
            let _ = self.mapper.unmap(gpa, req.len);
            let _ = self.backend.dax_unmap(nodeid, host_va, req.len);
            self.current_mapped_bytes.fetch_sub(req.len, Ordering::AcqRel);
            return Err(super::backend::EEXIST);
        }
        active.insert(req.moffset, Slot { host_va, len: req.len });
        Ok(())
    }

    /// Handle FUSE_REMOVEMAPPING. The payload contains a count followed
    /// by `count` RemoveMappingOne entries; the FuseServer parses the
    /// count and feeds them to us one by one.
    pub fn remove(&self, nodeid: u64, entry: &RemoveMappingOne) -> Result<(), Errno> {
        let mut active = self.active.lock().unwrap();
        let slot = active.remove(&entry.moffset).ok_or(super::backend::EINVAL)?;
        if slot.len != entry.len {
            // The kernel asked us to unmap a different len than we
            // mapped. Refuse — re-insert the slot.
            active.insert(entry.moffset, slot);
            return Err(super::backend::EINVAL);
        }
        let gpa = self.dax_base_gpa + entry.moffset;
        let released_len = slot.len;
        drop(active);
        // Best-effort: even if either step fails, we've removed the
        // slot from `active` so the next SETUPMAPPING at the same
        // moffset can proceed. Release the cap reservation
        // unconditionally — the slot is gone regardless.
        let mapper_res = self.mapper.unmap(gpa, entry.len);
        let backend_res = self.backend.dax_unmap(nodeid, slot.host_va, slot.len);
        self.current_mapped_bytes
            .fetch_sub(released_len, Ordering::AcqRel);
        mapper_res.and(backend_res)
    }

    /// Bytes currently consumed by active mappings. Used by tests +
    /// observability. Cheap atomic load.
    pub fn current_mapped_bytes(&self) -> u64 {
        self.current_mapped_bytes.load(Ordering::Acquire)
    }

    /// Number of currently-active DAX slots. Used by tests + density
    /// metrics.
    pub fn active_slot_count(&self) -> usize {
        self.active.lock().unwrap().len()
    }
}

fn log_dax<F: FnOnce() -> String>(make_msg: F) {
    use std::io::Write;
    let Some(target) = std::env::var_os("SUPERMACHINE_FUSE_TRACE") else { return };
    let s = make_msg();
    let target = target.to_string_lossy().into_owned();
    if target == "1" || target == "stderr" {
        eprintln!("[dax] {s}");
        return;
    }
    if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(&target) {
        let _ = writeln!(f, "[dax] {s}");
    }
}

/// Parse a FUSE_REMOVEMAPPING payload into a list of RemoveMappingOne
/// entries. Header carries `count`; entries follow contiguously.
pub fn parse_remove_payload(payload: &[u8]) -> Result<Vec<RemoveMappingOne>, Errno> {
    if payload.len() < core::mem::size_of::<RemoveMappingIn>() {
        return Err(super::backend::EINVAL);
    }
    let hdr: RemoveMappingIn = unsafe { core::ptr::read_unaligned(payload.as_ptr() as *const RemoveMappingIn) };
    let count = hdr.count as usize;
    let one_size = core::mem::size_of::<RemoveMappingOne>();
    let needed = core::mem::size_of::<RemoveMappingIn>() + count * one_size;
    if payload.len() < needed {
        return Err(super::backend::EINVAL);
    }
    let mut out = Vec::with_capacity(count);
    let base = payload.as_ptr();
    for i in 0..count {
        let off = core::mem::size_of::<RemoveMappingIn>() + i * one_size;
        let one: RemoveMappingOne =
            unsafe { core::ptr::read_unaligned(base.add(off) as *const RemoveMappingOne) };
        out.push(one);
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fuse::backend::MemoryFs;
    use std::sync::Arc;

    /// A backend with a dax_map impl that returns a synthetic pointer
    /// so SETUPMAPPING orchestration can run end-to-end.
    struct StubDaxBackend {
        inner: MemoryFs,
        next_va: Mutex<usize>,
    }
    impl StubDaxBackend {
        fn new() -> Self {
            Self {
                inner: MemoryFs::new(),
                next_va: Mutex::new(0x40_0000_0000usize),
            }
        }
    }

    impl FsBackend for StubDaxBackend {
        fn lookup(&self, p: u64, n: &std::ffi::OsStr) -> Result<crate::fuse::backend::Entry, Errno> { self.inner.lookup(p, n) }
        fn forget(&self, n: u64, nl: u64) { self.inner.forget(n, nl) }
        fn getattr(&self, n: u64, f: Option<u64>) -> Result<crate::fuse::protocol::Attr, Errno> { self.inner.getattr(n, f) }
        fn open(&self, n: u64, f: u32) -> Result<u64, Errno> { self.inner.open(n, f) }
        fn read(&self, n: u64, h: u64, o: u64, s: u32) -> Result<Vec<u8>, Errno> { self.inner.read(n, h, o, s) }
        fn release(&self, n: u64, h: u64) -> Result<(), Errno> { self.inner.release(n, h) }
        fn opendir(&self, n: u64, f: u32) -> Result<u64, Errno> { self.inner.opendir(n, f) }
        fn readdir(&self, n: u64, h: u64, o: u64, s: u32) -> Result<Vec<crate::fuse::backend::DirEntry>, Errno> { self.inner.readdir(n, h, o, s) }
        fn releasedir(&self, n: u64, h: u64) -> Result<(), Errno> { self.inner.releasedir(n, h) }
        fn statfs(&self, n: u64) -> Result<crate::fuse::backend::StatFs, Errno> { self.inner.statfs(n) }
        fn dax_map(&self, _nodeid: u64, _fh: u64, _foffset: u64, len: u64, _prot: u32) -> Result<*mut u8, Errno> {
            let mut g = self.next_va.lock().unwrap();
            let va = *g;
            *g += len as usize;
            Ok(va as *mut u8)
        }
        fn dax_unmap(&self, _nodeid: u64, _host_va: *mut u8, _len: u64) -> Result<(), Errno> {
            Ok(())
        }
    }

    fn make_session() -> (DaxSession, Arc<MockHvfMapper>) {
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::new(0x80_0000_0000, 0x4_0000_0000, backend, mapper.clone());
        (session, mapper)
    }

    #[test]
    fn setup_records_mapping_and_calls_mapper() {
        let (sess, m) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        assert_eq!(sess.active_slot_count(), 1);
        let calls = m.calls();
        assert_eq!(calls.len(), 1);
        match &calls[0] {
            MockCall::Map { gpa, len, prot, .. } => {
                assert_eq!(*gpa, 0x80_0000_0000);
                assert_eq!(*len, 0x4000);
                assert_eq!(*prot, DAX_PROT_READ);
            }
            _ => panic!("expected Map"),
        }
    }

    #[test]
    fn remove_undoes_a_prior_setup() {
        let (sess, m) = make_session();
        let setup = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0x1_0000,
        };
        sess.setup(7, &setup).unwrap();
        let rem = RemoveMappingOne {
            moffset: 0x1_0000,
            len: 0x4000,
        };
        sess.remove(7, &rem).unwrap();
        assert_eq!(sess.active_slot_count(), 0);
        let calls = m.calls();
        assert_eq!(calls.len(), 2);
        match &calls[1] {
            MockCall::Unmap { gpa, len } => {
                assert_eq!(*gpa, 0x80_0000_0000 + 0x1_0000);
                assert_eq!(*len, 0x4000);
            }
            _ => panic!("expected Unmap"),
        }
    }

    #[test]
    fn setup_rejects_unaligned() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0x100, // not 16 KiB aligned
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_out_of_window() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            // 4 GiB window; ask for the last 16 KiB + 1 KiB → overflow
            moffset: 0x4_0000_0000 - 0x4000 + 0x4000,
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_overlapping_existing_moffset() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        sess.setup(7, &req).unwrap();
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EEXIST);
    }

    #[test]
    fn setup_rejects_zero_perms() {
        let (sess, _) = make_session();
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: 0, // no R, no W
            moffset: 0,
        };
        assert_eq!(sess.setup(7, &req).unwrap_err(), super::super::backend::EINVAL);
    }

    #[test]
    fn setup_rejects_when_cap_would_be_exceeded() {
        // Cap of 8 KiB (2 pages). Two 4 KiB mappings fit; third rejects.
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::with_cap(
            0x80_0000_0000,
            0x4_0000_0000,
            0x8000,  // cap = 8 KiB
            backend,
            mapper.clone(),
        );
        let mk = |moff: u64| SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: moff,
        };
        session.setup(7, &mk(0)).unwrap();
        session.setup(7, &mk(0x4000)).unwrap();
        // Third one would push current_mapped to 12 KiB; cap rejects.
        let err = session.setup(7, &mk(0x8000)).unwrap_err();
        assert_eq!(err, super::super::backend::ENOSPC);
        assert_eq!(session.current_mapped_bytes(), 0x8000);
    }

    #[test]
    fn remove_releases_cap_reservation() {
        let backend = Arc::new(StubDaxBackend::new());
        let mapper = Arc::new(MockHvfMapper::new());
        let session = DaxSession::with_cap(
            0x80_0000_0000,
            0x4_0000_0000,
            0x4000, // cap = 4 KiB; only one slot at a time
            backend,
            mapper.clone(),
        );
        let req = SetupMappingIn {
            fh: 1,
            foffset: 0,
            len: 0x4000,
            flags: FUSE_SETUPMAPPING_FLAG_READ,
            moffset: 0,
        };
        session.setup(7, &req).unwrap();
        // Second mapping is rejected because the cap is now full —
        // the cap check fires before the overlap check, so ENOSPC.
        assert_eq!(
            session.setup(7, &req.clone()).unwrap_err(),
            super::super::backend::ENOSPC
        );
        // Remove the first; cap is freed.
        let rem = RemoveMappingOne { moffset: 0, len: 0x4000 };
        session.remove(7, &rem).unwrap();
        assert_eq!(session.current_mapped_bytes(), 0);
        // Now a fresh mapping fits.
        session.setup(7, &req).unwrap();
    }

    #[test]
    fn parse_remove_payload_round_trip() {
        let entries = vec![
            RemoveMappingOne { moffset: 0x4000, len: 0x4000 },
            RemoveMappingOne { moffset: 0x8000, len: 0x8000 },
        ];
        let mut buf = Vec::new();
        let hdr = RemoveMappingIn { count: entries.len() as u32, _pad: 0 };
        buf.extend_from_slice(unsafe {
            std::slice::from_raw_parts(
                &hdr as *const RemoveMappingIn as *const u8,
                core::mem::size_of::<RemoveMappingIn>(),
            )
        });
        for e in &entries {
            buf.extend_from_slice(unsafe {
                std::slice::from_raw_parts(
                    e as *const RemoveMappingOne as *const u8,
                    core::mem::size_of::<RemoveMappingOne>(),
                )
            });
        }
        let parsed = parse_remove_payload(&buf).unwrap();
        assert_eq!(parsed.len(), 2);
        assert_eq!(parsed[0].moffset, 0x4000);
        assert_eq!(parsed[1].moffset, 0x8000);
    }
}
