pub const PAGE_BYTES: usize = 4096;
pub const MAX_PINS: u8 = 8;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Owner,
    Capacity,
    NoMemory,
    Backend,
    BadAddress,
    DuplicateAddress,
    Erasure,
    Foreign,
    Stale,
    Pinned,
    PinCapacity,
    PinConsumed,
    Generation,
}

/// The backend exclusively owns frame authority and keeps allocated pages
/// mapped until release. read_byte must read the real backing storage; an
/// implementation which invents zeros cannot establish the erasure contract.
pub trait Backend {
    fn allocate(&mut self) -> Result<u64, Error>;
    fn release(&mut self, address: u64) -> Result<(), Error>;
    fn zero_page(&mut self, address: u64) -> Result<(), Error>;
    fn read_byte(&self, address: u64, offset: usize) -> Result<u8, Error>;
    fn available_pages(&self) -> u64;
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Handle {
    registry: u64,
    owner: u64,
    address: u64,
    slot: usize,
    generation: u64,
}
impl Handle {
    pub const fn address(self) -> u64 {
        self.address
    }
    pub const fn generation(self) -> u64 {
        self.generation
    }
}

/// This token cannot be cloned. A successfully consumed token stays consumed
/// even if its frame slot is reused with a new generation.
#[derive(Debug)]
pub struct Pin {
    handle: Handle,
    consumed: bool,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
    Free,
    Live,
    Quarantined,
    Retired,
}
#[derive(Clone, Copy, Debug)]
struct Record {
    state: State,
    generation: u64,
    address: u64,
    owner: u64,
    pins: u8,
}
const EMPTY: Record = Record {
    state: State::Free,
    generation: 0,
    address: 0,
    owner: 0,
    pins: 0,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot {
    pub free_metadata: usize,
    pub retired_metadata: usize,
    pub live_pages: usize,
    pub quarantined_pages: usize,
    pub pins: usize,
    pub backend_free_pages: u64,
}

/// Inert ownership information for a parent authority audit. This is not a
/// handle and cannot be used to release or pin any page.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OwnedRecord {
    pub address: u64,
    pub owner: u64,
    pub generation: u64,
    pub pins: u8,
    pub quarantined: bool,
}

/// registry_id is issued once by the boot's frame-authority owner and must
/// uniquely identify that authority. This registry takes ownership of B;
/// callers cannot reach around its records to free or allocate live backing.
pub struct Registry<B, const N: usize> {
    id: u64,
    backend: B,
    records: [Record; N],
}
impl<B: Backend, const N: usize> Registry<B, N> {
    pub fn new(registry_id: u64, backend: B) -> Result<Self, Error> {
        if registry_id == 0 {
            return Err(Error::Owner);
        }
        Ok(Self {
            id: registry_id,
            backend,
            records: [EMPTY; N],
        })
    }

    pub fn snapshot(&self) -> Snapshot {
        Snapshot {
            free_metadata: self
                .records
                .iter()
                .filter(|r| r.state == State::Free)
                .count(),
            retired_metadata: self
                .records
                .iter()
                .filter(|r| r.state == State::Retired)
                .count(),
            live_pages: self
                .records
                .iter()
                .filter(|r| r.state == State::Live)
                .count(),
            quarantined_pages: self
                .records
                .iter()
                .filter(|r| r.state == State::Quarantined)
                .count(),
            pins: self.records.iter().map(|r| usize::from(r.pins)).sum(),
            backend_free_pages: self.backend.available_pages(),
        }
    }

    pub fn owned_records(&self) -> impl Iterator<Item = OwnedRecord> + '_ {
        self.records.iter().filter_map(|record| {
            matches!(record.state, State::Live | State::Quarantined).then_some(OwnedRecord {
                address: record.address,
                owner: record.owner,
                generation: record.generation,
                pins: record.pins,
                quarantined: record.state == State::Quarantined,
            })
        })
    }

