//! C2a — the supervisor-owned guest process tree for the native-execution VM
//! (Stage 1; see `docs/design/native-execution-vm-stage1-spec-2026-06-15.md`).
//!
//! The sentry's only process-identity state today is `ring.pid` (a host pid per
//! ring slot) and `MAIN_CELL_PID`. Host pids are fine while the sandbox is live,
//! but they are **un-serializable identity**: a warm restore re-creates the guest
//! processes and the kernel hands them *new* host pids, so anything keyed by host
//! pid breaks across a restore. This module introduces the owned, stable identity
//! the snapshot needs — a **virtual pid** ([`Vpid`]) the supervisor allocates and
//! owns — and the bidirectional `vpid ↔ host_pid` mapping that lets the rest of
//! the supervisor keep using host pids at runtime while the *serialized* identity
//! is always a vpid.
//!
//! ## Additive by design (C2a)
//!
//! This is purely additive: it **shadows** the process identity alongside the
//! existing host-pid-keyed `fdt()`/`cwds()` tables, hooking the very same
//! lifecycle points (`fd_snapshot` at fork, `fd_adopt` at the child's first
//! request, `fd_drop` at death, plus the main-cell seed in `supervisor_main`). It
//! changes **no** runtime behaviour and does not re-key any existing map — so it
//! cannot regress the live fork/clone/exec paths. Making the *guest* observe
//! vpids (virtualizing `getpid`/`getppid`/`wait4`/`kill` returns+args) is the
//! follow-on step "C2b"; capturing thread register files into each record is C3/C4.
//!
//! ## Lifecycle (mirrors `pending_cwd` → `cwds`)
//!
//! ```text
//!   supervisor_main(host_pid)         → register_main()    : vpid 1 = sandbox init
//!   fork/clone (CTL_FORK_TABLE, slot) → alloc_for_slot()   : child vpid, parent linked, host pid unknown
//!   child's first request (fd_adopt)  → bind_slot()        : slot's vpid ← child host pid, now Running
//!   fork cancelled (CTL_FORK_CANCEL)  → cancel_slot()      : drop the born-but-unborn vpid
//!   death (CTL_REAP / fd_drop)        → mark_dead()         : remove from the live tree
//! ```

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

use super::state_snap::{GuestProcRecord, ProcState, ProcTreeSnapshot, Vpid};

/// vpid 1 is the sandbox's init (the main cell), mirroring a PID namespace's init.
const INIT_VPID: Vpid = 1;

/// One owned guest process: its stable virtual identity + topology + lifecycle.
/// The `host_pid` is a private runtime detail (re-assigned on every restore) and
/// is **never** the serialized identity — only `vpid`/`parent_vpid` cross the wire.
#[derive(Clone, Debug, PartialEq, Eq)]
struct GuestProc {
    vpid: Vpid,
    parent_vpid: Vpid,
    /// The live host pid, or 0 while the process is born-but-not-yet-bound (between
    /// the parent's pre-fork `alloc_for_slot` and the child's first `bind_slot`).
    host_pid: i32,
    state: ProcState,
}

/// The supervisor's system of record for the guest process tree, keyed by the
/// owned [`Vpid`]. A `Mutex` global mirroring `fdt()`/`cwds()`.
pub struct ProcTable {
    /// The live tree: every allocated, non-removed process by its vpid.
    by_vpid: HashMap<Vpid, GuestProc>,
    /// Reverse index for the hot translation `host_pid → vpid` (only live,
    /// bound processes; a vpid with `host_pid == 0` is not indexed here).
    by_host: HashMap<i32, Vpid>,
    /// Children born at a fork but not yet bound to a host pid, keyed by the ring
    /// slot the child will adopt (parallel to `pending()` / `pending_cwd()`).
    pending_by_slot: HashMap<u64, GuestProc>,
    /// Monotonic vpid allocator. Never reused within a sandbox lifetime, so a
    /// stale reference to a dead vpid can never silently alias a live process.
    next_vpid: Vpid,
}

