use crate::memory::cached_block::convert;
use crate::disk::data_block::{DataBlock, DoubleIndirectBlock, IndirectBlock, TripleIndirectBlock};
use crate::disk::inode::Inode;
use crate::AyaFS;

impl AyaFS {
    /// 为 Inode 分配新 block, 返回 block 的编号
    pub(crate) fn allocate_block(&mut self, inode: &mut Inode) -> Option<u32> {
        // 先看这个 inode 的 direct block 有没有空闲
        for index in inode.direct.iter_mut() {
            if !self.data_bitmap.query(*index as usize) {
                let block_index = self.data_bitmap.allocate().unwrap() as u32;
                *index = block_index;
                inode.n_blocks += 1;
                // 当调用 get_inode_mut 拿出 &mut Inode 的时候对应的 block 在 cache 里已经脏了
                return Some(block_index);
            }
        }

        // direct block 全部分配完了, 先检查 indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.single_indirect as usize) {
            inode.single_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 indirect block 里尝试分配
        if let Some(block_index) = self.allocate_in_indirect(inode.single_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct & indirect block 全部分配完了, 先检查 double indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.double_indirect as usize) {
            inode.double_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_double_indirect(inode.double_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }

        // direct, indirect & double indirect block 全部分配完了, 先检查 triple indirect block 有没有分配, 没有就分配一个
        if !self.data_bitmap.query(inode.triple_indirect as usize) {
            inode.triple_indirect = self
                .data_bitmap
                .allocate()
                .expect("No free space for new block") as u32;
        }
        // 在 double indirect block 里尝试分配
        if let Some(block_index) = self.alloc_in_triple_indirect(inode.triple_indirect) {
            inode.n_blocks += 1;
            return Some(block_index);
        }
        None
    }

    fn allocate_in_indirect(&mut self, indirect_entry: u32) -> Option<u32> {
        // 取出 single indirect block, 尝试在里面分配
        let indirect_entry = indirect_entry as usize;

        if let Some(block) = self
            .get_block(indirect_entry)
            .map(convert::<DataBlock, IndirectBlock>)
        {
            let mut indirect_block = block.clone();
            for entry in indirect_block.block.entries.iter_mut() {
                if self.data_bitmap.query(*entry as usize) == false {
                    indirect_block.dirty = true; // 把这个块标记为 dirty
                    let block_index = self.data_bitmap.allocate().expect("No free space") as u32;
                    *entry = block_index;
                    self.update_block(indirect_block);
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_double_indirect(&mut self, double_indirect_entry: u32) -> Option<u32> {
        let double_indirect_entry = double_indirect_entry as usize;

        if let Some(block) = self
            .get_block(double_indirect_entry)
            .map(convert::<DataBlock, DoubleIndirectBlock>)
        {
            let mut double_indirect_block = block.clone();
            for indirect_entry in double_indirect_block.block.indirect.iter_mut() {
                if self.data_bitmap.query(*indirect_entry as usize) == false {
                    double_indirect_block.dirty = true;
                    *indirect_entry = self.data_bitmap.allocate().expect("No free space") as u32;
                }

                if let Some(block_index) = self.allocate_in_indirect(*indirect_entry) {
                    if double_indirect_block.dirty {
                        self.update_block(double_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }

    fn alloc_in_triple_indirect(&mut self, triple_indirect_entry: u32) -> Option<u32> {
        let triple_indirect_entry = triple_indirect_entry as usize;

        if let Some(block) = self
            .get_block(triple_indirect_entry)
            .map(convert::<DataBlock, TripleIndirectBlock>)
        {
            let mut triple_indirect_block = block.clone();
            for double_indirect_entry in triple_indirect_block.block.double_indirect.iter_mut() {
                if self.data_bitmap.query(*double_indirect_entry as usize) == false {
                    triple_indirect_block.dirty = true;
                    *double_indirect_entry =
                        self.data_bitmap.allocate().expect("No free space") as u32;
                }
                if let Some(block_index) = self.alloc_in_double_indirect(*double_indirect_entry) {
                    if triple_indirect_block.dirty {
                        self.update_block(triple_indirect_block);
                    }
                    return Some(block_index);
                }
            }
        }
        None
    }
}