use std::collections::BTreeMap;
use turkos_os_address_spaces::spaces::{Backend, Configuration, Error, Spaces};

#[derive(Debug)]
struct Root {
    address: u64,
    owner: u64,
}
#[derive(Default)]
struct Model {
    next: u64,
    roots: BTreeMap<u64, u64>,
    active: u64,
    acquired: usize,
    retired: usize,
    failed_acquire: usize,
    failed_retire: bool,
    invalid: bool,
    invalidations: Vec<(u16, u64, u64)>,
    selections: Vec<(u64, u16)>,
    cache: BTreeMap<u16, u64>,
    selected_data: u64,
}
impl Backend for Model {
    type Root = Root;
    fn acquire_root(&mut self, owner: u64) -> Result<Root, Error> {
        self.acquired += 1;
        if self.failed_acquire != 0 {
            self.failed_acquire -= 1;
            return Err(Error::NoMemory);
        }
        self.next += 4096;
        let address = 0x4000_0000 + self.next;
        assert!(self.roots.insert(address, owner).is_none());
        Ok(Root { address, owner })
    }
    fn root_address(&self, root: &Root) -> u64 {
        root.address
    }
    fn validate_root(&self, root: &Root) -> Result<(), Error> {
        if self.invalid || self.roots.get(&root.address) != Some(&root.owner) {
            return Err(Error::Backing);
        }
        Ok(())
    }
    fn invalidate_recycled(&mut self, asid: u16, old: u64, new: u64) {
        assert!(old != 0 && new > old);
        self.invalidations.push((asid, old, new));
        self.cache.remove(&asid);
    }
    fn select(&mut self, root: &Root, asid: u16) -> Result<(), Error> {
        self.validate_root(root)?;
        self.active = root.address;
        self.selections.push((root.address, asid));
        self.selected_data = *self.cache.entry(asid).or_insert(root.owner);
        Ok(())
    }
    fn restore_kernel(&mut self) -> Result<(), Error> {
        self.active = 0;
        Ok(())
    }
    fn retire_root(&mut self, root: Root, _: u16, _: u64) -> Result<(), (Error, Root)> {
        if self.active == root.address {
            return Err((Error::Active, root));
        }
        if self.failed_retire {
            return Err((Error::Retirement, root));
        }
        assert_eq!(self.roots.remove(&root.address), Some(root.owner));
        self.retired += 1;
        Ok(())
    }
}
fn configuration() -> Configuration {
    Configuration::from_registers(2, 25).unwrap()
}
fn backend<B, const N: usize, const A: usize, R>(
    spaces: &Spaces<B, N, A>,
    observe: impl FnOnce(&B) -> R,
) -> R
where
    B: Backend,
{
    spaces.with_backend_observation(observe)
}