impl ProcTable {
    fn new() -> ProcTable {
        ProcTable {
            by_vpid: HashMap::new(),
            by_host: HashMap::new(),
            pending_by_slot: HashMap::new(),
            next_vpid: INIT_VPID,
        }
    }

    fn alloc_vpid(&mut self) -> Vpid {
        let v = self.next_vpid;
        // vpids are a u32 namespace; exhausting it would take 4 billion process
        // creations in one sandbox lifetime. Saturate rather than wrap (wrapping
        // could alias a live vpid); a saturated allocator is a hard error the
        // capture path would surface, never a silent identity collision.
        self.next_vpid = self.next_vpid.saturating_add(1);
        v
    }
}

/// The process table singleton (mirrors `fdt()` / `cwds()`).
pub fn table() -> &'static Mutex<ProcTable> {
    static T: OnceLock<Mutex<ProcTable>> = OnceLock::new();
    T.get_or_init(|| Mutex::new(ProcTable::new()))
}

// ─── lifecycle hooks ─────────────────────────────────────────────────────────

/// Register the main cell (sandbox init) as [`INIT_VPID`] bound to `host_pid`.
/// Idempotent on the vpid (a re-seed rebinds init to the new host pid — used when
/// the main cell is (re)started), but never re-allocates vpid 1.
pub fn register_main(host_pid: i32) -> Vpid {
    let mut t = table().lock().unwrap();
    // Clear any prior host-pid binding for init (re-seed case).
    if let Some(old) = t.by_vpid.get(&INIT_VPID).map(|p| p.host_pid) {
        if old != 0 {
            t.by_host.remove(&old);
        }
    }
    // Keep the allocator ahead of INIT_VPID.
    if t.next_vpid <= INIT_VPID {
        t.next_vpid = INIT_VPID + 1;
    }
    t.by_vpid.insert(
        INIT_VPID,
        GuestProc {
            vpid: INIT_VPID,
            parent_vpid: 0,
            host_pid,
            state: ProcState::Running,
        },
    );
    if host_pid != 0 {
        t.by_host.insert(host_pid, INIT_VPID);
    }
    INIT_VPID
}

/// Rebuild live vpid↔host bindings from a restored snapshot.
///
/// `host_bindings` maps snapshot vpids to the freshly-created host pids that now
/// run those guest processes. Only vpids with a host binding are inserted into
/// the live table; the allocator is still advanced past every vpid recorded in
/// the snapshot so subsequent forks cannot collide with restored identity.
pub fn restore_bindings(snapshot: &ProcTreeSnapshot, host_bindings: &[(Vpid, i32)]) {
    let mut t = table().lock().unwrap();
    restore_bindings_into(&mut t, snapshot, host_bindings);
}

fn restore_bindings_into(
    t: &mut ProcTable,
    snapshot: &ProcTreeSnapshot,
    host_bindings: &[(Vpid, i32)],
) {
    t.by_vpid.clear();
    t.by_host.clear();
    t.pending_by_slot.clear();

    let mut max_vpid = INIT_VPID;
    let host_by_vpid: HashMap<Vpid, i32> = host_bindings
        .iter()
        .copied()
        .filter(|(_, host_pid)| *host_pid > 0)
        .collect();

    for rec in &snapshot.procs {
        max_vpid = max_vpid.max(rec.vpid);
        let Some(&host_pid) = host_by_vpid.get(&rec.vpid) else {
            continue;
        };
        t.by_vpid.insert(
            rec.vpid,
            GuestProc {
                vpid: rec.vpid,
                parent_vpid: rec.parent_vpid,
                host_pid,
                state: rec.state,
            },
        );
        t.by_host.insert(host_pid, rec.vpid);
    }
    t.next_vpid = max_vpid.saturating_add(1);
}

