use crate::block_io::BlockIo;
use crate::cache::BlockCache;
use crate::error::FsError;
use lancer_core::fs::{BTREE_NODE_SIZE, BTreeNode};
use zerocopy::{FromBytes, IntoBytes};

const POOL_SIZE: usize = 256;
const BITMAP_WORDS: usize = POOL_SIZE / 64;

pub type NodeIdx = u16;

const POOL_SENTINEL: u64 = 1u64 << 63;

pub const fn is_pool_ref(physical_addr: u64) -> bool {
    physical_addr & POOL_SENTINEL != 0
}

pub const fn pool_idx_to_addr(idx: NodeIdx) -> u64 {
    POOL_SENTINEL | ((idx as u64) << 4)
}

pub const fn addr_to_pool_idx(addr: u64) -> NodeIdx {
    (((addr & lancer_core::fs::ADDR_MASK) & !POOL_SENTINEL) >> 4) as NodeIdx
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub enum NodeSlot {
    Free,
    Clean(u64),
    Dirty,
}

pub struct NodePool {
    nodes: [BTreeNode; POOL_SIZE],
    status: [NodeSlot; POOL_SIZE],
    free_bitmap: [u64; BITMAP_WORDS],
    free_count: u16,
}

impl NodePool {
    pub const fn new() -> Self {
        Self {
            nodes: [BTreeNode::ZEROED; POOL_SIZE],
            status: [NodeSlot::Free; POOL_SIZE],
            free_bitmap: [u64::MAX; BITMAP_WORDS],
            free_count: POOL_SIZE as u16,
        }
    }
}

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

impl NodePool {
    pub fn pool_alloc(&mut self) -> Result<NodeIdx, FsError> {
        let (word_idx, bit) = self
            .free_bitmap
            .iter()
            .enumerate()
            .find_map(|(i, &w)| match w {
                0 => None,
                _ => Some((i, w.trailing_zeros() as usize)),
            })
            .ok_or(FsError::PoolExhausted)?;

        let idx = (word_idx * 64 + bit) as NodeIdx;
        self.free_bitmap[word_idx] &= !(1u64 << bit);
        self.free_count -= 1;
        self.status[idx as usize] = NodeSlot::Dirty;
        self.nodes[idx as usize] = BTreeNode::EMPTY_LEAF;
        Ok(idx)
    }

    pub fn pool_free(&mut self, idx: NodeIdx) {
        let i = idx as usize;
        match self.status[i] {
            NodeSlot::Free => {}
            _ => {
                let word_idx = i / 64;
                let bit = i % 64;
                self.free_bitmap[word_idx] |= 1u64 << bit;
                self.free_count += 1;
                self.status[i] = NodeSlot::Free;
            }
        }
    }

    pub fn pool_load(
        &mut self,
        cache: &mut BlockCache,
        bio: &mut BlockIo,
        block_num: u64,
    ) -> Result<NodeIdx, FsError> {
        let idx = self.pool_alloc()?;
        let data = cache.cache_read(bio, block_num)?;
        let node = BTreeNode::read_from_bytes(data.as_slice()).map_err(|_| FsError::CorruptNode)?;
        self.nodes[idx as usize] = node;
        self.status[idx as usize] = NodeSlot::Clean(block_num);
        Ok(idx)
    }

    pub fn get(&self, idx: NodeIdx) -> &BTreeNode {
        &self.nodes[idx as usize]
    }

    pub fn get_mut(&mut self, idx: NodeIdx) -> &mut BTreeNode {
        self.status[idx as usize] = NodeSlot::Dirty;
        &mut self.nodes[idx as usize]
    }

    pub fn is_dirty(&self, idx: NodeIdx) -> bool {
        self.status[idx as usize] == NodeSlot::Dirty
    }

    pub fn status(&self, idx: NodeIdx) -> NodeSlot {
        self.status[idx as usize]
    }

    pub fn clear_all(&mut self) {
        self.free_bitmap = [u64::MAX; BITMAP_WORDS];
        self.free_count = POOL_SIZE as u16;
        self.status = [NodeSlot::Free; POOL_SIZE];
    }

    pub fn node_bytes(&self, idx: NodeIdx) -> &[u8; BTREE_NODE_SIZE] {
        self.nodes[idx as usize].as_bytes().try_into().unwrap()
    }

    pub fn dirty_indices_by_level(&self) -> ([NodeIdx; POOL_SIZE], usize) {
        let mut indices = [0u16; POOL_SIZE];
        let mut count = 0usize;

        (0..POOL_SIZE as NodeIdx).for_each(|idx| {
            if self.status[idx as usize] == NodeSlot::Dirty {
                indices[count] = idx;
                count += 1;
            }
        });

        let used = &mut indices[..count];
        used.sort_unstable_by_key(|&idx| self.nodes[idx as usize].header.level);

        (indices, count)
    }
}