#[test]
fn actual_width_uses_mmfr0_and_tcr_separately() {
    let eight = Configuration::from_registers(2, 25).unwrap();
    assert_eq!((eight.hardware_bits(), eight.active_bits()), (8, 8));
    let hardware16_active8 = Configuration::from_registers(0x22, 25).unwrap();
    assert_eq!(
        (
            hardware16_active8.hardware_bits(),
            hardware16_active8.active_bits()
        ),
        (16, 8)
    );
    let sixteen = Configuration::from_registers(0x22, 25 | (1 << 36)).unwrap();
    assert_eq!((sixteen.hardware_bits(), sixteen.active_bits()), (16, 16));
    assert_eq!(sixteen.user_asids(), 65535);
    assert_eq!(eight.user_asids(), 255);
    assert_eq!(
        Configuration::from_registers(2, 25 | (1 << 36)),
        Err(Error::Configuration)
    );
    assert_eq!(
        Configuration::from_registers(0x12, 25),
        Err(Error::Configuration)
    );
}
#[test]
fn unsupported_translation_regimes_are_refused() {
    for bits in [1 << 7, 1 << 22, 1 << 14, 1 << 37, 1 << 38, 1 << 59] {
        assert_eq!(
            Configuration::from_registers(2, 25 | bits),
            Err(Error::Configuration)
        );
    }
    assert_eq!(
        Configuration::from_registers(2, 24),
        Err(Error::Configuration)
    );
    assert_eq!(
        Configuration::from_registers(2 | (15 << 28), 25),
        Err(Error::Configuration)
    );
    assert_eq!(
        Configuration::from_registers(2, 25 | (5 << 32)),
        Err(Error::Configuration)
    );
    assert_eq!(
        Configuration::from_registers(5, 16 | (5 << 32))
            .unwrap()
            .active_bits(),
        8
    );
}
#[test]
fn kernel_asid_zero_and_illegal_root_bits_never_encode() {
    let config = configuration();
    assert_eq!(config.ttbr(0x4000_1000, 7), Ok(0x0007_0000_4000_1000));
    assert_eq!(config.tlbi_operand(7), Ok(0x0007_0000_0000_0000));
    for root in [0, 0x4000_1001, 1 << 48] {
        assert_eq!(config.ttbr(root, 1), Err(Error::Configuration));
    }
    for asid in [0, 256, u16::MAX] {
        assert_eq!(config.ttbr(0x4000_1000, asid), Err(Error::Configuration));
        assert_eq!(config.tlbi_operand(asid), Err(Error::Configuration));
    }
    let sixteen = Configuration::from_registers(0x22, 25 | (1 << 36)).unwrap();
    assert_eq!(sixteen.ttbr(0x4000_1000, 65535), Ok(0xffff_0000_4000_1000));
}
#[test]
fn software_quota_is_bounded_by_real_active_width() {
    assert!(matches!(
        Spaces::<Model, 1, 256>::new(1, configuration(), Model::default()),
        Err(Error::Configuration)
    ));
    assert!(matches!(
        Spaces::<Model, 1, 0>::new(1, configuration(), Model::default()),
        Err(Error::Configuration)
    ));
    assert!(matches!(
        Spaces::<Model, 0, 1>::new(1, configuration(), Model::default()),
        Err(Error::Configuration)
    ));
    assert!(matches!(
        Spaces::<Model, 1, 1>::new(0, configuration(), Model::default()),
        Err(Error::Configuration)
    ));
    let sixteen = Configuration::from_registers(0x22, 25 | (1 << 36)).unwrap();
    assert!(Spaces::<Model, 1, 256>::new(1, sixteen, Model::default()).is_ok());
}
#[test]
fn failed_root_acquisition_registers_neither_id_nor_asid() {
    let model = Model {
        failed_acquire: 1,
        ..Model::default()
    };
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), model).unwrap();
    assert_eq!(spaces.create(99), Err(Error::NoMemory));
    assert_eq!(spaces.owned_records().count(), 0);
    assert_eq!(
        backend(&spaces, |b| (b.roots.len(), b.invalidations.len())),
        (0, 0)
    );
    spaces.audit().unwrap();
    let retry = spaces.create(99).unwrap();
    let row = spaces.observe(retry, 99).unwrap();
    assert_eq!(
        (row.slot, row.generation, row.asid, row.asid_generation),
        (0, 1, 1, 1)
    );
}
#[test]
fn id_capacity_fails_before_root_acquisition() {
    let mut spaces = Spaces::<_, 1, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(99).unwrap();
    let before = spaces.observe(a, 99).unwrap();
    assert_eq!(spaces.create(100), Err(Error::Capacity));
    assert_eq!(backend(&spaces, |b| b.acquired), 1);
    assert_eq!(spaces.observe(a, 99), Ok(before));
}
#[test]
fn scarce_asid_capacity_fails_before_root_acquisition() {
    let mut spaces = Spaces::<_, 3, 1>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(99).unwrap();
    let before = spaces.observe(a, 99).unwrap();
    assert_eq!(spaces.create(100), Err(Error::AsidsFull));
    assert_eq!(backend(&spaces, |b| b.acquired), 1);
    assert_eq!(spaces.observe(a, 99), Ok(before));
}
#[test]
fn two_live_spaces_have_distinct_roots_and_asids() {
    let mut spaces = Spaces::<_, 3, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    let b = spaces.create(22).unwrap();
    let a = spaces.observe(a, 11).unwrap();
    let b = spaces.observe(b, 22).unwrap();
    assert_ne!(a.root, b.root);
    assert_eq!((a.asid, b.asid), (1, 2));
    assert_eq!(
        (
            a.generation,
            b.generation,
            a.asid_generation,
            b.asid_generation
        ),
        (1, 1, 1, 1)
    );
    spaces.audit().unwrap();
}
#[test]
fn a_b_a_selection_preserves_each_cached_private_value() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    let b = spaces.create(22).unwrap();
    for (handle, owner) in [(a, 11), (b, 22), (a, 11)] {
        spaces.activate(handle, owner).unwrap();
        assert_eq!(backend(&spaces, |b| b.selected_data), owner);
    }
    assert_eq!(backend(&spaces, |b| b.invalidations.len()), 0);
    assert_eq!(
        backend(&spaces, |b| b
            .selections
            .iter()
            .map(|v| v.1)
            .collect::<Vec<_>>()),
        [1, 2, 1]
    );
}
#[test]
fn recycled_asid_flushes_old_cached_value_before_new_activation() {
    let mut spaces = Spaces::<_, 2, 1>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    spaces.activate(a, 11).unwrap();
    spaces.restore_kernel().unwrap();
    spaces.release(a, 11).unwrap();
    let b = spaces.create(22).unwrap();
    let observed = spaces.observe(b, 22).unwrap();
    assert_eq!(
        (observed.asid, observed.asid_generation, observed.generation),
        (1, 2, 2)
    );
    assert_eq!(backend(&spaces, |b| b.invalidations.clone()), [(1, 1, 2)]);
    spaces.activate(b, 22).unwrap();
    assert_eq!(backend(&spaces, |b| b.selected_data), 22);
    assert_eq!(spaces.activate(a, 11), Err(Error::Stale));
}
#[test]
fn foreign_registry_and_wrong_owner_fail_before_backend_selection() {
    let mut a = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let mut b = Spaces::<_, 2, 2>::new(2, configuration(), Model::default()).unwrap();
    let handle = a.create(11).unwrap();
    b.create(11).unwrap();
    assert_eq!(b.activate(handle, 11), Err(Error::Foreign));
    assert_eq!(a.activate(handle, 22), Err(Error::Foreign));
    assert_eq!(a.release(handle, 22), Err(Error::Foreign));
    assert_eq!(backend(&a, |b| b.selections.len()), 0);
    assert_eq!(backend(&b, |b| b.selections.len()), 0);
    a.audit().unwrap();
    b.audit().unwrap();
}
#[test]
fn owner_zero_is_rejected_before_acquisition() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    assert_eq!(spaces.create(0), Err(Error::Foreign));
    assert_eq!(backend(&spaces, |b| b.acquired), 0);
}
#[test]
fn active_root_release_retains_identical_owner_token() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    let before = spaces.activate(a, 11).unwrap();
    assert_eq!(spaces.release(a, 11), Err(Error::Active));
    assert_eq!(spaces.observe(a, 11), Ok(before));
    assert_eq!(backend(&spaces, |b| (b.roots.len(), b.retired)), (1, 0));
    spaces.audit().unwrap();
    spaces.restore_kernel().unwrap();
    spaces.release(a, 11).unwrap();
    assert_eq!(spaces.owned_records().count(), 0);
}
#[test]
fn failed_retirement_restores_id_and_asid_exactly() {
    let model = Model {
        failed_retire: true,
        ..Model::default()
    };
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), model).unwrap();
    let a = spaces.create(11).unwrap();
    let before = spaces.observe(a, 11).unwrap();
    assert_eq!(spaces.release(a, 11), Err(Error::Retirement));
    assert_eq!(spaces.observe(a, 11), Ok(before));
    assert_eq!(backend(&spaces, |b| (b.roots.len(), b.retired)), (1, 0));
    let b = spaces.create(22).unwrap();
    assert_eq!(spaces.observe(b, 22).unwrap().asid, 2);
    spaces.audit().unwrap();
}
#[test]
fn successful_retirement_rejects_stale_and_double_release() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    spaces.release(a, 11).unwrap();
    assert_eq!(spaces.release(a, 11), Err(Error::Stale));
    assert_eq!(spaces.activate(a, 11), Err(Error::Stale));
    let b = spaces.create(11).unwrap();
    assert_eq!(b.slot(), a.slot());
    assert!(b.generation() > a.generation());
    assert_eq!(spaces.observe(a, 11), Err(Error::Stale));
    spaces.release(b, 11).unwrap();
    assert_eq!(backend(&spaces, |b| (b.roots.len(), b.retired)), (0, 2));
}
#[test]
fn immutable_root_observation_does_not_duplicate_or_release_owner() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    let values = spaces
        .with_root_observation(a, 11, |backend, root| {
            (root.owner, backend.root_address(root))
        })
        .unwrap();
    assert_eq!(values, (11, spaces.observe(a, 11).unwrap().root));
    assert_eq!(backend(&spaces, |b| (b.acquired, b.retired)), (1, 0));
}

