lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! Block allocator for thin provisioning.
//!
//! This module provides the block allocation and deallocation logic
//! for thin pools, including free space tracking and allocation policies.

use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec;
use alloc::vec::Vec;

use lazy_static::lazy_static;
use spin::Mutex;

use super::types::{AllocationPolicy, PhysicalBlock, ThinError, ThinResult};

// ═══════════════════════════════════════════════════════════════════════════════
// FREE SPACE BITMAP
// ═══════════════════════════════════════════════════════════════════════════════

/// Bitmap for tracking free blocks.
#[derive(Debug)]
pub struct FreeBitmap {
    /// Bitmap data (1 = free, 0 = allocated).
    bits: Vec<u64>,
    /// Total number of blocks.
    total_blocks: u64,
    /// Number of free blocks.
    free_count: u64,
    /// First potentially free block (optimization).
    first_free_hint: u64,
}

impl FreeBitmap {
    /// Create a new bitmap with all blocks free.
    pub fn new(total_blocks: u64) -> Self {
        let words = total_blocks.div_ceil(64) as usize;
        let mut bits = vec![u64::MAX; words];

        // Clear bits beyond total_blocks
        let remainder = (total_blocks % 64) as u32;
        if remainder > 0 && !bits.is_empty() {
            let last_idx = bits.len() - 1;
            bits[last_idx] = (1u64 << remainder) - 1;
        }

        Self {
            bits,
            total_blocks,
            free_count: total_blocks,
            first_free_hint: 0,
        }
    }

    /// Check if a block is free.
    pub fn is_free(&self, block: u64) -> bool {
        if block >= self.total_blocks {
            return false;
        }
        let word = (block / 64) as usize;
        let bit = (block % 64) as u32;
        (self.bits[word] & (1u64 << bit)) != 0
    }

    /// Mark a block as allocated.
    pub fn allocate(&mut self, block: u64) -> bool {
        if block >= self.total_blocks {
            return false;
        }
        let word = (block / 64) as usize;
        let bit = (block % 64) as u32;

        if (self.bits[word] & (1u64 << bit)) != 0 {
            self.bits[word] &= !(1u64 << bit);
            self.free_count -= 1;

            // Update hint if needed
            if block == self.first_free_hint {
                self.update_first_free_hint(block);
            }
            true
        } else {
            false
        }
    }

    /// Mark a block as free.
    pub fn free(&mut self, block: u64) {
        if block >= self.total_blocks {
            return;
        }
        let word = (block / 64) as usize;
        let bit = (block % 64) as u32;

        if (self.bits[word] & (1u64 << bit)) == 0 {
            self.bits[word] |= 1u64 << bit;
            self.free_count += 1;

            // Update hint
            if block < self.first_free_hint {
                self.first_free_hint = block;
            }
        }
    }

    /// Find the first free block.
    pub fn find_first_free(&self) -> Option<u64> {
        // Start from hint
        let start_word = (self.first_free_hint / 64) as usize;

        for (i, &word) in self.bits.iter().enumerate().skip(start_word) {
            if word != 0 {
                let bit = word.trailing_zeros() as u64;
                let block = (i as u64) * 64 + bit;
                if block < self.total_blocks {
                    return Some(block);
                }
            }
        }
        None
    }

    /// Find a contiguous range of free blocks.
    pub fn find_contiguous(&self, count: u64) -> Option<u64> {
        if count == 0 || count > self.free_count {
            return None;
        }

        let mut start = 0u64;
        let mut run = 0u64;

        for block in 0..self.total_blocks {
            if self.is_free(block) {
                if run == 0 {
                    start = block;
                }
                run += 1;
                if run >= count {
                    return Some(start);
                }
            } else {
                run = 0;
            }
        }
        None
    }

    /// Allocate a range of blocks.
    pub fn allocate_range(&mut self, start: u64, count: u64) -> bool {
        // Verify all blocks are free
        for i in 0..count {
            if !self.is_free(start + i) {
                return false;
            }
        }

        // Allocate them
        for i in 0..count {
            self.allocate(start + i);
        }
        true
    }

    /// Free a range of blocks.
    pub fn free_range(&mut self, start: u64, count: u64) {
        for i in 0..count {
            self.free(start + i);
        }
    }

    /// Get the number of free blocks.
    pub fn free_count(&self) -> u64 {
        self.free_count
    }