/// At fork/clone (CTL_FORK_TABLE, pre-fork, in the PARENT's context): allocate the
/// child's vpid keyed by its ring `slot`, parented to the forking process. The
/// child's host pid isn't known yet — it is filled in by [`bind_slot`] when the
/// child first delegates. If a stale pending entry already sits on `slot` (a slot
/// recycled without a clean cancel/bind), it is replaced — exactly as
/// `fd_snapshot` overwrites a stale fd snapshot on slot reuse.
pub fn alloc_for_slot(slot: u64, parent_host_pid: i32) {
    let mut t = table().lock().unwrap();
    let parent_vpid = t.by_host.get(&parent_host_pid).copied().unwrap_or(0);
    let vpid = t.alloc_vpid();
    t.pending_by_slot.insert(
        slot,
        GuestProc {
            vpid,
            parent_vpid,
            host_pid: 0,
            state: ProcState::Running,
        },
    );
}

/// The child's first delegated request (fd_adopt): its host pid is now known —
/// promote the slot's pending vpid into the live tree bound to `child_host_pid`.
/// A no-op if no pending entry exists for `slot` (e.g. a child that never went
/// through `alloc_for_slot`, like the zygote/restore paths handled elsewhere).
pub fn bind_slot(slot: u64, child_host_pid: i32) {
    let mut t = table().lock().unwrap();
    if let Some(mut p) = t.pending_by_slot.remove(&slot) {
        p.host_pid = child_host_pid;
        let vpid = p.vpid;
        t.by_vpid.insert(vpid, p);
        t.by_host.insert(child_host_pid, vpid);
    }
}

/// A fork that was cancelled (CTL_FORK_CANCEL): drop the slot's born-but-unborn
/// vpid so it never lingers as a phantom child.
pub fn cancel_slot(slot: u64) {
    table().lock().unwrap().pending_by_slot.remove(&slot);
}

/// A process died (CTL_REAP / `fd_drop`): remove it from the live tree by host
/// pid. Idempotent — CTL_REAP can arrive both from the process's own `exit_group`
/// and from its parent's `wait4`, and teardown paths also call `fd_drop`.
pub fn mark_dead(host_pid: i32) {
    let mut t = table().lock().unwrap();
    if let Some(vpid) = t.by_host.remove(&host_pid) {
        t.by_vpid.remove(&vpid);
    }
}

/// Translate a live host pid to its owned vpid (`None` if not a tracked process).
pub fn vpid_for(host_pid: i32) -> Option<Vpid> {
    table().lock().unwrap().by_host.get(&host_pid).copied()
}

/// Translate an owned vpid to its current live host pid (`None` if unknown/dead or
/// not yet bound). This is the seam restore re-points at the freshly-cloned hosts.
pub fn host_pid_for(vpid: Vpid) -> Option<i32> {
    table()
        .lock()
        .unwrap()
        .by_vpid
        .get(&vpid)
        .map(|p| p.host_pid)
        .filter(|&h| h != 0)
}

/// DIAGNOSTIC: snapshot the WHOLE proctree as `(vpid, parent_vpid, host_pid)` —
/// includes not-yet-bound (host_pid 0) and zombie entries still in the table. Used
/// by the hang-dump watchdog to reveal parent→child topology (e.g. which guest owns
/// a dead child, and whether a child-exit SIGCHLD reached the right parent).
pub fn dump_all() -> Vec<(i64, i64, i64)> {
    let t = table().lock().unwrap();
    let mut v: Vec<(i64, i64, i64)> = t
        .by_vpid
        .values()
        .map(|p| (p.vpid as i64, p.parent_vpid as i64, p.host_pid as i64))
        .collect();
    v.sort_unstable();
    v
}

/// Every live, bound host pid in the sandbox's process tree — the set a guest
/// `readdir("/proc")` enumerates (sorted ascending). Pids aren't virtualized yet
/// (the guest sees host pids), so these ARE the guest-visible pids; once C2b-pid
/// virtualization lands this returns vpids instead.
pub fn live_host_pids() -> Vec<i32> {
    let t = table().lock().unwrap();
    let mut v: Vec<i32> = t.by_host.keys().copied().collect();
    v.sort_unstable();
    v
}

