use crate::classify::{Access, Kind, UserFault};
use turkos_os_mapping::mapping::Owner;
use turkos_os_page_tables::descriptor::Permissions;
pub const PAGE: u64 = 4096;
pub const USER_RW: Permissions = Permissions {
    user: true,
    write: true,
    execute: false,
};
pub const USER_RO: Permissions = Permissions {
    user: true,
    write: false,
    execute: false,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Invalid,
    Foreign,
    Capacity,
    Quota,
    NoMemory,
    Backend,
    Conflict,
    Protection,
    Stack,
    Retained,
    Busy,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Value<R: Copy + Eq> {
    pub frame: R,
    pub permissions: Permissions,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Change<R: Copy + Eq> {
    pub owner: Owner,
    pub address: u64,
    pub before: Option<Value<R>>,
    pub after: Option<Value<R>>,
}
/// Lease is a nonclone private page/metadata owner. Reference is inert. Failed
/// release returns the actual token. prepare holds every draft table resource
/// before publication, preserving current PTEs and retained rollback tokens on
/// failure. commit completes the synchronized root switch and BBM/TLBI before
/// returning; callers may reclaim old unreferenced data only after that point.
pub trait Backend {
    type Lease;
    type Reference: Copy + Eq;
    type Prepared;
    fn current_owner(&self) -> Option<Owner>;
    fn validate_owner(&self, owner: Owner) -> Result<(), Error>;
    fn allocate_zeroed(&mut self) -> Result<Self::Lease, Error>;
    fn reference(&self, lease: &Self::Lease) -> Self::Reference;
    fn physical_page(&self, lease: &Self::Lease) -> u64;
    fn validate(&self, lease: &Self::Lease) -> Result<(), Error>;
    fn release(&mut self, lease: Self::Lease) -> Result<(), (Error, Self::Lease)>;
    fn initialize(&mut self, lease: &Self::Lease, bytes: &[u8; 4096]) -> Result<(), Error>;
    fn copy_page(&mut self, source: &Self::Lease, target: &Self::Lease) -> Result<(), Error>;
    fn current(&self, owner: Owner, address: u64) -> Result<Option<Value<Self::Reference>>, Error>;
    fn prepare(&mut self, changes: &[Change<Self::Reference>]) -> Result<Self::Prepared, Error>;
    fn commit(&mut self, prepared: Self::Prepared);
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RegionKind {
    Anonymous,
    Stack { resident_low: u64 },
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Region {
    pub owner: Owner,
    pub start: u64,
    pub end: u64,
    pub kind: RegionKind,
}
#[derive(Clone, Copy)]
struct Process {
    owner: Owner,
    quota: usize,
    used: usize,
}
#[derive(Clone, Copy)]
struct Entry {
    owner: Owner,
    address: u64,
    page: usize,
    cow: bool,
}
struct PageOwner<L> {
    lease: L,
    references: usize,
    quarantined: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot {
    pub processes: usize,
    pub regions: usize,
    pub mappings: usize,
    pub owned_pages: usize,
    pub quarantined_pages: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Resolved {
    Present,
    Anonymous,
    Stack,
    CopyOnWrite,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Resolution<R: Copy + Eq> {
    pub kind: Resolved,
    pub frame: R,
    pub allocated: bool,
    pub retained_cleanup: bool,
}
pub struct Pager<B: Backend, const N: usize> {
    backend: B,
    processes: [Option<Process>; N],
    regions: [Option<Region>; N],
    entries: [Option<Entry>; N],
    pages: [Option<PageOwner<B::Lease>>; N],
}
impl<B: Backend, const N: usize> Pager<B, N> {
    pub fn new(backend: B) -> Self {
        Self {
            backend,
            processes: [None; N],
            regions: [None; N],
            entries: [None; N],
            pages: core::array::from_fn(|_| None),
        }
    }
    pub fn with_backend_observation<R>(&self, callback: impl FnOnce(&B) -> R) -> R {
        callback(&self.backend)
    }
    pub fn snapshot(&self) -> Snapshot {
        Snapshot {
            processes: self.processes.iter().flatten().count(),
            regions: self.regions.iter().flatten().count(),
            mappings: self.entries.iter().flatten().count(),
            owned_pages: self.pages.iter().flatten().count(),
            quarantined_pages: self
                .pages
                .iter()
                .flatten()
                .filter(|p| p.quarantined)
                .count(),
        }
    }
    pub fn owned_pages(&self) -> impl Iterator<Item = (B::Reference, usize, bool)> + '_ {
        self.pages.iter().flatten().map(|p| {
            (
                self.backend.reference(&p.lease),
                p.references,
                p.quarantined,
            )
        })
    }
    pub fn mappings(&self) -> impl Iterator<Item = (Owner, u64, B::Reference, bool)> + '_ {
        self.entries.iter().flatten().map(|e| {
            (
                e.owner,
                e.address,
                self.backend
                    .reference(&self.pages[e.page].as_ref().unwrap().lease),
                e.cow,
            )
        })
    }
    pub fn register(&mut self, owner: Owner, quota: usize) -> Result<(), Error> {
        if owner.id == 0 || owner.generation == 0 || quota == 0 || quota > N {
            return Err(Error::Invalid);
        }
        self.backend.validate_owner(owner)?;
        if self.processes.iter().flatten().any(|p| p.owner == owner) {
            return Err(Error::Conflict);
        }
        let slot = self
            .processes
            .iter()
            .position(Option::is_none)
            .ok_or(Error::Capacity)?;
        self.processes[slot] = Some(Process {
            owner,
            quota,
            used: 0,
        });
        Ok(())
    }
    fn process(&self, owner: Owner) -> Result<usize, Error> {
        self.backend.validate_owner(owner)?;
        self.processes
            .iter()
            .position(|p| p.is_some_and(|p| p.owner == owner))
            .ok_or(Error::Foreign)
    }
    pub fn add_region(&mut self, region: Region) -> Result<(), Error> {
        self.process(region.owner)?;
        if region.start >= region.end
            || !region.start.is_multiple_of(PAGE)
            || !region.end.is_multiple_of(PAGE)
            || region.end > 1 << 48
        {
            return Err(Error::Invalid);
        }
        if let RegionKind::Stack { resident_low } = region.kind
            && (resident_low < region.start
                || resident_low > region.end
                || !resident_low.is_multiple_of(PAGE))
        {
            return Err(Error::Invalid);
        }
        if self
            .regions
            .iter()
            .flatten()
            .any(|r| r.owner == region.owner && region.start < r.end && r.start < region.end)
        {
            return Err(Error::Conflict);
        }
        if self
            .entries
            .iter()
            .flatten()
            .any(|e| e.owner == region.owner && region.start <= e.address && e.address < region.end)
        {
            return Err(Error::Conflict);
        }
        if matches!(region.kind, RegionKind::Stack { .. })
            && self
                .regions
                .iter()
                .flatten()
                .any(|r| r.owner == region.owner && matches!(r.kind, RegionKind::Stack { .. }))
        {
            return Err(Error::Conflict);
        }
        let slot = self
            .regions
            .iter()
            .position(Option::is_none)
            .ok_or(Error::Capacity)?;
        self.regions[slot] = Some(region);
        Ok(())
    }
    fn new_page(&mut self) -> Result<usize, Error> {
        let slot = self
            .pages
            .iter()
            .position(Option::is_none)
            .ok_or(Error::Capacity)?;
        let lease = self.backend.allocate_zeroed()?;
        self.pages[slot] = Some(PageOwner {
            lease,
            references: 0,
            quarantined: false,
        });
        if self
            .backend
            .validate(&self.pages[slot].as_ref().unwrap().lease)
            .is_err()
        {
            self.pages[slot].as_mut().unwrap().quarantined = true;
            return Err(Error::Retained);
        }
        let reference = self
            .backend
            .reference(&self.pages[slot].as_ref().unwrap().lease);
        let address = self
            .backend
            .physical_page(&self.pages[slot].as_ref().unwrap().lease);
        if !address.is_multiple_of(PAGE) || address >= 1 << 48 {
            self.pages[slot].as_mut().unwrap().quarantined = true;
            return Err(Error::Retained);
        }
        if self.pages.iter().enumerate().any(|(i, p)| {
            i != slot
                && p.as_ref().is_some_and(|p| {
                    self.backend.reference(&p.lease) == reference
                        || self.backend.physical_page(&p.lease) == address
                })
        }) {
            self.pages[slot].as_mut().unwrap().quarantined = true;
            return Err(Error::Retained);
        }
        Ok(slot)
    }
    fn discard(&mut self, slot: usize) -> bool {
        let page = self.pages[slot].take().unwrap();
        match self.backend.release(page.lease) {
            Ok(()) => false,
            Err((_, lease)) => {
                let _ = lease;
                self.pages[slot] = None;
                true
            }
        }
    }
    fn rollback(&mut self, slot: usize, error: Error) -> Error {
        if self.discard(slot) {
            Error::Retained
        } else {
            error
        }
    }
    fn checked_value(&self, entry: Entry) -> Result<Value<B::Reference>, Error> {
        let page = self
            .pages
            .get(entry.page)
            .and_then(Option::as_ref)
            .ok_or(Error::Backend)?;
        if page.quarantined || page.references == 0 {
            return Err(Error::Backend);
        }
        self.backend.validate(&page.lease)?;
        let value = Value {
            frame: self.backend.reference(&page.lease),
            permissions: if entry.cow { USER_RO } else { USER_RW },
        };
        if self.backend.current(entry.owner, entry.address)? != Some(value) {
            return Err(Error::Backend);
        }
        Ok(value)
    }
    pub fn resolve(
        &mut self,
        owner: Owner,
        fault: UserFault,
    ) -> Result<Resolution<B::Reference>, Error> {
        let process = self.process(owner)?;
        if self.backend.current_owner() != Some(owner) {
            return Err(Error::Foreign);
        }
        if let Some(stack) = self
            .regions
            .iter()
            .flatten()
            .find(|r| r.owner == owner && matches!(r.kind, RegionKind::Stack { .. }))
        {
            let sp = fault.stack_pointer();
            if !sp.is_multiple_of(16) || sp < stack.start || sp > stack.end {
                return Err(Error::Stack);
            }
        }
        if fault.access() == Access::Execute || fault.kind() == Kind::AccessFlag {
            return Err(Error::Protection);
        }
        let address = fault.address() & !(PAGE - 1);
        if let Some(entry_slot) = self
            .entries
            .iter()
            .position(|e| e.is_some_and(|e| e.owner == owner && e.address == address))
        {
            let entry = self.entries[entry_slot].unwrap();
            let before = self.checked_value(entry)?;
            if entry.cow && fault.kind() == Kind::Permission && fault.access() == Access::Write {
                return self.cow(entry_slot, before);
            }
            if fault.kind() != Kind::Translation {
                return Err(Error::Protection);
            }
            return Ok(Resolution {
                kind: Resolved::Present,
                frame: before.frame,
                allocated: false,
                retained_cleanup: false,
            });
        }
        if fault.kind() != Kind::Translation {
            return Err(Error::Protection);
        }
        let region_slot = self
            .regions
            .iter()
            .position(|r| {
                r.is_some_and(|r| r.owner == owner && r.start <= address && address < r.end)
            })
            .ok_or(Error::Protection)?;
        let region = self.regions[region_slot].unwrap();
        let kind = match region.kind {
            RegionKind::Anonymous => Resolved::Anonymous,
            RegionKind::Stack { resident_low } => {
                let sp = fault.stack_pointer();
                if !sp.is_multiple_of(16)
                    || resident_low.checked_sub(PAGE) != Some(address)
                    || sp < region.start
                    || sp > region.end
                    || sp < address
                    || sp > address + PAGE
                    || fault.address() < sp
                    || fault.address() - sp > 256
                {
                    return Err(Error::Stack);
                }
                Resolved::Stack
            }
        };
        if self.processes[process].unwrap().used >= self.processes[process].unwrap().quota {
            return Err(Error::Quota);
        }
        let entry_slot = self
            .entries
            .iter()
            .position(Option::is_none)
            .ok_or(Error::Capacity)?;
        if self.backend.current(owner, address)?.is_some() {
            return Err(Error::Conflict);
        }
        let page_slot = self.new_page()?;
        let frame = self
            .backend
            .reference(&self.pages[page_slot].as_ref().unwrap().lease);
        let change = Change {
            owner,
            address,
            before: None,
            after: Some(Value {
                frame,
                permissions: USER_RW,
            }),
        };
        let prepared = match self.backend.prepare(&[change]) {
            Ok(p) => p,
            Err(e) => return Err(self.rollback(page_slot, e)),
        };
        self.backend.commit(prepared);
        self.pages[page_slot].as_mut().unwrap().references = 1;
        self.entries[entry_slot] = Some(Entry {
            owner,
            address,
            page: page_slot,
            cow: false,
        });
        self.processes[process].as_mut().unwrap().used += 1;
        if matches!(kind, Resolved::Stack) {
            self.regions[region_slot].as_mut().unwrap().kind = RegionKind::Stack {
                resident_low: address,
            };
        }
        Ok(Resolution {
            kind,
            frame,
            allocated: true,
            retained_cleanup: false,
        })
    }
    fn cow(
        &mut self,
        entry_slot: usize,
        before: Value<B::Reference>,
    ) -> Result<Resolution<B::Reference>, Error> {
        let entry = self.entries[entry_slot].unwrap();
        let old = entry.page;
        let new = self.new_page()?;
        let result = {
            let (source, target) = if old < new {
                let (a, b) = self.pages.split_at(new);
                (a[old].as_ref().unwrap(), b[0].as_ref().unwrap())
            } else {
                let (a, b) = self.pages.split_at(old);
                (b[0].as_ref().unwrap(), a[new].as_ref().unwrap())
            };
            self.backend.copy_page(&source.lease, &target.lease)
        };
        if let Err(e) = result {
            return Err(self.rollback(new, e));
        }
        let frame = self
            .backend
            .reference(&self.pages[new].as_ref().unwrap().lease);
        let change = Change {
            owner: entry.owner,
            address: entry.address,
            before: Some(before),
            after: Some(Value {
                frame,
                permissions: USER_RW,
            }),
        };
        let prepared = match self.backend.prepare(&[change]) {
            Ok(p) => p,
            Err(e) => return Err(self.rollback(new, e)),
        };
        self.backend.commit(prepared);
        self.entries[entry_slot] = Some(Entry {
            page: new,
            cow: false,
            ..entry
        });
        self.pages[new].as_mut().unwrap().references = 1;
        self.pages[old].as_mut().unwrap().references -= 1;
        let retained = self.pages[old].as_ref().unwrap().references == 0 && self.discard(old);
        Ok(Resolution {
            kind: Resolved::CopyOnWrite,
            frame,
            allocated: true,
            retained_cleanup: retained,
        })
    }
    pub fn share_cow(
        &mut self,
        a: Owner,
        b: Owner,
        address_a: u64,
        address_b: u64,
        seed: &[u8; 4096],
    ) -> Result<(), Error> {
        if a == b
            || !address_a.is_multiple_of(PAGE)
            || !address_b.is_multiple_of(PAGE)
            || address_a >= 1 << 48
            || address_b >= 1 << 48
        {
            return Err(Error::Invalid);
        }
        let pa = self.process(a)?;
        let pb = self.process(b)?;
        for (owner, address, p) in [(a, address_a, pa), (b, address_b, pb)] {
            if self.processes[p].unwrap().used >= self.processes[p].unwrap().quota {
                return Err(Error::Quota);
            }
            if self
                .entries
                .iter()
                .flatten()
                .any(|e| e.owner == owner && e.address == address)
                || self.backend.current(owner, address)?.is_some()
            {
                return Err(Error::Conflict);
            }
            if self
                .regions
                .iter()
                .flatten()
                .any(|r| r.owner == owner && r.start <= address && address < r.end)
            {
                return Err(Error::Conflict);
            }
        }
        let mut free = self
            .entries
            .iter()
            .enumerate()
            .filter(|(_, e)| e.is_none())
            .map(|(i, _)| i);
        let ea = free.next().ok_or(Error::Capacity)?;
        let eb = free.next().ok_or(Error::Capacity)?;
        let page = self.new_page()?;
        if let Err(e) = self
            .backend
            .initialize(&self.pages[page].as_ref().unwrap().lease, seed)
        {
            return Err(self.rollback(page, e));
        }
        let frame = self
            .backend
            .reference(&self.pages[page].as_ref().unwrap().lease);
        let value = Some(Value {
            frame,
            permissions: USER_RO,
        });
        let changes = [
            Change {
                owner: a,
                address: address_a,
                before: None,
                after: value,
            },
            Change {
                owner: b,
                address: address_b,
                before: None,
                after: value,
            },
        ];
        let prepared = match self.backend.prepare(&changes) {
            Ok(p) => p,
            Err(e) => return Err(self.rollback(page, e)),
        };
        self.backend.commit(prepared);
        self.pages[page].as_mut().unwrap().references = 2;
        self.entries[ea] = Some(Entry {
            owner: a,
            address: address_a,
            page,
            cow: true,
        });
        self.entries[eb] = Some(Entry {
            owner: b,
            address: address_b,
            page,
            cow: true,
        });
        self.processes[pa].as_mut().unwrap().used += 1;
        self.processes[pb].as_mut().unwrap().used += 1;
        Ok(())
    }
    pub fn reap_quarantine(&mut self) -> Result<usize, Error> {
        let mut count = 0;
        for slot in 0..N {
            if self.pages[slot].as_ref().is_some_and(|p| p.quarantined) {
                if self.pages[slot].as_ref().unwrap().references != 0 {
                    return Err(Error::Backend);
                }
                if self.discard(slot) {
                    return Err(Error::Retained);
                }
                count += 1;
            }
        }
        Ok(count)
    }
    pub fn detach_owner(&mut self, owner: Owner) -> Result<(), Error> {
        let process = self.process(owner)?;
        self.audit()?;
        let mut changes: [Option<Change<B::Reference>>; N] = [None; N];
        let mut count = 0;
        for entry in self.entries.iter().flatten().filter(|e| e.owner == owner) {
            changes[count] = Some(Change {
                owner,
                address: entry.address,
                before: Some(self.checked_value(*entry)?),
                after: None,
            });
            count += 1;
        }
        // A bounded compact change array avoids publishing a partial owner exit.
        let mut compact: [Change<B::Reference>; N] =
            if let Some(first) = changes.first().and_then(|v| *v) {
                [first; N]
            } else {
                if count != 0 {
                    return Err(Error::Backend);
                }
                for region in &mut self.regions {
                    if region.is_some_and(|r| r.owner == owner) {
                        *region = None;
                    }
                }
                self.processes[process] = None;
                return Ok(());
            };
        for i in 0..count {
            compact[i] = changes[i].unwrap();
        }
        let prepared = self.backend.prepare(&compact[..count])?;
        self.backend.commit(prepared);
        for entry in &mut self.entries {
            if let Some(e) = entry.filter(|e| e.owner == owner) {
                self.pages[e.page].as_mut().unwrap().references -= 1;
                *entry = None;
            }
        }
        let mut retained = false;
        for slot in 0..N {
            if self.pages[slot]
                .as_ref()
                .is_some_and(|p| p.references == 0 && !p.quarantined)
                && self.discard(slot)
            {
                retained = true;
            }
        }
        for region in &mut self.regions {
            if region.is_some_and(|r| r.owner == owner) {
                *region = None;
            }
        }
        self.processes[process] = None;
        if retained {
            Err(Error::Retained)
        } else {
            Ok(())
        }
    }
    pub fn audit(&self) -> Result<(), Error> {
        let mut refs = [0usize; N];
        let mut used = [0usize; N];
        for (index, entry) in self
            .entries
            .iter()
            .enumerate()
            .filter_map(|(i, e)| e.map(|e| (i, e)))
        {
            if !entry.address.is_multiple_of(PAGE)
                || entry.address >= 1 << 48
                || self.entries[..index]
                    .iter()
                    .flatten()
                    .any(|e| e.owner == entry.owner && e.address == entry.address)
            {
                return Err(Error::Backend);
            }
            self.checked_value(entry)?;
            let process = self.process(entry.owner)?;
            used[process] = used[process].checked_add(1).ok_or(Error::Capacity)?;
            refs[entry.page] = refs[entry.page].checked_add(1).ok_or(Error::Capacity)?;
        }
        for (i, page) in self
            .pages
            .iter()
            .enumerate()
            .filter_map(|(i, p)| p.as_ref().map(|p| (i, p)))
        {
            self.backend.validate(&page.lease)?;
            if page.references != refs[i]
                || page.quarantined && page.references != 0
                || !page.quarantined && page.references == 0
            {
                return Err(Error::Backend);
            }
            let frame = self.backend.reference(&page.lease);
            let address = self.backend.physical_page(&page.lease);
            if !address.is_multiple_of(PAGE) || address >= 1 << 48 {
                return Err(Error::Backend);
            }
            if self.pages[..i].iter().flatten().any(|p| {
                self.backend.reference(&p.lease) == frame
                    || self.backend.physical_page(&p.lease) == address
            }) {
                return Err(Error::Backend);
            }
        }
        for (i, process) in self
            .processes
            .iter()
            .enumerate()
            .filter_map(|(i, p)| p.map(|p| (i, p)))
        {
            if process.used != used[i] || process.used > process.quota {
                return Err(Error::Backend);
            }
        }
        Ok(())
    }
}