    /// Get the total number of blocks.
    pub fn total_blocks(&self) -> u64 {
        self.total_blocks
    }

    /// Get usage percentage.
    pub fn usage_percent(&self) -> u8 {
        if self.total_blocks == 0 {
            return 0;
        }
        let used = self.total_blocks - self.free_count;
        ((used as u128 * 100) / self.total_blocks as u128) as u8
    }

    /// Update the first free hint.
    fn update_first_free_hint(&mut self, from: u64) {
        let start_word = (from / 64) as usize;
        for (i, &word) in self.bits.iter().enumerate().skip(start_word) {
            if word != 0 {
                self.first_free_hint = (i as u64) * 64 + word.trailing_zeros() as u64;
                return;
            }
        }
        self.first_free_hint = self.total_blocks; // No free blocks
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// EXTENT TREE
// ═══════════════════════════════════════════════════════════════════════════════

/// A contiguous extent of blocks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Extent {
    /// Starting block.
    pub start: u64,
    /// Number of blocks.
    pub count: u64,
}

impl Extent {
    /// Create a new extent.
    pub fn new(start: u64, count: u64) -> Self {
        Self { start, count }
    }

    /// Get the end block (exclusive).
    pub fn end(&self) -> u64 {
        self.start + self.count
    }

    /// Check if this extent contains a block.
    pub fn contains(&self, block: u64) -> bool {
        block >= self.start && block < self.end()
    }

    /// Check if this extent overlaps with another.
    pub fn overlaps(&self, other: &Extent) -> bool {
        self.start < other.end() && other.start < self.end()
    }

    /// Check if this extent is adjacent to another.
    pub fn adjacent(&self, other: &Extent) -> bool {
        self.end() == other.start || other.end() == self.start
    }

    /// Merge with an adjacent extent.
    pub fn merge(&self, other: &Extent) -> Option<Extent> {
        if self.end() == other.start {
            Some(Extent::new(self.start, self.count + other.count))
        } else if other.end() == self.start {
            Some(Extent::new(other.start, self.count + other.count))
        } else {
            None
        }
    }
}

/// Tree of free extents for efficient allocation.
#[derive(Debug)]
pub struct ExtentTree {
    /// Extents indexed by start block.
    by_start: BTreeMap<u64, Extent>,
    /// Extent starts indexed by size (for best-fit).
    by_size: BTreeMap<u64, BTreeSet<u64>>,
    /// Total free blocks.
    free_blocks: u64,
}

impl ExtentTree {
    /// Create an empty extent tree.
    pub fn new() -> Self {
        Self {
            by_start: BTreeMap::new(),
            by_size: BTreeMap::new(),
            free_blocks: 0,
        }
    }

    /// Create with initial extent.
    pub fn with_extent(start: u64, count: u64) -> Self {
        let mut tree = Self::new();
        tree.add_extent(Extent::new(start, count));
        tree
    }

    /// Add a free extent.
    pub fn add_extent(&mut self, extent: Extent) {
        // Try to merge with adjacent extents
        let mut merged = extent;
        let mut added_blocks = extent.count;

        // Check for preceding extent - extract data first
        let prev_info: Option<(u64, u64)> = self
            .by_start
            .range(..extent.start)
            .next_back()
            .filter(|(_, prev)| prev.end() == extent.start)
            .map(|(&start, ext)| (start, ext.count));

        if let Some((prev_start, prev_count)) = prev_info {
            // Remove the extent but don't subtract from free_blocks (we'll add correctly at the end)
            if let Some(ext) = self.by_start.remove(&prev_start) {
                if let Some(set) = self.by_size.get_mut(&ext.count) {
                    set.remove(&prev_start);
                    if set.is_empty() {
                        self.by_size.remove(&ext.count);
                    }
                }
            }
            merged = Extent::new(prev_start, prev_count + merged.count);
        }

        // Check for following extent - extract data first
        let next_info: Option<(u64, u64)> = self
            .by_start
            .range(merged.end()..)
            .next()
            .filter(|(_, next)| merged.end() == next.start)
            .map(|(&start, ext)| (start, ext.count));

        if let Some((next_start, next_count)) = next_info {
            // Remove the extent but don't subtract from free_blocks
            if let Some(ext) = self.by_start.remove(&next_start) {
                if let Some(set) = self.by_size.get_mut(&ext.count) {
                    set.remove(&next_start);
                    if set.is_empty() {
                        self.by_size.remove(&ext.count);
                    }
                }
            }
            merged = Extent::new(merged.start, merged.count + next_count);
        }

        // Insert merged extent
        self.by_start.insert(merged.start, merged);
        self.by_size
            .entry(merged.count)
            .or_default()
            .insert(merged.start);
        // Only add the new blocks (merged extents already counted)
        self.free_blocks += added_blocks;
    }

