//! Randomized check of the derivation tree against a model of the rules in docs/ipc.md.

extern crate std;

use std::collections::HashMap;
use std::vec::Vec;

use super::*;
use crate::{SyscallOutcome, SyscallResult};
use domain::syscall::Syscall;
use types::address::{DataRef, VirtualAddress};
use types::badge::Badge;
use types::bounded::{CapSlot, PageCount};
use types::id::ThreadId;
use types::mapping::PagePerms;
use types::permission::{Permission, PermissionSet};

const STEPS: usize = 1500;

struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        self.0
    }

    fn below(&mut self, n: usize) -> usize {
        (self.next() % n as u64) as usize
    }

    fn pick<T: Copy>(&mut self, xs: &[T]) -> Option<T> {
        (!xs.is_empty()).then(|| xs[self.below(xs.len())])
    }
}

#[derive(Clone, Copy)]
struct Node {
    parent: Option<usize>,
    is_region: bool,
}

/// The reference model: logical capabilities and where each one lives.
#[derive(Default)]
struct Model {
    nodes: Vec<Node>,
    alive: HashMap<usize, (ProcessId, CapSlot)>,
}

impl Model {
    fn add(&mut self, at: (ProcessId, CapSlot), parent: Option<usize>, is_region: bool) -> usize {
        self.nodes.push(Node { parent, is_region });
        let id = self.nodes.len() - 1;
        self.alive.insert(id, at);
        id
    }

    fn delete(&mut self, id: usize) {
        let grandparent = self.nodes[id].parent;
        for n in 0..self.nodes.len() {
            if self.alive.contains_key(&n) && self.nodes[n].parent == Some(id) {
                self.nodes[n].parent = grandparent;
            }
        }
        self.alive.remove(&id);
    }

    fn is_descendant(&self, mut n: usize, root: usize) -> bool {
        while let Some(p) = self.nodes[n].parent {
            if p == root {
                return true;
            }
            n = p;
        }
        false
    }

    fn revoke(&mut self, root: usize) {
        let doomed: Vec<usize> = self
            .alive
            .keys()
            .copied()
            .filter(|&n| self.is_descendant(n, root))
            .collect();
        for n in doomed {
            self.alive.remove(&n);
        }
    }

    /// Sorted, so a seed always means the same run (HashMap order is
    /// random per process).
    fn held_by(&self, pid: ProcessId) -> Vec<(usize, CapSlot)> {
        let mut held: Vec<(usize, CapSlot)> = self
            .alive
            .iter()
            .filter(|(_, (p, _))| *p == pid)
            .map(|(id, (_, s))| (*id, *s))
            .collect();
        held.sort_by_key(|(id, _)| *id);
        held
    }
}

fn all() -> PermissionSet {
    PermissionSet::empty()
        .with(Permission::Read)
        .with(Permission::Write)
        .with(Permission::Execute)
        .with(Permission::Delegate)
}

struct World {
    k: Kernel<MockPlatform>,
    m: Model,
    pids: [ProcessId; 3],
    tids: [ThreadId; 3],
}

impl World {
    fn new() -> Self {
        let (k, a) = kernel_with_process();
        let ta = k.current_thread().unwrap();
        let mut w = Self {
            k,
            m: Model::default(),
            pids: [a; 3],
            tids: [ta; 3],
        };
        // init's own boot capability (ProcessCreate) is a kernel root.
        let create = CapSlot::new(0).unwrap();
        w.m.add((a, create), None, false);
        for i in 1..3 {
            w.spawn(i);
        }
        for call in [Syscall::CreateEndpoint, Syscall::CreateNotification] {
            let s = w.slot(0, call);
            w.m.add((a, s), None, false);
        }
        let r = w.slot(
            0,
            Syscall::CreateRegion {
                pages: PageCount::new(1).unwrap(),
            },
        );
        w.m.add((a, r), None, true);
        w
    }

    fn slot(&mut self, who: usize, call: Syscall) -> CapSlot {
        self.k.switch_to_thread_for_test(self.tids[who]);
        match self.k.handle_syscall(call) {
            SyscallOutcome::Complete(SyscallResult::Slot(s)) => s,
            _ => panic!("expected a slot"),
        }
    }

    /// (Re)creates process `i` from init; init gets a root process cap.
    fn spawn(&mut self, i: usize) {
        let cap = self.slot(
            0,
            Syscall::CreateProcess {
                create_cap: CapSlot::new(0).unwrap(),
                image: DataRef::empty(),
                limits: super::helpers::CHILD_LIMITS,
            },
        );
        self.pids[i] = self.k.process_of_cap_for_test(cap);
        self.tids[i] =
            self.k
                .spawn_thread_for_test(self.pids[i], Priority::DEFAULT, TimeSlice::DEFAULT);
        self.m.add((self.pids[0], cap), None, false);
    }