/// Live bound processes in the subtree rooted at `root_host_pid`, returned as
/// `(vpid, parent_vpid, host_pid)` sorted by vpid with the root first. This is
/// the C4 live-capture target set: host pids are runtime-only, while vpids carry
/// the topology that crosses the snapshot boundary.
pub fn live_subtree(root_host_pid: i32) -> Vec<(Vpid, Vpid, i32)> {
    let t = table().lock().unwrap();
    live_subtree_from_table(&t, root_host_pid)
}

fn live_subtree_from_table(t: &ProcTable, root_host_pid: i32) -> Vec<(Vpid, Vpid, i32)> {
    let Some(root_vpid) = t.by_host.get(&root_host_pid).copied() else {
        return Vec::new();
    };
    let mut out = Vec::new();
    for p in t.by_vpid.values() {
        if p.host_pid == 0 {
            continue;
        }
        let mut cur = Some(p.vpid);
        while let Some(vpid) = cur {
            if vpid == root_vpid {
                out.push((p.vpid, p.parent_vpid, p.host_pid));
                break;
            }
            cur = t.by_vpid.get(&vpid).and_then(|proc| {
                if proc.parent_vpid == 0 {
                    None
                } else {
                    Some(proc.parent_vpid)
                }
            });
        }
    }
    out.sort_by_key(|(vpid, _, _)| (*vpid != root_vpid, *vpid));
    out
}

/// Whether the process occupying `slot` belongs to one of the protected host-pid
/// roots. Handles both a fully-bound child and the short pre-first-syscall window
/// where its process record still lives in `pending_by_slot`.
pub fn slot_descends_from_any(slot: u64, host_pid: i32, roots: &[i32]) -> bool {
    let t = table().lock().unwrap();
    let root_vpids: Vec<Vpid> = roots
        .iter()
        .filter_map(|p| t.by_host.get(p).copied())
        .collect();
    if root_vpids.is_empty() {
        return false;
    }
    let mut current = t
        .by_host
        .get(&host_pid)
        .copied()
        .or_else(|| t.pending_by_slot.get(&slot).map(|p| p.vpid));
    while let Some(vpid) = current {
        if root_vpids.contains(&vpid) {
            return true;
        }
        current = t
            .by_vpid
            .get(&vpid)
            .or_else(|| t.pending_by_slot.values().find(|p| p.vpid == vpid))
            .map(|p| p.parent_vpid)
            .filter(|p| *p != 0);
    }
    false
}