    /// Remove an extent by its start.
    fn remove_extent(&mut self, start: u64) -> Option<Extent> {
        if let Some(extent) = self.by_start.remove(&start) {
            if let Some(set) = self.by_size.get_mut(&extent.count) {
                set.remove(&start);
                if set.is_empty() {
                    self.by_size.remove(&extent.count);
                }
            }
            self.free_blocks -= extent.count;
            Some(extent)
        } else {
            None
        }
    }

    /// Allocate blocks using first-fit.
    pub fn allocate_first_fit(&mut self, count: u64) -> Option<u64> {
        // Find first extent large enough
        let found: Option<(u64, Extent)> = self
            .by_start
            .iter()
            .find(|(_, extent)| extent.count >= count)
            .map(|(&start, &extent)| (start, extent));

        if let Some((start, extent)) = found {
            self.remove_extent(start);

            // If extent is larger, add remainder back
            if extent.count > count {
                // Directly insert the remainder without going through add_extent
                // to avoid double-counting issues
                let remainder = Extent::new(start + count, extent.count - count);
                self.by_start.insert(remainder.start, remainder);
                self.by_size
                    .entry(remainder.count)
                    .or_default()
                    .insert(remainder.start);
                self.free_blocks += remainder.count;
            }

            return Some(start);
        }
        None
    }

    /// Allocate blocks using best-fit.
    pub fn allocate_best_fit(&mut self, count: u64) -> Option<u64> {
        // Find smallest extent that's large enough
        let found: Option<(u64, u64)> = self
            .by_size
            .range(count..)
            .next()
            .and_then(|(&size, starts)| starts.iter().next().map(|&start| (start, size)));

        if let Some((start, size)) = found {
            self.remove_extent(start);

            // If extent is larger, add remainder back
            if size > count {
                // Directly insert the remainder without going through add_extent
                let remainder = Extent::new(start + count, size - count);
                self.by_start.insert(remainder.start, remainder);
                self.by_size
                    .entry(remainder.count)
                    .or_default()
                    .insert(remainder.start);
                self.free_blocks += remainder.count;
            }

            return Some(start);
        }
        None
    }

    /// Free a range of blocks.
    pub fn free_range(&mut self, start: u64, count: u64) {
        self.add_extent(Extent::new(start, count));
    }

    /// Get total free blocks.
    pub fn free_blocks(&self) -> u64 {
        self.free_blocks
    }

    /// Get number of extents.
    pub fn extent_count(&self) -> usize {
        self.by_start.len()
    }

