use std::{
    cell::RefCell,
    collections::{BTreeMap, BTreeSet},
    rc::Rc,
};
use turkos_os_frame_authority::authority::{Allocator, Lease, Plan};
use turkos_os_frame_metadata::metadata::{self, Registry};
use turkos_os_page_tables::{
    descriptor::{Geometry, Memory, Permissions},
    tree::{Backend, Error, Tree},
};
use turkos_os_runtime_platform::ram::Region;

const BASE: u64 = 0x4000_0000;
const RW: Permissions = Permissions {
    user: true,
    write: true,
    execute: false,
};
const RX: Permissions = Permissions {
    user: true,
    write: false,
    execute: true,
};
const RO: Permissions = Permissions {
    user: true,
    write: false,
    execute: false,
};
#[derive(Clone, Debug, PartialEq, Eq)]
enum Event {
    Allocate(u64),
    Store(u64, usize, u64),
    Dsb,
    Tlbi(u64, Option<u64>),
    Isb,
    Release(u64),
    Select(u64, u16),
}
struct State {
    authority: Allocator,
    leases: Vec<Lease>,
    bytes: BTreeMap<u64, Box<[u8; 4096]>>,
    shared: BTreeMap<u64, [u64; 512]>,
    shared_pages: BTreeSet<u64>,
    events: Vec<Event>,
    budget: Option<usize>,
    fail_release: BTreeSet<u64>,
    fail_write: Option<(u64, usize)>,
    active: u64,
}
impl State {
    fn new(pages: u64) -> Self {
        Self {
            authority: Allocator::new(
                Plan::already_admitted(&[Region {
                    start: BASE,
                    end: BASE + pages * 4096,
                }])
                .unwrap(),
                43,
            )
            .unwrap(),
            leases: Vec::new(),
            bytes: BTreeMap::new(),
            shared: BTreeMap::new(),
            shared_pages: BTreeSet::new(),
            events: Vec::new(),
            budget: None,
            fail_release: BTreeSet::new(),
            fail_write: None,
            active: 0,
        }
    }
    fn raw(&self, address: u64, index: usize) -> u64 {
        u64::from_le_bytes(
            self.bytes[&address][index * 8..index * 8 + 8]
                .try_into()
                .unwrap(),
        )
    }
    fn live(&self) -> Vec<u64> {
        self.authority.live_addresses().collect()
    }
}
#[derive(Clone)]
struct Frames(Rc<RefCell<State>>);
impl metadata::Backend for Frames {
    fn allocate(&mut self) -> Result<u64, metadata::Error> {
        let mut s = self.0.borrow_mut();
        if let Some(remaining) = s.budget.as_mut() {
            if *remaining == 0 {
                return Err(metadata::Error::NoMemory);
            }
            *remaining -= 1;
        }
        let lease = s
            .authority
            .allocate()
            .map_err(|_| metadata::Error::NoMemory)?;
        let address = lease.address();
        s.leases.push(lease);
        s.bytes.insert(address, Box::new([0xa5; 4096]));
        s.events.push(Event::Allocate(address));
        Ok(address)
    }
    fn release(&mut self, address: u64) -> Result<(), metadata::Error> {
        let mut s = self.0.borrow_mut();
        if s.fail_release.contains(&address) {
            return Err(metadata::Error::Backend);
        }
        let slot = s
            .leases
            .iter()
            .position(|l| l.address() == address)
            .ok_or(metadata::Error::Backend)?;
        let lease = s.leases.remove(slot);
        if s.authority.release(&lease).is_err() {
            s.leases.insert(slot, lease);
            return Err(metadata::Error::Backend);
        }
        s.events.push(Event::Release(address));
        Ok(())
    }
    fn zero_page(&mut self, address: u64) -> Result<(), metadata::Error> {
        self.0
            .borrow_mut()
            .bytes
            .get_mut(&address)
            .ok_or(metadata::Error::Backend)?
            .fill(0);
        Ok(())
    }
    fn read_byte(&self, address: u64, offset: usize) -> Result<u8, metadata::Error> {
        self.0
            .borrow()
            .bytes
            .get(&address)
            .and_then(|p| p.get(offset))
            .copied()
            .ok_or(metadata::Error::Backend)
    }
    fn available_pages(&self) -> u64 {
        self.0
            .borrow()
            .authority
            .free_regions()
            .iter()
            .map(|r| (r.end - r.start) / 4096)
            .sum()
    }
}
struct Tables {
    registry: Registry<Frames, 64>,
    state: Rc<RefCell<State>>,
}
impl Backend for Tables {
    type Handle = metadata::Handle;
    fn allocate_zeroed(&mut self) -> Result<Self::Handle, Error> {
        self.registry.allocate_zeroed(81).map_err(|e| {
            if e == metadata::Error::NoMemory {
                Error::NoMemory
            } else {
                Error::Backend
            }
        })
    }
    fn address(&self, h: Self::Handle) -> u64 {
        h.address()
    }
    fn validate(&self, h: Self::Handle) -> Result<(), Error> {
        self.registry
            .read_byte(h, 81, 4095)
            .map(|_| ())
            .map_err(|_| Error::Stale)
    }
    fn release(&mut self, h: Self::Handle) -> Result<(), Error> {
        self.registry.release(h, 81).map_err(|_| Error::Backend)
    }
    fn entry(&self, h: Self::Handle, index: usize) -> Result<u64, Error> {
        if index >= 512 {
            return Err(Error::Backend);
        }
        let mut bytes = [0; 8];
        for (i, b) in bytes.iter_mut().enumerate() {
            *b = self
                .registry
                .read_byte(h, 81, index * 8 + i)
                .map_err(|_| Error::Backend)?;
        }
        Ok(u64::from_le_bytes(bytes))
    }
    fn write_entry(&mut self, h: Self::Handle, index: usize, value: u64) -> Result<(), Error> {
        self.validate(h)?;
        if index >= 512 {
            return Err(Error::Backend);
        }
        let mut s = self.state.borrow_mut();
        if s.fail_write == Some((h.address(), index)) {
            return Err(Error::Backend);
        }
        s.bytes.get_mut(&h.address()).ok_or(Error::Backend)?[index * 8..index * 8 + 8]
            .copy_from_slice(&value.to_le_bytes());
        s.events.push(Event::Store(h.address(), index, value));
        Ok(())
    }
    fn validate_shared_table(&self, address: u64) -> Result<(), Error> {
        if self.state.borrow().shared.contains_key(&address) {
            Ok(())
        } else {
            Err(Error::Shared)
        }
    }
    fn validate_shared_page(&self, address: u64) -> Result<(), Error> {
        if self.state.borrow().shared_pages.contains(&address) {
            Ok(())
        } else {
            Err(Error::Shared)
        }
    }
    fn shared_entry(&self, address: u64, index: usize) -> Result<u64, Error> {
        self.state
            .borrow()
            .shared
            .get(&address)
            .and_then(|p| p.get(index))
            .copied()
            .ok_or(Error::Shared)
    }
    fn active_root(&self) -> u64 {
        self.state.borrow().active
    }
    fn invalidate(&mut self, root: u64, va: Option<u64>) {
        self.state.borrow_mut().events.extend([
            Event::Dsb,
            Event::Tlbi(root, va),
            Event::Dsb,
            Event::Isb,
        ]);
    }
    fn select_root(&mut self, root: u64, _geometry: Geometry, asid: u16) {
        let mut s = self.state.borrow_mut();
        s.active = root;
        s.events.push(Event::Select(root, asid));
    }
}
fn fixture<const N: usize>(bits: u8, pages: u64) -> (Tree<Tables, N>, Rc<RefCell<State>>) {
    let state = Rc::new(RefCell::new(State::new(pages)));
    let backend = Tables {
        registry: Registry::new(991, Frames(state.clone())).unwrap(),
        state: state.clone(),
    };
    (
        Tree::new(821, Geometry::new(bits, 40).unwrap(), backend).unwrap(),
        state,
    )
}

