//! Bounded slab/large-object metadata owning every backing lease.
//! The serialized caller chooses task context explicitly. IRQ and fatal paths
//! fail before touching the backing allocator. Metadata construction never
//! allocates; the moved backend owns every physical lease it issues.
use core::alloc::Layout;

pub const PAGE: usize = 4096;
pub const MAX_ORDER: u8 = 3;
pub const CLASSES: [usize; 8] = [16, 32, 64, 128, 256, 512, 1024, 2048];

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    InvalidLayout,
    InvalidOwner,
    Overflow,
    MetadataFull,
    NoMemory,
    WouldBlock,
    ForeignHandle,
    ForeignPointer,
    StaleHandle,
    Backing,
    Audit,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Context {
    Task,
    Interrupt,
    Fatal,
}

/// Only the trusted backing adapter creates this record. Lease identity and
/// successful release are the adapter's responsibility; Heap checks bounds.
pub struct Block<L> {
    pub address: usize,
    pub pages: usize,
    pub lease: L,
}
pub trait Backend {
    type Lease;
    fn acquire(&mut self, order: u8) -> Result<Block<Self::Lease>, Error>;
    /// Failure returns the same still-owned token, never a fabricated token.
    fn relinquish(&mut self, lease: Self::Lease) -> Result<(), (Error, Self::Lease)>;
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
    Slab(usize),
    Large,
    Quarantined,
}
struct Backing<L> {
    block: Block<L>,
    kind: Kind,
}

#[derive(Clone, Copy)]
struct Record {
    live: bool,
    generation: u64,
    owner: u32,
    address: usize,
    requested: usize,
    align: usize,
    capacity: usize,
    backing: usize,
    cell: usize,
}
impl Record {
    const EMPTY: Self = Self {
        live: false,
        generation: 0,
        owner: 0,
        address: 0,
        requested: 0,
        align: 0,
        capacity: 0,
        backing: 0,
        cell: 0,
    };
}

