//! Aligned contiguous blocks over exclusively borrowed central page authority.
//!
//! Caller-owned Workspace belongs in BSS in the guest. Free block records never
//! establish page ownership: the central allocator owns every acquired page and
//! its non-copy generation Lease remains private here. All acquisition/release
//! operations preflight both ledgers before publishing either inventory.

use core::sync::atomic::{AtomicU64, Ordering};
use turkos_os_frame_authority::authority::{
    Allocator, Error as FrameError, Lease, MAX_LIVE, MAX_REGIONS, PAGE,
};
use turkos_os_runtime_platform::ram::Region;

pub const MAX_ORDER: u8 = 3;
pub const MAX_FREE_BLOCKS: usize = 512;
const MAX_INTERNAL_ORDER: u8 = 51;
static NEXT_WORKSPACE: AtomicU64 = AtomicU64::new(1);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    InvalidOrder,
    Misaligned,
    Overflow,
    Capacity,
    NoMemory,
    ForeignOwner,
    StaleLease,
    AlreadyInitialized,
    NotFree,
    NotBuddy,
    Audit,
    Backend(FrameError),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FreeBlock {
    address: u64,
    order: u8,
}
const EMPTY_BLOCK: FreeBlock = FreeBlock {
    address: 0,
    order: 0,
};
impl FreeBlock {
    pub fn new(address: u64, order: u8) -> Result<Self, Error> {
        if order > MAX_INTERNAL_ORDER {
            return Err(Error::InvalidOrder);
        }
        let bytes = PAGE.checked_shl(order as u32).ok_or(Error::Overflow)?;
        if !address.is_multiple_of(bytes) {
            return Err(Error::Misaligned);
        }
        address.checked_add(bytes).ok_or(Error::Overflow)?;
        Ok(Self { address, order })
    }
    pub fn address(self) -> u64 {
        self.address
    }
    pub fn order(self) -> u8 {
        self.order
    }
    pub fn pages(self) -> u64 {
        1u64 << self.order
    }
    pub fn region(self) -> Region {
        Region {
            start: self.address,
            end: self.address + self.pages() * PAGE,
        }
    }
    pub fn split(self) -> Result<(Self, Self), Error> {
        if self.order == 0 {
            return Err(Error::InvalidOrder);
        }
        let order = self.order - 1;
        let first = Self::new(self.address, order)?;
        let second = Self::new(first.region().end, order)?;
        if first.region().end != second.address || second.region().end != self.region().end {
            return Err(Error::Audit);
        }
        Ok((first, second))
    }
    pub fn merge(self, other: Self) -> Result<Self, Error> {
        if self.order != other.order || self == other || self.order == MAX_INTERNAL_ORDER {
            return Err(Error::NotBuddy);
        }
        if self.address ^ (self.pages() * PAGE) != other.address {
            return Err(Error::NotBuddy);
        }
        Self::new(self.address.min(other.address), self.order + 1)
    }
}

