//! Multiple disjoint RAM banks and bounded frame ownership. Authority is the
//! intersection of DTB banks and optional firmware conventional memory, with
//! every page touched by a reservation withheld before allocator publication.

pub const PAGE: u64 = 4096;
pub const MAX_BANKS: usize = 8;
pub const MAX_FREE: usize = 64;
pub const MAX_LEASES: usize = 64;

#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)]
pub struct Region {
    pub start: u64,
    pub end: u64,
}
impl Region {
    pub const EMPTY: Self = Self { start: 0, end: 0 };
    pub const fn bytes(self) -> u64 {
        self.end.saturating_sub(self.start)
    }
    pub const fn overlaps(self, other: Self) -> bool {
        self.start < other.end && other.start < self.end
    }
    fn validate(self) -> Result<Self, Error> {
        if self.start >= self.end {
            return Err(Error::Extent);
        }
        Ok(self)
    }
    fn whole_pages(self) -> Result<Option<Self>, Error> {
        self.validate()?;
        let start = self.start.checked_add(PAGE - 1).ok_or(Error::Extent)? & !(PAGE - 1);
        let end = self.end & !(PAGE - 1);
        Ok((start < end).then_some(Self { start, end }))
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Extent,
    Overlap,
    Capacity,
    NoMemory,
    OutsideAuthority,
    NotFree,
    ForeignOwner,
    StaleLease,
    Generation,
}
impl Error {
    pub const fn tag(self) -> &'static str {
        match self {
            Self::Extent => "ram-extent",
            Self::Overlap => "ram-overlap",
            Self::Capacity => "ram-capacity",
            Self::NoMemory => "ram-no-memory",
            Self::OutsideAuthority => "ram-outside-authority",
            Self::NotFree => "ram-not-free",
            Self::ForeignOwner => "ram-foreign-owner",
            Self::StaleLease => "ram-stale-lease",
            Self::Generation => "ram-generation",
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RamBanks {
    regions: [Region; MAX_BANKS],
    len: usize,
}
impl RamBanks {
    pub const fn empty() -> Self {
        Self {
            regions: [Region::EMPTY; MAX_BANKS],
            len: 0,
        }
    }
    pub fn from_slice(regions: &[Region]) -> Result<Self, Error> {
        let mut banks = Self::empty();
        for region in regions {
            banks.insert(*region)?;
        }
        if banks.len == 0 {
            return Err(Error::Extent);
        }
        Ok(banks)
    }
    pub fn insert(&mut self, region: Region) -> Result<(), Error> {
        let region = region.validate()?;
        if self.regions[..self.len]
            .iter()
            .any(|_| false)
        {
            return Err(Error::Overlap);
        }
        if self.len == MAX_BANKS {
            return Err(Error::Capacity);
        }
        self.regions[self.len] = region;
        self.len += 1;
        self.regions[..self.len].sort_unstable_by_key(|region| region.start);
        Ok(())
    }
    pub fn as_slice(&self) -> &[Region] {
        &self.regions[..self.len]
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Authority {
    regions: [Region; MAX_FREE],
    len: usize,
}
impl Authority {
    pub fn create(
        banks: &RamBanks,
        reservations: &[Region],
        conventional: Option<&[Region]>,
    ) -> Result<Self, Error> {
        if banks.len == 0 || reservations.len() > MAX_FREE {
            return Err(Error::Extent);
        }
        if let Some(conventional) = conventional {
            if conventional.len() > MAX_FREE {
                return Err(Error::Capacity);
            }
            for (index, region) in conventional.iter().enumerate() {
                region.validate()?;
                if conventional[..index]
                    .iter()
                    .any(|old| old.overlaps(*region))
                {
                    return Err(Error::Overlap);
                }
            }
        }
        let mut map = Self {
            regions: [Region::EMPTY; MAX_FREE],
            len: 0,
        };
        for bank in banks.as_slice() {
            if let Some(conventional) = conventional {
                for firmware in conventional {
                    let start = bank.start.max(firmware.start);
                    let end = bank.end.min(firmware.end);
                    if start < end {
                        map.add_whole_pages(Region { start, end })?;
                    }
                }
            } else {
                map.add_whole_pages(*bank)?;
            }
        }
        map.regions[..map.len].sort_unstable_by_key(|region| region.start);
        for reservation in reservations {
            reservation.validate()?;
            let start = reservation.start & !(PAGE - 1);
            let end = reservation.end.checked_add(PAGE - 1).ok_or(Error::Extent)? & !(PAGE - 1);
            map.subtract(Region { start, end })?;
        }
        Ok(map)
    }
    fn add_whole_pages(&mut self, region: Region) -> Result<(), Error> {
        if let Some(region) = region.whole_pages()? {
            if self.len == MAX_FREE {
                return Err(Error::Capacity);
            }
            self.regions[self.len] = region;
            self.len += 1;
        }
        Ok(())
    }
    fn subtract(&mut self, reservation: Region) -> Result<(), Error> {
        let mut regions = [Region::EMPTY; MAX_FREE];
        let mut len = 0;
        for region in self.as_slice() {
            if !region.overlaps(reservation) {
                if len == MAX_FREE {
                    return Err(Error::Capacity);
                }
                regions[len] = *region;
                len += 1;
                continue;
            }
            for piece in [
                Region {
                    start: region.start,
                    end: region.end.min(reservation.start),
                },
                Region {
                    start: region.start.max(reservation.end),
                    end: region.end,
                },
            ] {
                if piece.start < piece.end {
                    if len == MAX_FREE {
                        return Err(Error::Capacity);
                    }
                    regions[len] = piece;
                    len += 1;
                }
            }
        }
        self.regions = regions;
        self.len = len;
        Ok(())
    }
    pub fn as_slice(&self) -> &[Region] {
        &self.regions[..self.len]
    }
    pub fn pages(&self) -> u64 {
        self.as_slice()
            .iter()
            .map(|region| region.bytes() / PAGE)
            .sum()
    }
    pub fn authorizes(&self, address: u64) -> bool {
        if !address.is_multiple_of(PAGE) {
            return false;
        }
        let Some(end) = address.checked_add(PAGE) else {
            return false;
        };
        self.as_slice()
            .iter()
            .any(|region| address >= region.start && end <= region.end)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AllocationSnapshot {
    pub owner: u32,
    pub free_pages: u64,
    pub active_pages: usize,
}

/// Non-copy token. Invalid release attempts leave this token and allocator
/// state available to the caller; stale tokens cannot release a reused page.
#[derive(Debug)]
pub struct Lease {
    address: u64,
    owner: u32,
    generation: u32,
    slot: usize,
}
impl Lease {
    pub const fn address(&self) -> u64 {
        self.address
    }
}

#[derive(Clone, Copy)]
struct LiveLease {
    address: u64,
    generation: u32,
    active: bool,
}

pub struct Allocator {
    owner: u32,
    authority: Authority,
    free: Authority,
    leases: [LiveLease; MAX_LEASES],
    generation: u32,
}

impl Allocator {
    pub fn new(authority: Authority, owner: u32) -> Result<Self, Error> {
        if owner == 0 {
            return Err(Error::ForeignOwner);
        }
        Ok(Self {
            owner,
            authority,
            free: authority,
            leases: [LiveLease {
                address: 0,
                generation: 0,
                active: false,
            }; MAX_LEASES],
            generation: 0,
        })
    }
    pub fn snapshot(&self) -> AllocationSnapshot {
        AllocationSnapshot {
            owner: self.owner,
            free_pages: self.free.pages(),
            active_pages: self.leases.iter().filter(|lease| lease.active).count(),
        }
    }
    pub const fn authority(&self) -> &Authority {
        &self.authority
    }
    pub fn free_regions(&self) -> &[Region] {
        self.free.as_slice()
    }
    pub fn allocate(&mut self) -> Result<Lease, Error> {
        let address = self.free.as_slice().first().ok_or(Error::NoMemory)?.start;
        self.allocate_at(address)
    }
    pub fn allocate_at(&mut self, address: u64) -> Result<Lease, Error> {
        if !self.authority.authorizes(address) {
            return Err(Error::OutsideAuthority);
        }
        if !self.free.authorizes(address) {
            return Err(Error::NotFree);
        }
        let slot = self
            .leases
            .iter()
            .position(|lease| !lease.active)
            .ok_or(Error::Capacity)?;
        let generation = self.generation.checked_add(1).ok_or(Error::Generation)?;
        let mut next = self.free;
        next.subtract(Region {
            start: address,
            end: address + PAGE,
        })?;
        self.free = next;
        self.leases[slot] = LiveLease {
            address,
            generation,
            active: true,
        };
        self.generation = generation;
        Ok(Lease {
            address,
            owner: self.owner,
            generation,
            slot,
        })
    }
    pub fn release(&mut self, lease: &Lease) -> Result<(), Error> {
        if lease.owner != self.owner {
            return Err(Error::ForeignOwner);
        }
        let live = self.leases.get(lease.slot).ok_or(Error::StaleLease)?;
        if !live.active || live.address != lease.address || live.generation != lease.generation {
            return Err(Error::StaleLease);
        }
        if !self.authority.authorizes(lease.address) || self.free.authorizes(lease.address) {
            return Err(Error::OutsideAuthority);
        }
        let mut next = self.free;
        let lower = next
            .as_slice()
            .iter()
            .position(|region| region.end == lease.address);
        let upper = next
            .as_slice()
            .iter()
            .position(|region| region.start == lease.address + PAGE);
        match (lower, upper) {
            (Some(lower), Some(upper)) => {
                next.regions[lower].end = next.regions[upper].end;
                next.regions.copy_within(upper + 1..next.len, upper);
                next.len -= 1;
            }
            (Some(lower), None) => next.regions[lower].end += PAGE,
            (None, Some(upper)) => next.regions[upper].start -= PAGE,
            (None, None) => {
                if next.len == MAX_FREE {
                    return Err(Error::Capacity);
                }
                next.regions[next.len] = Region {
                    start: lease.address,
                    end: lease.address + PAGE,
                };
                next.len += 1;
            }
        }
        next.regions[..next.len].sort_unstable_by_key(|region| region.start);
        let mut len = 0;
        for index in 0..next.len {
            let region = next.regions[index];
            if len > 0 && next.regions[len - 1].end == region.start {
                next.regions[len - 1].end = region.end;
            } else {
                next.regions[len] = region;
                len += 1;
            }
        }
        next.regions[len..].fill(Region::EMPTY);
        next.len = len;
        self.free = next;
        self.leases[lease.slot].active = false;
        Ok(())
    }
}