    /// Get largest free extent size.
    pub fn largest_extent(&self) -> u64 {
        self.by_size.keys().next_back().copied().unwrap_or(0)
    }
}

impl Default for ExtentTree {
    fn default() -> Self {
        Self::new()
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// BLOCK ALLOCATOR
// ═══════════════════════════════════════════════════════════════════════════════

/// Block allocator for a thin pool.
#[derive(Debug)]
pub struct BlockAllocator {
    /// Extent tree for free space.
    extents: ExtentTree,
    /// Total blocks in pool.
    total_blocks: u64,
    /// Block size in bytes.
    block_size: u64,
    /// Allocation policy.
    policy: AllocationPolicy,
    /// Next allocation hint (for NextFit policy).
    next_hint: u64,
    /// Reserved blocks (for metadata, etc.).
    reserved: u64,
}

impl BlockAllocator {
    /// Create a new block allocator.
    pub fn new(total_blocks: u64, block_size: u64, reserved: u64) -> Self {
        let usable = total_blocks.saturating_sub(reserved);
        let extents = ExtentTree::with_extent(reserved, usable);

        Self {
            extents,
            total_blocks,
            block_size,
            policy: AllocationPolicy::default(),
            next_hint: reserved,
            reserved,
        }
    }

    /// Set allocation policy.
    pub fn set_policy(&mut self, policy: AllocationPolicy) {
        self.policy = policy;
    }

    /// Allocate a single block.
    pub fn allocate(&mut self) -> ThinResult<PhysicalBlock> {
        self.allocate_contiguous(1)
    }

    /// Allocate contiguous blocks.
    pub fn allocate_contiguous(&mut self, count: u64) -> ThinResult<PhysicalBlock> {
        if count == 0 {
            return Err(ThinError::InvalidConfig("count must be > 0".into()));
        }

        let start = match self.policy {
            AllocationPolicy::FirstFit => self.extents.allocate_first_fit(count),
            AllocationPolicy::BestFit => self.extents.allocate_best_fit(count),
            AllocationPolicy::NextFit => {
                // Try from hint, then from start
                self.allocate_next_fit(count)
            }
            AllocationPolicy::Contiguous | AllocationPolicy::Striped => {
                self.extents.allocate_first_fit(count)
            }
        };

        match start {
            Some(block) => {
                self.next_hint = block + count;
                Ok(PhysicalBlock::new(block))
            }
            None => Err(ThinError::AllocationFailed),
        }
    }

    /// Next-fit allocation.
    fn allocate_next_fit(&mut self, count: u64) -> Option<u64> {
        // This is a simplified next-fit that falls back to first-fit
        // A full implementation would track position across allocations
        self.extents.allocate_first_fit(count)
    }

    /// Free a block.
    pub fn free(&mut self, block: PhysicalBlock) {
        self.free_contiguous(block, 1);
    }

    /// Free contiguous blocks.
    pub fn free_contiguous(&mut self, block: PhysicalBlock, count: u64) {
        if block.0 >= self.reserved && block.0 + count <= self.total_blocks {
            self.extents.free_range(block.0, count);
        }
    }

    /// Get number of free blocks.
    pub fn free_blocks(&self) -> u64 {
        self.extents.free_blocks()
    }

    /// Get total usable blocks.
    pub fn usable_blocks(&self) -> u64 {
        self.total_blocks.saturating_sub(self.reserved)
    }

    /// Get number of allocated blocks.
    pub fn allocated_blocks(&self) -> u64 {
        self.usable_blocks() - self.free_blocks()
    }

    /// Get usage percentage.
    pub fn usage_percent(&self) -> u8 {
        let usable = self.usable_blocks();
        if usable == 0 {
            return 0;
        }
        let allocated = self.allocated_blocks();
        ((allocated as u128 * 100) / usable as u128) as u8
    }

    /// Get free bytes.
    pub fn free_bytes(&self) -> u64 {
        self.free_blocks() * self.block_size
    }

    /// Get largest contiguous free extent.
    pub fn largest_contiguous(&self) -> u64 {
        self.extents.largest_extent()
    }

    /// Get fragmentation percentage (0 = not fragmented, 100 = fully fragmented).
    pub fn fragmentation(&self) -> u8 {
        let extent_count = self.extents.extent_count() as u64;
        let free = self.free_blocks();

        if free == 0 || extent_count == 0 {
            return 0;
        }

        // Ideal: 1 extent for all free space
        // Fragmentation = (extents - 1) / (free_blocks - 1) * 100
        // Simplified: more extents = more fragmentation
        let frag = ((extent_count.saturating_sub(1)) as u128 * 100) / free as u128;
        (frag as u8).min(100)
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_free_bitmap_new() {
        let bitmap = FreeBitmap::new(100);
        assert_eq!(bitmap.total_blocks(), 100);
        assert_eq!(bitmap.free_count(), 100);
        assert_eq!(bitmap.usage_percent(), 0);
    }

    #[test]
    fn test_free_bitmap_allocate() {
        let mut bitmap = FreeBitmap::new(100);

        assert!(bitmap.is_free(0));
        assert!(bitmap.allocate(0));
        assert!(!bitmap.is_free(0));
        assert_eq!(bitmap.free_count(), 99);

        // Double allocate should fail
        assert!(!bitmap.allocate(0));
    }

    #[test]
    fn test_free_bitmap_free() {
        let mut bitmap = FreeBitmap::new(100);

        bitmap.allocate(50);
        assert!(!bitmap.is_free(50));

        bitmap.free(50);
        assert!(bitmap.is_free(50));
    }

    #[test]
    fn test_free_bitmap_find_first_free() {
        let mut bitmap = FreeBitmap::new(100);

        assert_eq!(bitmap.find_first_free(), Some(0));

        bitmap.allocate(0);
        assert_eq!(bitmap.find_first_free(), Some(1));

        for i in 0..50 {
            bitmap.allocate(i);
        }
        assert_eq!(bitmap.find_first_free(), Some(50));
    }

    #[test]
    fn test_free_bitmap_contiguous() {
        let mut bitmap = FreeBitmap::new(100);

        // All free, should find 10 contiguous at start
        assert_eq!(bitmap.find_contiguous(10), Some(0));

        // Allocate blocks 5-15, leaving holes
        for i in 5..15 {
            bitmap.allocate(i);
        }

        // Can't find 10 contiguous at start anymore
        assert_eq!(bitmap.find_contiguous(5), Some(0));
        assert_eq!(bitmap.find_contiguous(10), Some(15));
    }

    #[test]
    fn test_extent() {
        let e1 = Extent::new(0, 10);
        let e2 = Extent::new(10, 5);
        let e3 = Extent::new(20, 5);

        assert!(e1.contains(5));
        assert!(!e1.contains(10));

        assert!(e1.adjacent(&e2));
        assert!(!e1.adjacent(&e3));

        let merged = e1.merge(&e2).unwrap();
        assert_eq!(merged.start, 0);
        assert_eq!(merged.count, 15);
    }

    #[test]
    fn test_extent_tree_add() {
        let mut tree = ExtentTree::new();

        tree.add_extent(Extent::new(0, 10));
        assert_eq!(tree.free_blocks(), 10);
        assert_eq!(tree.extent_count(), 1);

        // Add non-adjacent extent
        tree.add_extent(Extent::new(20, 5));
        assert_eq!(tree.free_blocks(), 15);
        assert_eq!(tree.extent_count(), 2);

        // Add extent that merges with first
        tree.add_extent(Extent::new(10, 10)); // Connects 0-10 with 20-25
        assert_eq!(tree.free_blocks(), 25);
        assert_eq!(tree.extent_count(), 1); // All merged
    }

    #[test]
    fn test_extent_tree_allocate() {
        let mut tree = ExtentTree::with_extent(0, 100);

        let block = tree.allocate_first_fit(10).unwrap();
        assert_eq!(block, 0);
        assert_eq!(tree.free_blocks(), 90);

        let block2 = tree.allocate_best_fit(5).unwrap();
        assert_eq!(block2, 10);
        assert_eq!(tree.free_blocks(), 85);
    }

    #[test]
    fn test_block_allocator() {
        let mut alloc = BlockAllocator::new(1000, 4096, 10);

        // 10 blocks reserved, 990 usable
        assert_eq!(alloc.usable_blocks(), 990);
        assert_eq!(alloc.free_blocks(), 990);

        // Allocate a block
        let block = alloc.allocate().unwrap();
        assert!(block.0 >= 10); // Should be after reserved
        assert_eq!(alloc.free_blocks(), 989);
    }

    #[test]
    fn test_block_allocator_contiguous() {
        let mut alloc = BlockAllocator::new(1000, 4096, 0);

        let start = alloc.allocate_contiguous(100).unwrap();
        assert_eq!(alloc.free_blocks(), 900);

        // Largest extent should be the remainder
        assert_eq!(alloc.largest_contiguous(), 900);

        // Free the blocks
        alloc.free_contiguous(start, 100);
        assert_eq!(alloc.free_blocks(), 1000);
        assert_eq!(alloc.largest_contiguous(), 1000);
    }

    #[test]
    fn test_block_allocator_full() {
        let mut alloc = BlockAllocator::new(10, 4096, 0);

        // Allocate all blocks
        for _ in 0..10 {
            alloc.allocate().unwrap();
        }

        // Should fail now
        assert!(alloc.allocate().is_err());
        assert_eq!(alloc.usage_percent(), 100);
    }

    #[test]
    fn test_block_allocator_fragmentation() {
        let mut alloc = BlockAllocator::new(100, 4096, 0);

        // Allocate every other block (creates fragmentation)
        for i in (0..100).step_by(2) {
            let block = alloc.allocate().unwrap();
            // Skip a block by freeing it
        }

        // Some fragmentation expected
        let frag = alloc.fragmentation();
        // Can't be too specific, but should be non-zero after alternating allocations
    }
}