/// A non-copy, non-clone token. Addresses and orders cannot be edited by callers.
#[derive(Debug)]
pub struct BlockLease {
    block: FreeBlock,
    owner: u32,
    workspace: u64,
    generation: u64,
    slot: usize,
}
impl BlockLease {
    pub fn address(&self) -> u64 {
        self.block.address
    }
    pub fn order(&self) -> u8 {
        self.block.order
    }
    pub fn pages(&self) -> u64 {
        self.block.pages()
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct LiveBlock {
    block: FreeBlock,
    generation: u64,
    pages: [u8; 8],
    active: bool,
}
const EMPTY_LIVE: LiveBlock = LiveBlock {
    block: EMPTY_BLOCK,
    generation: 0,
    pages: [0; 8],
    active: false,
};

pub struct Workspace {
    initialized: bool,
    owner: u32,
    identity: u64,
    generation: u64,
    pool: [Region; MAX_REGIONS],
    pool_len: usize,
    free: [FreeBlock; MAX_FREE_BLOCKS],
    free_len: usize,
    live: [LiveBlock; MAX_LIVE],
    pages: [Option<Lease>; MAX_LIVE],
}
impl Workspace {
    pub const fn new() -> Self {
        Self {
            initialized: false,
            owner: 0,
            identity: 0,
            generation: 0,
            pool: [Region::EMPTY; MAX_REGIONS],
            pool_len: 0,
            free: [EMPTY_BLOCK; MAX_FREE_BLOCKS],
            free_len: 0,
            live: [EMPTY_LIVE; MAX_LIVE],
            pages: [const { None }; MAX_LIVE],
        }
    }
}
impl Default for Workspace {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Audit {
    pub pool_pages: u64,
    pub free_pages: u64,
    pub live_pages: u64,
    pub live_blocks: usize,
    pub central_free_pages: u64,
    pub central_live_pages: u64,
}

pub struct Buddy<'a> {
    workspace: &'a mut Workspace,
    authority: &'a mut Allocator,
}
impl<'a> Buddy<'a> {
    pub fn initialize(
        workspace: &'a mut Workspace,
        authority: &'a mut Allocator,
        owner: u32,
    ) -> Result<Self, Error> {
        let mut regions = [Region::EMPTY; MAX_REGIONS];
        let len = authority.free_regions().len();
        regions[..len].copy_from_slice(authority.free_regions());
        Self::initialize_pool(workspace, authority, owner, &regions[..len])
    }
    pub fn initialize_region(
        workspace: &'a mut Workspace,
        authority: &'a mut Allocator,
        owner: u32,
        region: Region,
    ) -> Result<Self, Error> {
        Self::initialize_pool(workspace, authority, owner, &[region])
    }
    fn initialize_pool(
        workspace: &'a mut Workspace,
        authority: &'a mut Allocator,
        owner: u32,
        regions: &[Region],
    ) -> Result<Self, Error> {
        if workspace.initialized {
            return Err(Error::AlreadyInitialized);
        }
        if owner == 0 {
            return Err(Error::ForeignOwner);
        }
        authority.audit().map_err(Error::Backend)?;
        if regions.is_empty() {
            return Err(Error::NoMemory);
        }
        if regions.len() > MAX_REGIONS {
            return Err(Error::Capacity);
        }
        let mut count = 0;
        for (index, region) in regions.iter().enumerate() {
            if region.start >= region.end
                || !region.start.is_multiple_of(PAGE)
                || !region.end.is_multiple_of(PAGE)
            {
                return Err(Error::Misaligned);
            }
            if !authority
                .free_regions()
                .iter()
                .any(|free| free.start <= region.start && free.end >= region.end)
                || regions[..index].iter().any(|prior| prior.overlaps(*region))
            {
                return Err(Error::NotFree);
            }
            let mut address = region.start;
            while address < region.end {
                let block = seed_block(address, region.end)?;
                address = block.region().end;
                count += 1;
                if count > MAX_FREE_BLOCKS {
                    return Err(Error::Capacity);
                }
            }
        }
        let identity = NEXT_WORKSPACE
            .try_update(Ordering::Relaxed, Ordering::Relaxed, |next| {
                next.checked_add(1)
            })
            .map_err(|_| Error::Overflow)?;
        workspace.pool.fill(Region::EMPTY);
        workspace.pool[..regions.len()].copy_from_slice(regions);
        workspace.pool_len = regions.len();
        workspace.free.fill(EMPTY_BLOCK);
        workspace.free_len = 0;
        for region in regions {
            let mut address = region.start;
            while address < region.end {
                let block = seed_block(address, region.end)?;
                address = block.region().end;
                workspace.free[workspace.free_len] = block;
                workspace.free_len += 1;
            }
        }
        workspace.live.fill(EMPTY_LIVE);
        for page in &mut workspace.pages {
            *page = None;
        }
        workspace.owner = owner;
        workspace.identity = identity;
        workspace.generation = 0;
        workspace.initialized = true;
        Ok(Self {
            workspace,
            authority,
        })
    }
    pub fn resume(
        workspace: &'a mut Workspace,
        authority: &'a mut Allocator,
        owner: u32,
    ) -> Result<Self, Error> {
        if !workspace.initialized || workspace.owner != owner {
            return Err(Error::ForeignOwner);
        }
        let buddy = Self {
            workspace,
            authority,
        };
        buddy.audit()?;
        Ok(buddy)
    }
    pub fn free_blocks(&self) -> &[FreeBlock] {
        &self.workspace.free[..self.workspace.free_len]
    }
    /// Observe the central ledger without allowing a mutation behind Buddy's
    /// exclusively borrowed page inventory. References cannot escape the call.
    pub fn with_authority_observation<R>(&self, callback: impl FnOnce(&Allocator) -> R) -> R {
        callback(self.authority)
    }
    pub fn live_blocks(&self) -> impl Iterator<Item = (u64, u8)> + '_ {
        self.workspace
            .live
            .iter()
            .filter(|block| block.active)
            .map(|block| (block.block.address, block.block.order))
    }
    fn contains_pool(&self, block: FreeBlock) -> bool {
        self.workspace.pool[..self.workspace.pool_len]
            .iter()
            .any(|region| region.start <= block.address && region.end >= block.region().end)
    }
    pub fn split_free(&mut self, parent: FreeBlock) -> Result<(FreeBlock, FreeBlock), Error> {
        self.audit()?;
        let index = self
            .free_blocks()
            .iter()
            .position(|block| *block == parent)
            .ok_or(Error::NotFree)?;
        if self.workspace.free_len == MAX_FREE_BLOCKS {
            return Err(Error::Capacity);
        }
        let (first, second) = parent.split()?;
        self.remove_free(index);
        self.push_free(first);
        self.push_free(second);
        Ok((first, second))
    }
    pub fn merge_free(&mut self, first: FreeBlock, second: FreeBlock) -> Result<FreeBlock, Error> {
        self.audit()?;
        let parent = first.merge(second)?;
        if !self.contains_pool(parent) {
            return Err(Error::NotBuddy);
        }
        let a = self
            .free_blocks()
            .iter()
            .position(|block| *block == first)
            .ok_or(Error::NotFree)?;
        let b = self
            .free_blocks()
            .iter()
            .position(|block| *block == second)
            .ok_or(Error::NotFree)?;
        self.remove_free(a.max(b));
        self.remove_free(a.min(b));
        self.push_free(parent);
        Ok(parent)
    }
    pub fn acquire(&mut self, order: u8) -> Result<BlockLease, Error> {
        if order > MAX_ORDER {
            return Err(Error::InvalidOrder);
        }
        self.audit()?;
        let candidate = self
            .free_blocks()
            .iter()
            .filter(|block| block.order >= order)
            .min_by_key(|block| (block.order, block.address))
            .copied()
            .ok_or(Error::NoMemory)?;
        self.acquire_at(candidate.address, order)
    }
    pub fn acquire_at(&mut self, address: u64, order: u8) -> Result<BlockLease, Error> {
        if order > MAX_ORDER {
            return Err(Error::InvalidOrder);
        }
        let wanted = FreeBlock::new(address, order)?;
        self.audit()?;
        let slot = self
            .workspace
            .live
            .iter()
            .position(|block| !block.active)
            .ok_or(Error::Capacity)?;
        let index = self
            .free_blocks()
            .iter()
            .position(|block| {
                block.address <= address
                    && block.region().end >= wanted.region().end
                    && block.order >= order
            })
            .ok_or(Error::NoMemory)?;
        let mut selected = self.workspace.free[index];
        let mut siblings = [EMPTY_BLOCK; MAX_INTERNAL_ORDER as usize];
        let mut split_len = 0;
        while selected.order > order {
            let (first, second) = selected.split()?;
            if address < second.address {
                selected = first;
                siblings[split_len] = second;
            } else {
                selected = second;
                siblings[split_len] = first;
            }
            split_len += 1;
        }
        if selected != wanted {
            return Err(Error::Audit);
        }
        if self.workspace.free_len - 1 + split_len > MAX_FREE_BLOCKS {
            return Err(Error::Capacity);
        }
        let generation = self
            .workspace
            .generation
            .checked_add(1)
            .ok_or(Error::Overflow)?;
        let count = wanted.pages() as usize;
        let mut page_slots = [0u8; 8];
        let mut found = 0;
        for (index, page) in self.workspace.pages.iter().enumerate() {
            if page.is_none() && found < count {
                page_slots[found] = index as u8;
                found += 1;
            }
        }
        if found != count {
            return Err(Error::Capacity);
        }
        // Central authority preflights all pages, free-map edits and generations
        // before returning any token. Every later metadata edit is infallible.
        let mut leases = self
            .authority
            .allocate_many_at(address, count)
            .map_err(Error::Backend)?;
        self.remove_free(index);
        for sibling in &siblings[..split_len] {
            self.push_free(*sibling);
        }
        for (index, page_slot) in page_slots[..count].iter().enumerate() {
            self.workspace.pages[*page_slot as usize] = leases[index].take();
        }
        self.workspace.live[slot] = LiveBlock {
            block: wanted,
            generation,
            pages: page_slots,
            active: true,
        };
        self.workspace.generation = generation;
        Ok(BlockLease {
            block: wanted,
            owner: self.workspace.owner,
            workspace: self.workspace.identity,
            generation,
            slot,
        })
    }
    pub fn relinquish(&mut self, lease: BlockLease) -> Result<(), (Error, BlockLease)> {
        match self.release_checked(&lease) {
            Ok(()) => Ok(()),
            Err(error) => Err((error, lease)),
        }
    }
    fn release_checked(&mut self, lease: &BlockLease) -> Result<(), Error> {
        self.audit()?;
        if lease.owner != self.workspace.owner {
            return Err(Error::ForeignOwner);
        }
        let record = *self
            .workspace
            .live
            .get(lease.slot)
            .ok_or(Error::StaleLease)?;
        if !record.active || record.block != lease.block || record.generation != lease.generation {
            return Err(Error::StaleLease);
        }
        let mut parent = record.block;
        let mut removals = [0usize; MAX_INTERNAL_ORDER as usize];
        let mut len = 0;
        while parent.order < MAX_INTERNAL_ORDER {
            let buddy_address = parent.address ^ (parent.pages() * PAGE);
            let Some(index) = self
                .free_blocks()
                .iter()
                .position(|block| block.order == parent.order && block.address == buddy_address)
            else {
                break;
            };
            let combined = parent.merge(self.workspace.free[index])?;
            if !self.contains_pool(combined) {
                break;
            }
            removals[len] = index;
            len += 1;
            parent = combined;
        }
        if self.workspace.free_len - len + 1 > MAX_FREE_BLOCKS {
            return Err(Error::Capacity);
        }
        let count = record.block.pages() as usize;
        let first = self.workspace.pages[record.pages[0] as usize]
            .as_ref()
            .ok_or(Error::Audit)?;
        let mut pages = [first; 8];
        for (index, slot) in record.pages[..count].iter().enumerate() {
            pages[index] = self.workspace.pages[*slot as usize]
                .as_ref()
                .ok_or(Error::Audit)?;
        }
        self.authority
            .release_many_leases(&pages[..count])
            .map_err(Error::Backend)?;
        removals[..len].sort_unstable();
        for index in removals[..len].iter().rev() {
            self.remove_free(*index);
        }
        self.push_free(parent);
        for slot in &record.pages[..count] {
            self.workspace.pages[*slot as usize] = None;
        }
        self.workspace.live[lease.slot] = EMPTY_LIVE;
        Ok(())
    }
    fn remove_free(&mut self, index: usize) {
        self.workspace
            .free
            .copy_within(index + 1..self.workspace.free_len, index);
        self.workspace.free_len -= 1;
        self.workspace.free[self.workspace.free_len] = EMPTY_BLOCK;
    }
    fn push_free(&mut self, block: FreeBlock) {
        self.workspace.free[self.workspace.free_len] = block;
        self.workspace.free_len += 1;
    }
    pub fn close(self) -> Result<(), (Error, Self)> {
        let audit = match self.audit() {
            Ok(value) => value,
            Err(error) => return Err((error, self)),
        };
        if audit.live_pages != 0 {
            return Err((Error::NotFree, self));
        }
        self.workspace.initialized = false;
        self.workspace.owner = 0;
        self.workspace.identity = 0;
        self.workspace.pool.fill(Region::EMPTY);
        self.workspace.pool_len = 0;
        self.workspace.free.fill(EMPTY_BLOCK);
        self.workspace.free_len = 0;
        Ok(())
    }
    pub fn audit(&self) -> Result<Audit, Error> {
        if !self.workspace.initialized
            || self.workspace.owner == 0
            || self.workspace.identity == 0
            || self.workspace.pool_len > MAX_REGIONS
            || self.workspace.free_len > MAX_FREE_BLOCKS
        {
            return Err(Error::Audit);
        }
        let central = self.authority.audit().map_err(Error::Backend)?;
        let pool = &self.workspace.pool[..self.workspace.pool_len];
        let mut pool_pages = 0u64;
        let mut observed_free = 0u64;
        for (index, region) in pool.iter().enumerate() {
            if region.start >= region.end
                || !region.start.is_multiple_of(PAGE)
                || !region.end.is_multiple_of(PAGE)
                || pool[..index].iter().any(|prior| prior.overlaps(*region))
            {
                return Err(Error::Audit);
            }
            pool_pages = pool_pages
                .checked_add((region.end - region.start) / PAGE)
                .ok_or(Error::Overflow)?;
            for free in self.authority.free_regions() {
                let start = region.start.max(free.start);
                let end = region.end.min(free.end);
                if start < end {
                    observed_free += (end - start) / PAGE;
                }
            }
        }
        let mut free_pages = 0u64;
        for (index, block) in self.free_blocks().iter().enumerate() {
            FreeBlock::new(block.address, block.order)?;
            if !self.contains_pool(*block)
                || self.free_blocks()[..index]
                    .iter()
                    .any(|prior| prior.region().overlaps(block.region()))
                || !self
                    .authority
                    .free_regions()
                    .iter()
                    .any(|free| free.start <= block.address && free.end >= block.region().end)
            {
                return Err(Error::Audit);
            }
            free_pages = free_pages
                .checked_add(block.pages())
                .ok_or(Error::Overflow)?;
        }
        if self.workspace.free[self.workspace.free_len..]
            .iter()
            .any(|block| *block != EMPTY_BLOCK)
        {
            return Err(Error::Audit);
        }
        let mut live_pages = 0u64;
        let mut live_blocks = 0;
        let mut seen = [false; MAX_LIVE];
        for (index, record) in self.workspace.live.iter().enumerate() {
            if !record.active {
                if *record != EMPTY_LIVE {
                    return Err(Error::Audit);
                }
                continue;
            }
            if record.generation == 0
                || record.generation > self.workspace.generation
                || record.block.order > MAX_ORDER
                || !self.contains_pool(record.block)
                || self.workspace.live[..index]
                    .iter()
                    .any(|old| old.active && old.block.region().overlaps(record.block.region()))
                || self
                    .free_blocks()
                    .iter()
                    .any(|free| free.region().overlaps(record.block.region()))
            {
                return Err(Error::Audit);
            }
            FreeBlock::new(record.block.address, record.block.order)?;
            let count = record.block.pages() as usize;
            for (page, slot) in record.pages[..count].iter().enumerate() {
                let slot = *slot as usize;
                if slot >= MAX_LIVE || seen[slot] {
                    return Err(Error::Audit);
                }
                let lease = self.workspace.pages[slot].as_ref().ok_or(Error::Audit)?;
                if lease.address() != record.block.address + page as u64 * PAGE {
                    return Err(Error::Audit);
                }
                self.authority
                    .validate_lease(lease)
                    .map_err(Error::Backend)?;
                seen[slot] = true;
            }
            if record.pages[count..].iter().any(|slot| *slot != 0) {
                return Err(Error::Audit);
            }
            live_pages += count as u64;
            live_blocks += 1;
        }
        for (index, lease) in self.workspace.pages.iter().enumerate() {
            if seen[index] != lease.is_some() {
                return Err(Error::Audit);
            }
        }
        let observed_live = self
            .authority
            .live_addresses()
            .filter(|address| {
                pool.iter()
                    .any(|region| *address >= region.start && *address < region.end)
            })
            .count() as u64;
        if free_pages != observed_free
            || live_pages != observed_live
            || free_pages + live_pages != pool_pages
        {
            return Err(Error::Audit);
        }
        Ok(Audit {
            pool_pages,
            free_pages,
            live_pages,
            live_blocks,
            central_free_pages: central.free_pages,
            central_live_pages: central.live_pages,
        })
    }
}
fn seed_block(address: u64, end: u64) -> Result<FreeBlock, Error> {
    for order in (0..=MAX_INTERNAL_ORDER).rev() {
        if let Ok(block) = FreeBlock::new(address, order)
            && block.region().end <= end
        {
            return Ok(block);
        }
    }
    Err(Error::Misaligned)
}