/// Copying a handle cannot copy ownership: every operation validates the live
/// heap/owner/slot/generation ledger. A freed handle is inert.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Handle {
    heap: u64,
    slot: usize,
    generation: u64,
    owner: u32,
    address: usize,
}
impl Handle {
    pub const fn address(self) -> usize {
        self.address
    }
    pub const fn generation(self) -> u64 {
        self.generation
    }
    pub const fn owner(self) -> u32 {
        self.owner
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct LayoutClass {
    pub class: Option<usize>,
    pub order: u8,
    pub capacity: usize,
}
pub fn classify(size: usize, align: usize) -> Result<LayoutClass, Error> {
    if align == 0 || !align.is_power_of_two() {
        return Err(Error::InvalidLayout);
    }
    Layout::from_size_align(size, align).map_err(|_| Error::InvalidLayout)?;
    let demand = size.max(align);
    if let Some(class) = CLASSES.iter().copied().find(|class| *class >= demand) {
        return Ok(LayoutClass {
            class: Some(class),
            order: 0,
            capacity: class,
        });
    }
    let pages = demand.checked_add(PAGE - 1).ok_or(Error::Overflow)? / PAGE;
    let pages = pages.checked_next_power_of_two().ok_or(Error::Overflow)?;
    let order = pages.trailing_zeros() as u8;
    if order > MAX_ORDER {
        return Err(Error::NoMemory);
    }
    Ok(LayoutClass {
        class: None,
        order,
        capacity: pages * PAGE,
    })
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Snapshot {
    pub objects: usize,
    pub requested_bytes: usize,
    pub capacity_bytes: usize,
    pub backing_pages: usize,
    pub slab_pages: usize,
    pub large_pages: usize,
    pub quarantined_pages: usize,
    pub retired_slots: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ObjectRecord {
    pub address: usize,
    pub owner: u32,
    pub generation: u64,
    pub requested: usize,
    pub capacity: usize,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BackingRecord {
    pub address: usize,
    pub pages: usize,
    pub kind: Kind,
}

pub struct Heap<B: Backend, const N: usize, const P: usize> {
    backend: B,
    id: u64,
    objects: [Record; N],
    backing: [Option<Backing<B::Lease>>; P],
}
impl<B: Backend, const N: usize, const P: usize> Heap<B, N, P> {
    /// id is issued uniquely by the runtime owner for this boot. The backend
    /// is moved here; clients cannot independently relinquish its leases.
    pub fn new(backend: B, id: u64) -> Result<Self, Error> {
        if id == 0 || N == 0 || P == 0 {
            return Err(Error::InvalidOwner);
        }
        Ok(Self {
            backend,
            id,
            objects: [Record::EMPTY; N],
            backing: core::array::from_fn(|_| None),
        })
    }
    /// Initialize caller-owned BSS without materializing a large Heap on stack.
    ///
    /// # Safety
    /// storage must be aligned, writable, uninitialized, exclusively owned and
    /// outlive the heap and its handles. The caller serializes all later access.
    pub unsafe fn init_in_place(storage: *mut Self, backend: B, id: u64) -> Result<(), Error> {
        if storage.is_null() || id == 0 || N == 0 || P == 0 {
            return Err(Error::InvalidOwner);
        }
        unsafe {
            core::ptr::addr_of_mut!((*storage).backend).write(backend);
            core::ptr::addr_of_mut!((*storage).id).write(id);
            let objects = core::ptr::addr_of_mut!((*storage).objects).cast::<Record>();
            for index in 0..N {
                objects.add(index).write(Record::EMPTY);
            }
            let backing =
                core::ptr::addr_of_mut!((*storage).backing).cast::<Option<Backing<B::Lease>>>();
            for index in 0..P {
                backing.add(index).write(None);
            }
        }
        Ok(())
    }
    pub fn with_backend_observation<R>(&self, observe: impl FnOnce(&B) -> R) -> R {
        observe(&self.backend)
    }
    /// Return the exclusive backend only after all objects and backing tokens
    /// have been released. A quarantine remains owned and prevents shutdown.
    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        if self.objects.iter().any(|row| row.live) || self.backing.iter().any(Option::is_some) {
            return Err((Error::Backing, self));
        }
        Ok(self.backend)
    }
    pub fn object_records(&self) -> impl Iterator<Item = ObjectRecord> + '_ {
        self.objects
            .iter()
            .filter(|row| row.live)
            .map(|row| ObjectRecord {
                address: row.address,
                owner: row.owner,
                generation: row.generation,
                requested: row.requested,
                capacity: row.capacity,
            })
    }
    pub fn backing_records(&self) -> impl Iterator<Item = BackingRecord> + '_ {
        self.backing.iter().flatten().map(|row| BackingRecord {
            address: row.block.address,
            pages: row.block.pages,
            kind: row.kind,
        })
    }
    pub fn snapshot(&self) -> Snapshot {
        let mut result = Snapshot::default();
        for row in &self.objects {
            if row.live {
                result.objects += 1;
                result.requested_bytes += row.requested;
                result.capacity_bytes += row.capacity;
            } else if row.generation == u64::MAX {
                result.retired_slots += 1;
            }
        }
        for row in self.backing.iter().flatten() {
            result.backing_pages += row.block.pages;
            match row.kind {
                Kind::Slab(_) => result.slab_pages += row.block.pages,
                Kind::Large => result.large_pages += row.block.pages,
                Kind::Quarantined => result.quarantined_pages += row.block.pages,
            }
        }
        result
    }
    fn context(context: Context) -> Result<(), Error> {
        if context == Context::Task {
            Ok(())
        } else {
            Err(Error::WouldBlock)
        }
    }
    fn unused_cell(&self, backing: usize, class: usize) -> Option<usize> {
        (0..PAGE / class).find(|cell| {
            !self
                .objects
                .iter()
                .any(|row| row.live && row.backing == backing && row.cell == *cell)
        })
    }
    fn validate_block(&self, block: &Block<B::Lease>, order: u8) -> Result<(), Error> {
        let bytes = PAGE << order;
        if block.pages != 1usize << order
            || block.address == 0
            || !block.address.is_multiple_of(bytes)
        {
            return Err(Error::Backing);
        }
        let end = block.address.checked_add(bytes).ok_or(Error::Overflow)?;
        if self.backing.iter().flatten().any(|old| {
            block.address < old.block.address + old.block.pages * PAGE && old.block.address < end
        }) {
            return Err(Error::Backing);
        }
        Ok(())
    }
    pub fn allocate(
        &mut self,
        owner: u32,
        size: usize,
        align: usize,
        context: Context,
    ) -> Result<Handle, Error> {
        Self::context(context)?;
        if owner == 0 {
            return Err(Error::InvalidOwner);
        }
        let layout = classify(size, align)?;
        let slot = self
            .objects
            .iter()
            .position(|row| !row.live && row.generation != u64::MAX)
            .ok_or(Error::MetadataFull)?;
        let generation = self.objects[slot].generation + 1;
        let existing = layout.class.and_then(|class| {
            self.backing.iter().enumerate().find_map(|(index, row)| {
                row.as_ref()
                    .filter(|row| row.kind == Kind::Slab(class))
                    .and_then(|_| self.unused_cell(index, class))
                    .map(|cell| (index, cell))
            })
        });
        let (backing, cell) = if let Some(existing) = existing {
            existing
        } else {
            let backing = self
                .backing
                .iter()
                .position(Option::is_none)
                .ok_or(Error::MetadataFull)?;
            let block = self.backend.acquire(layout.order)?;
            if let Err(error) = self.validate_block(&block, layout.order) {
                let Block {
                    address,
                    pages,
                    lease,
                } = block;
                if let Err((_, lease)) = self.backend.relinquish(lease) {
                    self.backing[backing] = Some(Backing {
                        block: Block {
                            address,
                            pages,
                            lease,
                        },
                        kind: Kind::Quarantined,
                    });
                }
                return Err(error);
            }
            self.backing[backing] = Some(Backing {
                block,
                kind: layout.class.map_or(Kind::Large, Kind::Slab),
            });
            (backing, 0)
        };
        let address = self.backing[backing]
            .as_ref()
            .ok_or(Error::Audit)?
            .block
            .address
            + cell * layout.capacity;
        self.objects[slot] = Record {
            live: true,
            generation,
            owner,
            address,
            requested: size,
            align,
            capacity: layout.capacity,
            backing,
            cell,
        };
        Ok(Handle {
            heap: self.id,
            slot,
            generation,
            owner,
            address,
        })
    }
    fn record(&self, handle: Handle, owner: u32) -> Result<Record, Error> {
        if handle.heap != self.id || owner == 0 || owner != handle.owner || handle.slot >= N {
            return Err(Error::ForeignHandle);
        }
        let row = self.objects[handle.slot];
        if !row.live || row.generation != handle.generation {
            return Err(Error::StaleHandle);
        }
        if row.owner != owner || row.address != handle.address {
            return Err(Error::ForeignHandle);
        }
        Ok(row)
    }
    pub fn release_pointer(
        &mut self,
        handle: Handle,
        owner: u32,
        pointer: usize,
        context: Context,
    ) -> Result<(), Error> {
        Self::context(context)?;
        let row = self.record(handle, owner)?;
        if pointer != row.address {
            return Err(Error::ForeignPointer);
        }
        let others = self.objects.iter().any(|other| {
            other.live && other.backing == row.backing && other.address != row.address
        });
        if !others {
            let backing = self.backing[row.backing].take().ok_or(Error::Audit)?;
            let Backing {
                block:
                    Block {
                        address,
                        pages,
                        lease,
                    },
                kind,
            } = backing;
            if let Err((error, lease)) = self.backend.relinquish(lease) {
                self.backing[row.backing] = Some(Backing {
                    block: Block {
                        address,
                        pages,
                        lease,
                    },
                    kind,
                });
                return Err(error);
            }
        }
        self.objects[handle.slot].live = false;
        Ok(())
    }
    pub fn release(&mut self, handle: Handle, owner: u32, context: Context) -> Result<(), Error> {
        self.release_pointer(handle, owner, handle.address, context)
    }
    pub fn audit(&self) -> Result<Snapshot, Error> {
        for (index, row) in self.objects.iter().enumerate().filter(|(_, row)| row.live) {
            let backing = self
                .backing
                .get(row.backing)
                .and_then(Option::as_ref)
                .ok_or(Error::Audit)?;
            if row.generation == 0
                || row.owner == 0
                || row.requested == 0
                || row.requested > row.capacity
                || row.align == 0
                || !row.address.is_multiple_of(row.align)
            {
                return Err(Error::Audit);
            }
            let expected = classify(row.requested, row.align).map_err(|_| Error::Audit)?;
            if row.capacity != expected.capacity
                || backing.block.pages != 1 << expected.order
                || backing.kind != expected.class.map_or(Kind::Large, Kind::Slab)
                || row.address != backing.block.address + row.cell * row.capacity
                || row.address + row.capacity > backing.block.address + backing.block.pages * PAGE
            {
                return Err(Error::Audit);
            }
            for other in self.objects[index + 1..].iter().filter(|other| other.live) {
                if row.address < other.address + other.capacity
                    && other.address < row.address + row.capacity
                {
                    return Err(Error::Audit);
                }
            }
        }
        for (index, row) in self
            .backing
            .iter()
            .enumerate()
            .filter_map(|(index, row)| row.as_ref().map(|row| (index, row)))
        {
            if row.kind != Kind::Quarantined
                && !self
                    .objects
                    .iter()
                    .any(|object| object.live && object.backing == index)
            {
                return Err(Error::Audit);
            }
            if row.kind == Kind::Quarantined
                && self
                    .objects
                    .iter()
                    .any(|object| object.live && object.backing == index)
            {
                return Err(Error::Audit);
            }
        }
        Ok(self.snapshot())
    }
}

#[cfg(test)]
mod generation_tests {
    use super::*;
    struct Never;
    impl Backend for Never {
        type Lease = ();
        fn acquire(&mut self, _: u8) -> Result<Block<()>, Error> {
            panic!("retired slot must not acquire")
        }
        fn relinquish(&mut self, _: ()) -> Result<(), (Error, ())> {
            panic!("retired slot must not release")
        }
    }
    #[test]
    fn generation_wrap_retires_slot_before_physical_acquire() {
        let mut heap = Heap::<_, 1, 1>::new(Never, 1).unwrap();
        heap.objects[0].generation = u64::MAX;
        assert_eq!(
            heap.allocate(1, 24, 8, Context::Task),
            Err(Error::MetadataFull)
        );
        assert_eq!(heap.snapshot().retired_slots, 1);
    }
}