use libk::error::Error;

use crate::{BlockGroupDescriptor, Ext2Fs, ExtendedSuperblock, Inode};

pub enum InodeBlock<'a> {
    Hole,
    Data(&'a [u8]),
}

impl InodeBlock<'_> {
    pub fn copy_to(&self, offset: usize, destination: &mut [u8]) {
        match self {
            Self::Hole => destination.fill(0),
            Self::Data(data) => {
                destination.copy_from_slice(&data[offset..offset + destination.len()])
            }
        }
    }
}

impl Ext2Fs {
    pub async fn with_block<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        index: u32,
        mapper: F,
    ) -> Result<T, Error> {
        if index < 1 || index >= self.total_blocks {
            return Err(Error::InvalidFile);
        }
        self.mapper
            .try_with(self.block_address(index), mapper)
            .await
    }

    pub async fn with_block_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
        &self,
        index: u32,
        write_size: usize,
        mapper: F,
    ) -> Result<T, Error> {
        if index < 1 || index >= self.total_blocks {
            return Err(Error::InvalidFile);
        }
        self.mapper
            .try_with_mut(self.block_address(index), write_size, mapper)
            .await
    }

    pub async fn with_superblock_mut<T, F: FnOnce(&mut ExtendedSuperblock) -> Result<T, Error>>(
        &self,
        mapper: F,
    ) -> Result<T, Error> {
        let position = self.sb_position & !(self.block_size as u64 - 1);
        let offset = (self.sb_position & (self.block_size as u64 - 1)) as usize;

        self.mapper
            .try_with_mut(position, size_of::<ExtendedSuperblock>(), |block| {
                let sb = bytemuck::from_bytes_mut(
                    &mut block[offset..offset + size_of::<ExtendedSuperblock>()],
                );
                mapper(sb)
            })
            .await
    }

    pub async fn with_inode_block<T, F: FnOnce(InodeBlock) -> Result<T, Error>>(
        &self,
        inode: &Inode,
        block: u32,
        mapper: F,
    ) -> Result<T, Error> {
        let block_index = self.inode_block_index(inode, block).await?;
        if block_index == 0 {
            mapper(InodeBlock::Hole)
        } else {
            self.with_block(block_index, |block| mapper(InodeBlock::Data(block)))
                .await
        }
    }

    pub async fn with_inode_block_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
        &self,
        inode: &Inode,
        block: u32,
        write_size: usize,
        mapper: F,
    ) -> Result<T, Error> {
        let block_index = self.inode_block_index(inode, block).await?;
        self.with_block_mut(block_index, write_size, mapper).await
    }

    pub(crate) async fn with_bgdt_entry<T, F: FnOnce(&BlockGroupDescriptor) -> Result<T, Error>>(
        &self,
        block_group: u32,
        mapper: F,
    ) -> Result<T, Error> {
        let offset = block_group as usize * size_of::<BlockGroupDescriptor>();
        let block = offset / self.block_size;
        if block >= self.bgdt.block_count {
            log::warn!("ext2: bgdt block out of bounds: {block}");
            return Err(Error::InvalidArgument);
        }
        let offset_in_block = offset % self.block_size;

        self.with_block(block as u32 + self.bgdt.base, |block| {
            let descriptor = bytemuck::from_bytes(
                &block[offset_in_block..offset_in_block + size_of::<BlockGroupDescriptor>()],
            );
            mapper(descriptor)
        })
        .await
    }

    pub(crate) async fn with_bgdt_entry_mut<
        T,
        F: FnOnce(&mut BlockGroupDescriptor) -> Result<T, Error>,
    >(
        &self,
        block_group: u32,
        mapper: F,
    ) -> Result<T, Error> {
        let offset = block_group as usize * size_of::<BlockGroupDescriptor>();
        let block = offset / self.block_size;
        if block >= self.bgdt.block_count {
            log::warn!("ext2: bgdt block out of bounds: {block}");
            return Err(Error::InvalidArgument);
        }
        let offset_in_block = offset % self.block_size;

        self.with_block_mut(
            block as u32 + self.bgdt.base,
            size_of::<BlockGroupDescriptor>(),
            |block| {
                let descriptor = bytemuck::from_bytes_mut(
                    &mut block
                        [offset_in_block..offset_in_block + size_of::<BlockGroupDescriptor>()],
                );
                mapper(descriptor)
            },
        )
        .await
    }
}