    pub fn allocate_zeroed(&mut self, owner: u64) -> Result<Handle, Error> {
        if owner == 0 {
            return Err(Error::Owner);
        }
        // Find metadata first: ENOSPC must never acquire a physical frame.
        let slot = self
            .records
            .iter()
            .position(|r| r.state == State::Free)
            .ok_or(Error::Capacity)?;
        let generation = match self.records[slot].generation.checked_add(1) {
            Some(value) => value,
            None => {
                self.records[slot].state = State::Retired;
                return Err(Error::Generation);
            }
        };
        let address = self.backend.allocate()?;
        // An overlapping result from a broken backend must never be freed:
        // it might still belong to an existing live lease. Retain authority.
        if self
            .records
            .iter()
            .any(|r| matches!(r.state, State::Live | State::Quarantined) && r.address == address)
        {
            return Err(Error::DuplicateAddress);
        }
        let failure = if address & (PAGE_BYTES as u64 - 1) != 0
            || address.checked_add(PAGE_BYTES as u64).is_none()
        {
            Some(Error::BadAddress)
        } else if self.backend.zero_page(address).is_err() {
            Some(Error::Backend)
        } else {
            let mut error = None;
            for offset in 0..PAGE_BYTES {
                match self.backend.read_byte(address, offset) {
                    Ok(0) => {}
                    Ok(_) => {
                        error = Some(Error::Erasure);
                        break;
                    }
                    Err(_) => {
                        error = Some(Error::Backend);
                        break;
                    }
                }
            }
            error
        };
        if let Some(error) = failure {
            if self.backend.release(address).is_err() {
                self.records[slot] = Record {
                    state: State::Quarantined,
                    generation,
                    address,
                    owner,
                    pins: 0,
                };
            }
            return Err(error);
        }
        let handle = Handle {
            registry: self.id,
            owner,
            address,
            slot,
            generation,
        };
        self.records[slot] = Record {
            state: State::Live,
            generation,
            address,
            owner,
            pins: 0,
        };
        Ok(handle)
    }

    fn record(&self, handle: Handle, owner: u64) -> Result<&Record, Error> {
        if handle.registry != self.id || handle.owner != owner || owner == 0 {
            return Err(Error::Foreign);
        }
        let record = self.records.get(handle.slot).ok_or(Error::Stale)?;
        if record.state != State::Live
            || record.generation != handle.generation
            || record.address != handle.address
            || record.owner != owner
        {
            return Err(Error::Stale);
        }
        Ok(record)
    }

    pub fn read_byte(&self, handle: Handle, owner: u64, offset: usize) -> Result<u8, Error> {
        self.record(handle, owner)?;
        if offset >= PAGE_BYTES {
            return Err(Error::BadAddress);
        }
        self.backend.read_byte(handle.address, offset)
    }

    pub fn release(&mut self, handle: Handle, owner: u64) -> Result<(), Error> {
        if false && self.record(handle, owner)?.pins != 0 {
            return Err(Error::Pinned);
        }
        // A failed backend release retains the complete live record.
        self.backend.release(handle.address)?;
        self.records[handle.slot].state = State::Free;
        self.records[handle.slot].address = 0;
        self.records[handle.slot].owner = 0;
        Ok(())
    }

    pub fn pin(&mut self, handle: Handle, owner: u64) -> Result<Pin, Error> {
        if self.record(handle, owner)?.pins == MAX_PINS {
            return Err(Error::PinCapacity);
        }
        self.records[handle.slot].pins += 1;
        Ok(Pin {
            handle,
            consumed: false,
        })
    }

    pub fn unpin(&mut self, pin: &mut Pin, owner: u64) -> Result<(), Error> {
        if pin.consumed {
            return Err(Error::PinConsumed);
        }
        if self.record(pin.handle, owner)?.pins == 0 {
            return Err(Error::PinConsumed);
        }
        self.records[pin.handle.slot].pins -= 1;
        pin.consumed = true;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    struct NoAllocation;
    impl Backend for NoAllocation {
        fn allocate(&mut self) -> Result<u64, Error> {
            panic!("generation wrap must not acquire a frame")
        }
        fn release(&mut self, _: u64) -> Result<(), Error> {
            panic!("no frame exists")
        }
        fn zero_page(&mut self, _: u64) -> Result<(), Error> {
            panic!("no frame exists")
        }
        fn read_byte(&self, _: u64, _: usize) -> Result<u8, Error> {
            panic!("no frame exists")
        }
        fn available_pages(&self) -> u64 {
            4
        }
    }
    #[test]
    fn exhausted_generation_retires_metadata_without_physical_acquisition() {
        let mut registry = Registry::<_, 1>::new(1, NoAllocation).unwrap();
        registry.records[0].generation = u64::MAX;
        assert_eq!(registry.allocate_zeroed(17), Err(Error::Generation));
        let snapshot = registry.snapshot();
        assert_eq!(snapshot.retired_metadata, 1);
        assert_eq!(snapshot.backend_free_pages, 4);
        assert_eq!(snapshot.live_pages, 0);
        assert_eq!(registry.allocate_zeroed(18), Err(Error::Capacity));
        assert_eq!(registry.snapshot(), snapshot);
    }
}