#[test]
fn backend_extraction_requires_every_root_id_retired() {
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), Model::default()).unwrap();
    let a = spaces.create(11).unwrap();
    let (error, mut spaces) = match spaces.into_backend() {
        Err(failure) => failure,
        Ok(_) => panic!("live root backend escaped"),
    };
    assert_eq!(error, Error::Active);
    assert_eq!(spaces.observe(a, 11).unwrap().owner, 11);
    spaces.release(a, 11).unwrap();
    let model = spaces.into_backend().ok().expect("all IDs retired");
    assert_eq!(
        (model.roots.len(), model.acquired, model.retired),
        (0, 1, 1)
    );
}

#[test]
fn invalid_backing_stops_observation_activation_and_release() {
    let model = Model {
        invalid: true,
        ..Model::default()
    };
    let mut spaces = Spaces::<_, 2, 2>::new(1, configuration(), model).unwrap();
    let a = spaces.create(11).unwrap();
    assert_eq!(spaces.observe(a, 11), Err(Error::Backing));
    assert_eq!(spaces.activate(a, 11), Err(Error::Backing));
    assert_eq!(spaces.release(a, 11), Err(Error::Backing));
    assert_eq!(spaces.audit(), Err(Error::Backing));
    assert_eq!(
        backend(&spaces, |b| (b.roots.len(), b.selections.len(), b.retired)),
        (1, 0, 0)
    );
}