#[test]
fn owned_four_level_chain_and_known_physical_translation() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let va = (7 << 39) | (13 << 30) | (29 << 21) | (31 << 12);
    tree.map(&root, va, &page, Memory::Normal, RX).unwrap();
    tree.audit().unwrap();
    let result = tree.walk(&root, va + 0xabc).unwrap();
    assert_eq!(result.physical, page.address() + 0xabc);
    assert_eq!(result.permissions, RX);
    assert_eq!(tree.owned_records().filter(|r| r.table).count(), 4);
    assert_eq!(
        tree.with_backend_observation(|b| b.registry.snapshot().live_pages),
        5
    );
    assert_eq!(state.borrow().live().len(), 5);
}
#[test]
fn second_table_enospc_publishes_no_prefix_and_returns_exact_inventory() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let before = state.borrow().live();
    let owners = tree.owned_records().collect::<Vec<_>>();
    state.borrow_mut().budget = Some(1);
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW),
        Err(Error::NoMemory)
    );
    assert_eq!(state.borrow().raw(root.address(), 0), 0);
    assert_eq!(state.borrow().live(), before);
    assert_eq!(tree.owned_records().collect::<Vec<_>>(), owners);
    tree.audit().unwrap();
}
#[test]
fn failed_rollback_keeps_actual_registry_handle_without_published_prefix() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    state.borrow_mut().budget = Some(1);
    state.borrow_mut().fail_release.insert(BASE + 2 * 4096);
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW),
        Err(Error::Retained)
    );
    assert_eq!(state.borrow().raw(root.address(), 0), 0);
    assert_eq!(state.borrow().live().len(), 3);
    let held = tree
        .owned_records()
        .filter(|r| r.quarantined)
        .collect::<Vec<_>>();
    assert_eq!(held.len(), 1);
    assert_eq!(held[0].address, BASE + 2 * 4096);
    assert_eq!(
        tree.with_backend_observation(|b| b.registry.snapshot().live_pages),
        3
    );
    tree.audit().unwrap();
    assert_eq!(tree.reap_quarantine(), Err(Error::Retained));
    state.borrow_mut().fail_release.clear();
    assert_eq!(tree.reap_quarantine(), Ok(1));
    assert_eq!(state.borrow().live().len(), 2);
}
#[test]
fn final_parent_store_failure_reclaims_unreachable_staged_tables() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let before = state.borrow().live();
    state.borrow_mut().fail_write = Some((root.address(), 0));
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW),
        Err(Error::Backend)
    );
    assert_eq!(state.borrow().raw(root.address(), 0), 0);
    assert_eq!(state.borrow().live(), before);
    tree.audit().unwrap();
}
#[test]
fn metadata_capacity_preflight_never_acquires_partial_table_chain() {
    let (mut tree, state) = fixture::<4>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let before = state.borrow().events.clone();
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW),
        Err(Error::Capacity)
    );
    assert_eq!(state.borrow().events, before);
    tree.audit().unwrap();
}
#[test]
fn detach_and_complete_tlbi_precede_each_intermediate_reclaim() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    state.borrow_mut().events.clear();
    tree.unmap(&root, 0x8000_0000).unwrap();
    let events = &state.borrow().events;
    assert_eq!(
        events
            .iter()
            .filter(|e| matches!(e, Event::Release(_)))
            .count(),
        3
    );
    for (index, event) in events.iter().enumerate() {
        if matches!(event, Event::Release(_)) {
            assert_eq!(
                &events[index - 4..index],
                [
                    Event::Dsb,
                    Event::Tlbi(root.address(), None),
                    Event::Dsb,
                    Event::Isb
                ]
            );
            assert!(matches!(events[index - 5], Event::Store(_, _, 0)));
        }
    }
    tree.audit().unwrap();
    assert_eq!(tree.owned_records().count(), 2);
}
#[test]
fn active_root_cannot_be_detached_or_freed_and_actual_token_is_returned() {
    let (mut tree, state) = fixture::<8>(48, 8);
    let root = tree.new_root().unwrap();
    tree.activate(&root, 11).unwrap();
    let address = root.address();
    let before = state.borrow().events.clone();
    assert_eq!(tree.detach_all(&root), Err(Error::ActiveRoot));
    let (error, root) = tree.release_root(root).unwrap_err();
    assert_eq!(error, Error::ActiveRoot);
    assert_eq!(root.address(), address);
    assert_eq!(state.borrow().events, before);
    state.borrow_mut().active = 0;
    tree.release_root(root).unwrap();
    assert!(state.borrow().live().is_empty());
}
#[test]
fn replacement_root_preserves_page_owner_and_live_mapping_after_old_retirement() {
    let (mut tree, state) = fixture::<20>(48, 20);
    let old = tree.new_root().unwrap();
    let new = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&old, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    tree.map_reference(&new, 0x8000_0000, page.reference(), Memory::Normal, RO)
        .unwrap();
    tree.activate(&new, 17).unwrap();
    tree.detach_all(&old).unwrap();
    tree.release_root(old).unwrap();
    tree.audit().unwrap();
    assert_eq!(
        tree.walk(&new, 0x8000_0123).unwrap().physical,
        page.address() + 0x123
    );
    let (error, page) = tree.release_page(page).unwrap_err();
    assert_eq!(error, Error::Mapped);
    assert!(state.borrow().live().contains(&page.address()));
    state.borrow_mut().active = 0;
    tree.detach_all(&new).unwrap();
    tree.release_root(new).unwrap();
    tree.release_page(page).unwrap();
    assert!(state.borrow().live().is_empty());
}
#[test]
fn stale_page_reference_and_foreign_owner_cannot_map_reused_storage() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let old = page.reference();
    tree.release_page(page).unwrap();
    let fresh = tree.allocate_page().unwrap();
    assert_eq!(fresh.address(), old.address());
    let before = state.borrow().live();
    assert_eq!(
        tree.map_reference(&root, 0x8000_0000, old, Memory::Normal, RW),
        Err(Error::Stale)
    );
    assert_eq!(state.borrow().live(), before);
    let other_state = Rc::new(RefCell::new(State::new(16)));
    let backend = Tables {
        registry: Registry::new(992, Frames(other_state.clone())).unwrap(),
        state: other_state,
    };
    let mut other = Tree::<_, 16>::new(822, Geometry::new(48, 40).unwrap(), backend).unwrap();
    let foreign = other.allocate_page().unwrap();
    assert_eq!(
        tree.map(&root, 0x8000_0000, &foreign, Memory::Normal, RW),
        Err(Error::Foreign)
    );
    tree.audit().unwrap();
}
#[test]
fn shared_kernel_tables_stay_borrowed_and_cannot_be_mutated_or_reclaimed() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    let shared = 0x7000_0000;
    state.borrow_mut().shared.insert(shared, [0; 512]);
    tree.share_table(&root, 0, 0, shared).unwrap();
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW),
        Err(Error::Shared)
    );
    tree.audit().unwrap();
    tree.detach_all(&root).unwrap();
    tree.release_root(root).unwrap();
    tree.release_page(page).unwrap();
    assert!(state.borrow().shared.contains_key(&shared));
    assert!(
        !state
            .borrow()
            .events
            .iter()
            .any(|e| matches!(e,Event::Release(a)if *a==shared))
    );
}
#[test]
fn reclaim_failure_retains_detached_table_and_page_release_retry_keeps_token() {
    let (mut tree, state) = fixture::<16>(48, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    let table = tree
        .owned_records()
        .filter(|r| r.table && !r.root)
        .last()
        .unwrap()
        .address;
    state.borrow_mut().fail_release.insert(table);
    assert_eq!(tree.unmap(&root, 0x8000_0000), Err(Error::Retained));
    tree.audit().unwrap();
    assert_eq!(tree.owned_records().filter(|r| r.quarantined).count(), 1);
    state.borrow_mut().fail_release.insert(page.address());
    let (error, page) = tree.release_page(page).unwrap_err();
    assert_eq!(error, Error::Backend);
    assert!(state.borrow().live().contains(&page.address()));
    state.borrow_mut().fail_release.clear();
    tree.release_page(page).unwrap();
    tree.reap_quarantine().unwrap();
    tree.release_root(root).unwrap();
    assert!(state.borrow().live().is_empty());
}
#[test]
fn three_level_geometry_uses_two_intermediates_and_refuses_duplicate_leaf() {
    let (mut tree, state) = fixture::<16>(39, 16);
    let root = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&root, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    let before = state.borrow().events.clone();
    assert_eq!(tree.owned_records().filter(|r| r.table).count(), 3);
    assert_eq!(
        tree.map(&root, 0x8000_0000, &page, Memory::Normal, RO),
        Err(Error::AlreadyMapped)
    );
    assert_eq!(state.borrow().events, before);
    tree.audit().unwrap();
}

#[test]
fn user_writable_executable_aliases_are_refused_across_all_activatable_roots() {
    let (mut tree, state) = fixture::<24>(48, 24);
    let a = tree.new_root().unwrap();
    let b = tree.new_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&a, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    let before = state.borrow().events.clone();
    assert_eq!(
        tree.map(&a, 0x8000_1000, &page, Memory::Normal, RX),
        Err(Error::AliasPermissions)
    );
    assert_eq!(
        tree.map(&b, 0x8000_0000, &page, Memory::Normal, RX),
        Err(Error::AliasPermissions)
    );
    assert_eq!(state.borrow().events, before);
    tree.map(&b, 0x8000_0000, &page, Memory::Normal, RO)
        .unwrap();
    tree.map(
        &a,
        0x4000_0000,
        &page,
        Memory::Normal,
        Permissions {
            user: false,
            write: true,
            execute: false,
        },
    )
    .unwrap();
    tree.audit().unwrap();
}
#[test]
fn replacement_permissions_publish_only_after_old_root_is_quiescent() {
    let (mut tree, state) = fixture::<24>(48, 24);
    let old = tree.new_root().unwrap();
    let draft = tree.new_draft_root().unwrap();
    let page = tree.allocate_page().unwrap();
    tree.map(&old, 0x8000_0000, &page, Memory::Normal, RW)
        .unwrap();
    tree.map(&draft, 0x8000_0000, &page, Memory::Normal, RX)
        .unwrap();
    tree.activate(&old, 7).unwrap();
    assert_eq!(tree.activate(&draft, 8), Err(Error::Quiescent));
    assert_eq!(tree.publish_root(&draft), Err(Error::AliasPermissions));
    assert_eq!(tree.quiesce(&old), Err(Error::ActiveRoot));
    state.borrow_mut().active = 0;
    tree.quiesce(&old).unwrap();
    assert_eq!(tree.activate(&old, 7), Err(Error::Quiescent));
    tree.publish_root(&draft).unwrap();
    tree.activate(&draft, 8).unwrap();
    assert_eq!(tree.publish_root(&old), Err(Error::AliasPermissions));
    tree.detach_all(&old).unwrap();
    tree.release_root(old).unwrap();
    assert_eq!(tree.walk(&draft, 0x8000_0000).unwrap().permissions, RX);
    tree.audit().unwrap();
}