    fn run(&mut self, who: usize, call: Syscall) -> bool {
        self.k.switch_to_thread_for_test(self.tids[who]);
        !matches!(
            self.k.handle_syscall(call),
            SyscallOutcome::Complete(SyscallResult::Error(_))
        )
    }

    fn step(&mut self, rng: &mut Rng) {
        let who = rng.below(3);
        let pid = self.pids[who];
        let held = self.m.held_by(pid);
        let Some((id, slot)) = rng.pick(&held) else {
            return;
        };

        match rng.below(10) {
            0..=2 => {
                // Badges only go on endpoints and notifications; the kernel
                // must refuse the rest, so try both.
                let badge = Badge::new(rng.below(2));
                let before = self.k.cap_slots_for_test(pid);
                let ok = self.run(
                    who,
                    Syscall::Mint {
                        slot,
                        badge,
                        permissions: all(),
                    },
                );
                if ok {
                    let new = newly_used(before, self.k.cap_slots_for_test(pid));
                    let is_region = self.m.nodes[id].is_region;
                    self.m.add((pid, new), Some(id), is_region);
                }
            }
            3..=4 => {
                // Delegate through any process cap the caller holds.
                let proc_caps: Vec<(CapSlot, ProcessId)> = held
                    .iter()
                    .filter_map(|(_, s)| {
                        self.k.process_cap_target_for_test(pid, *s).map(|t| (*s, t))
                    })
                    .collect();
                if let Some((to, target)) = rng.pick(&proc_caps) {
                    let before = self.k.cap_slots_for_test(target);
                    let ok = self.run(
                        who,
                        Syscall::Delegate {
                            slot,
                            to_endpoint: to,
                            permissions: all(),
                        },
                    );
                    if ok {
                        let new = newly_used(before, self.k.cap_slots_for_test(target));
                        let is_region = self.m.nodes[id].is_region;
                        self.m.add((target, new), Some(id), is_region);
                    }
                }
            }
            // init keeps its ProcessCreate cap so processes can be re-made.
            5 if !(who == 0 && slot == CapSlot::new(0).unwrap()) => {
                if self.run(who, Syscall::DropCap { slot }) {
                    self.m.delete(id);
                }
            }
            6..=7 => {
                assert!(
                    self.run(who, Syscall::Revoke { slot }),
                    "revoking a held cap failed"
                );
                self.m.revoke(id);
            }
            8 => {
                let va = VirtualAddress::from_raw(0x30000 + rng.below(16) * 0x1000);
                let _ = self.run(
                    who,
                    Syscall::MapRegion {
                        region: slot,
                        va,
                        perms: PagePerms::ReadWrite,
                    },
                );
            }
            _ if who == 0 && rng.below(4) == 0 => {
                let victim = 1 + rng.below(2);
                let vp = self.pids[victim];
                // init's cap to the victim may already be revoked; kill it
                // directly, as a supervisor holding a cap would.
                for (n, _) in self.m.held_by(vp) {
                    self.m.delete(n);
                }
                self.k.destroy_process(vp);
                self.spawn(victim);
            }
            _ => {}
        }
    }

    fn check(&self) {
        for &pid in &self.pids {
            let used = self.k.cap_slots_for_test(pid);
            let mut expected = [false; 64];
            for (_, s) in self.m.held_by(pid) {
                expected[s.value() as usize] = true;
            }
            assert_eq!(used, expected, "{pid:?} holds the wrong capabilities");
        }
        // Every capability-backed mapping is backed by a live capability
        // of the mapping process.
        for (_, region) in self.k.regions.iter() {
            for m in region.mappings() {
                if let Some(via) = m.via {
                    assert!(self.k.resolves(via), "mapping outlived its capability");
                    assert_eq!(via.process, m.process);
                }
            }
        }
    }
}

fn newly_used(before: [bool; 64], after: [bool; 64]) -> CapSlot {
    (0..64)
        .find(|&i| after[i] && !before[i])
        .and_then(|i| CapSlot::new(i as u8))
        .expect("a new slot")
}

#[test]
fn derivation_tree_matches_the_model() {
    let extra: u64 = std::env::var("MODEL_SEEDS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0);
    for seed in [5u64, 0xBADC0DE, 31337, 0x5EED, 2024]
        .into_iter()
        .chain(1000..1000 + extra)
    {
        let mut world = World::new();
        let mut rng = Rng(seed);
        let mut revokes = 0;
        for _ in 0..STEPS {
            let before = world.m.alive.len();
            world.step(&mut rng);
            if world.m.alive.len() < before {
                revokes += 1;
            }
            world.check();
        }
        if seed >= 1000 {
            continue; // extra seeds hunt for mismatches, not coverage
        }
        assert!(
            revokes > 20,
            "seed {seed:#x}: too little deletion exercised"
        );
        assert!(
            world.m.nodes.len() > 100,
            "seed {seed:#x}: too little derivation"
        );
    }
}