/// Serialize the live process tree into the wire [`ProcTreeSnapshot`].
///
/// `query_pgrp(host_pid) -> (pgid, sid)` is supplied by the caller (the
/// supervisor, which owns the raw-syscall layer) so this module stays free of host
/// syscalls and is unit-testable. Records carry **no host pid** (only vpid
/// identity). `threads` is left empty here — the capture orchestrator (C4) fills
/// each record's register files from the SENTINEL register capture (C3), keyed by
/// vpid. Born-but-unbound children (mid-fork) are excluded; capture runs only at a
/// quiescent checkpoint where no fork is in flight.
pub fn capture(query_pgrp: impl Fn(i32) -> (u32, u32)) -> ProcTreeSnapshot {
    let t = table().lock().unwrap();
    let mut procs: Vec<GuestProcRecord> = t
        .by_vpid
        .values()
        .filter(|p| p.host_pid != 0)
        .map(|p| {
            let (pgid, sid) = query_pgrp(p.host_pid);
            GuestProcRecord {
                vpid: p.vpid,
                parent_vpid: p.parent_vpid,
                pgid,
                sid,
                threads: Vec::new(),
                state: p.state,
            }
        })
        .collect();
    // Deterministic order (HashMap iteration is not) so a captured snapshot is
    // byte-stable and diff-able against a control run.
    procs.sort_by_key(|r| r.vpid);
    ProcTreeSnapshot { procs }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Each test drives a FRESH ProcTable instance directly (not the global
    // singleton) so the tests are independent and order-free.
    fn fresh() -> ProcTable {
        ProcTable::new()
    }

    // Re-implement the hook bodies against an explicit table so we can unit-test
    // the topology logic without the process-global singleton. (The pub fns above
    // are thin lock-and-delegate wrappers over exactly this logic.)
    fn t_register_main(t: &mut ProcTable, host_pid: i32) -> Vpid {
        if t.next_vpid <= INIT_VPID {
            t.next_vpid = INIT_VPID + 1;
        }
        t.by_vpid.insert(
            INIT_VPID,
            GuestProc {
                vpid: INIT_VPID,
                parent_vpid: 0,
                host_pid,
                state: ProcState::Running,
            },
        );
        t.by_host.insert(host_pid, INIT_VPID);
        INIT_VPID
    }
    fn t_alloc(t: &mut ProcTable, slot: u64, parent_host: i32) -> Vpid {
        let parent_vpid = t.by_host.get(&parent_host).copied().unwrap_or(0);
        let vpid = t.alloc_vpid();
        t.pending_by_slot.insert(
            slot,
            GuestProc {
                vpid,
                parent_vpid,
                host_pid: 0,
                state: ProcState::Running,
            },
        );
        vpid
    }
    fn t_bind(t: &mut ProcTable, slot: u64, child_host: i32) {
        if let Some(mut p) = t.pending_by_slot.remove(&slot) {
            p.host_pid = child_host;
            let v = p.vpid;
            t.by_vpid.insert(v, p);
            t.by_host.insert(child_host, v);
        }
    }
    fn t_dead(t: &mut ProcTable, host_pid: i32) {
        if let Some(v) = t.by_host.remove(&host_pid) {
            t.by_vpid.remove(&v);
        }
    }
    fn t_capture(t: &ProcTable, q: impl Fn(i32) -> (u32, u32)) -> ProcTreeSnapshot {
        let mut procs: Vec<GuestProcRecord> = t
            .by_vpid
            .values()
            .filter(|p| p.host_pid != 0)
            .map(|p| {
                let (pgid, sid) = q(p.host_pid);
                GuestProcRecord {
                    vpid: p.vpid,
                    parent_vpid: p.parent_vpid,
                    pgid,
                    sid,
                    threads: Vec::new(),
                    state: p.state,
                }
            })
            .collect();
        procs.sort_by_key(|r| r.vpid);
        ProcTreeSnapshot { procs }
    }

    fn sample_snapshot() -> ProcTreeSnapshot {
        ProcTreeSnapshot {
            procs: vec![
                GuestProcRecord {
                    vpid: 1,
                    parent_vpid: 0,
                    pgid: 1,
                    sid: 1,
                    threads: Vec::new(),
                    state: ProcState::Running,
                },
                GuestProcRecord {
                    vpid: 2,
                    parent_vpid: 1,
                    pgid: 1,
                    sid: 1,
                    threads: Vec::new(),
                    state: ProcState::Running,
                },
                GuestProcRecord {
                    vpid: 4,
                    parent_vpid: 2,
                    pgid: 1,
                    sid: 1,
                    threads: Vec::new(),
                    state: ProcState::Running,
                },
            ],
        }
    }

    #[test]
    fn main_is_vpid_one() {
        let mut t = fresh();
        assert_eq!(t_register_main(&mut t, 1000), INIT_VPID);
        assert_eq!(t.by_host.get(&1000), Some(&1));
    }

    #[test]
    fn fork_allocs_child_with_parent_linkage_and_binds_host_pid() {
        let mut t = fresh();
        t_register_main(&mut t, 1000); // init = vpid 1, host 1000
        let child = t_alloc(&mut t, /*slot*/ 7, /*parent_host*/ 1000);
        assert_eq!(child, 2, "second process gets vpid 2");
        // Not yet live — pending, host pid unknown.
        assert!(t.by_host.get(&2000).is_none());
        t_bind(&mut t, 7, 2000);
        assert_eq!(t.by_host.get(&2000), Some(&2));
        let rec = &t.by_vpid[&2];
        assert_eq!(rec.parent_vpid, 1, "child's parent is init");
        assert_eq!(rec.host_pid, 2000);
    }

    #[test]
    fn grandchild_topology_is_recorded() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        t_alloc(&mut t, 7, 1000);
        t_bind(&mut t, 7, 2000); // vpid 2, parent 1
        t_alloc(&mut t, 8, 2000); // child of vpid 2
        t_bind(&mut t, 8, 3000); // vpid 3, parent 2
        assert_eq!(t.by_vpid[&3].parent_vpid, 2);
    }

    #[test]
    fn death_removes_from_live_tree_and_is_idempotent() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        t_alloc(&mut t, 7, 1000);
        t_bind(&mut t, 7, 2000);
        t_dead(&mut t, 2000);
        assert!(t.by_host.get(&2000).is_none());
        assert!(!t.by_vpid.contains_key(&2));
        t_dead(&mut t, 2000); // idempotent: second reap is a no-op, not a panic
        assert!(!t.by_vpid.contains_key(&2));
    }

    #[test]
    fn cancel_drops_pending_and_realloc_does_not_collide() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        let v1 = t_alloc(&mut t, 7, 1000);
        t.pending_by_slot.remove(&7); // cancel
        assert!(t.pending_by_slot.is_empty());
        let v2 = t_alloc(&mut t, 7, 1000); // slot reused
        assert!(v2 > v1, "vpids are monotonic, never reused");
        t_bind(&mut t, 7, 2000);
        assert_eq!(t.by_vpid[&v2].host_pid, 2000);
        assert!(
            !t.by_vpid.contains_key(&v1),
            "cancelled vpid never went live"
        );
    }

    #[test]
    fn capture_excludes_unbound_and_is_vpid_ordered() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        t_alloc(&mut t, 7, 1000);
        t_bind(&mut t, 7, 2000);
        t_alloc(&mut t, 8, 1000); // born but NOT bound — must be excluded
        let snap = t_capture(&t, |host| (host as u32, 1)); // pgid=host_pid, sid=1 (stub)
        assert_eq!(
            snap.procs.len(),
            2,
            "only the 2 bound procs (unbound excluded)"
        );
        assert_eq!(snap.procs[0].vpid, 1);
        assert_eq!(snap.procs[1].vpid, 2);
        assert_eq!(
            snap.procs[0].pgid, 1000,
            "pgid filled from query_pgrp(host)"
        );
        assert_eq!(snap.procs[1].parent_vpid, 1);
        // No host pid leaks into the serialized record.
        for r in &snap.procs {
            assert!(r.threads.is_empty(), "threads filled later by C3/C4");
        }
    }

    #[test]
    fn live_subtree_returns_root_and_descendants_only() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        t_alloc(&mut t, 7, 1000);
        t_bind(&mut t, 7, 2000); // vpid 2
        t_alloc(&mut t, 8, 2000);
        t_bind(&mut t, 8, 3000); // vpid 3
        t_alloc(&mut t, 9, 1000);
        t_bind(&mut t, 9, 4000); // vpid 4, sibling of root target

        let got = live_subtree_from_table(&t, 2000);

        assert_eq!(got, vec![(2, 1, 2000), (3, 2, 3000)]);
    }

    #[test]
    fn restore_bindings_rekeys_snapshot_vpids_to_new_host_pids() {
        let mut t = fresh();
        t_register_main(&mut t, 1000);
        t_alloc(&mut t, 7, 1000);
        t_bind(&mut t, 7, 2000);

        restore_bindings_into(&mut t, &sample_snapshot(), &[(1, 9001), (4, 9004)]);

        assert_eq!(t.by_host.get(&9001), Some(&1));
        assert_eq!(t.by_host.get(&9004), Some(&4));
        assert!(!t.by_host.contains_key(&1000));
        assert!(!t.by_host.contains_key(&2000));
        assert_eq!(t.by_vpid[&4].parent_vpid, 2);
        assert!(
            !t.by_vpid.contains_key(&2),
            "unlaunched restored children are not inserted as live processes"
        );

        let v = t_alloc(&mut t, 8, 9001);
        assert_eq!(v, 5, "allocator advances past snapshot max vpid");
    }
}
