use libk::{error::Error, time::real_time};
use yggdrasil_abi::io::FileType;

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

impl Ext2Fs {
    async fn free<
        F: FnOnce(&mut BlockGroupDescriptor) -> u32,
        G: FnOnce(&mut ExtendedSuperblock),
    >(
        &self,
        element_index: u32,
        group_index: u32,
        descriptor_mapper: F,
        superblock_mapper: G,
    ) -> Result<(), Error> {
        let bitmap = self
            .with_bgdt_entry_mut(group_index, |descriptor| Ok(descriptor_mapper(descriptor)))
            .await?;
        self.with_block_mut(bitmap, size_of::<u32>(), |bitmap| {
            let index = (element_index / 8) as usize;
            let bit = 1u8 << (element_index % 8);

            if bitmap[index] & bit == 0 {
                log::warn!("Trying to free an unallocated element #{element_index}");
            }
            bitmap[index] &= !bit;
            Ok(())
        })
        .await?;
        self.with_superblock_mut(|sb| {
            superblock_mapper(sb);
            Ok(())
        })
        .await?;
        Ok(())
    }

    async fn allocate<
        F: Fn(&mut BlockGroupDescriptor) -> Option<(u32, u32)>,
        G: FnOnce(&mut ExtendedSuperblock),
    >(
        &self,
        descriptor_mapper: F,
        superblock_mapper: G,
    ) -> Result<u32, Error> {
        let bit_per_block = self.block_size * 8;
        for group_index in 0..self.bgdt.entry_count as u32 {
            let bitmap = self
                .with_bgdt_entry_mut(group_index, |descriptor| {
                    match descriptor_mapper(descriptor) {
                        Some(val) => Ok(Some(val)),
                        _ => Ok(None),
                    }
                })
                .await?;

            if let Some((bitmap, group_item_count)) = bitmap {
                let no = self
                    .with_block_mut(bitmap, 0, |bitmap| {
                        for i in 0..bit_per_block.min(group_item_count as usize) {
                            let index = i / 8;
                            let bit = 1u8 << (i % 8);

                            if bitmap[index] & bit == 0 {
                                bitmap[index] |= bit;
                                return Ok(Some(i as u32));
                            }
                        }
                        Ok(None)
                    })
                    .await?
                    .expect("TODO: bgdt says there're things, but bitmap says there aren't");

                self.with_superblock_mut(|sb| {
                    superblock_mapper(sb);
                    Ok(())
                })
                .await?;

                return Ok(group_index * group_item_count + no);
            }
        }

        todo!("ENOSPC")
    }

    pub(crate) async fn allocate_inode(&self, directory: bool) -> Result<u32, Error> {
        let ino = self
            .allocate(
                |descriptor| {
                    if descriptor.unallocated_inodes == 0
                        || (directory && descriptor.directories == u16::MAX)
                    {
                        None
                    } else {
                        if directory {
                            descriptor.directories += 1;
                        }
                        descriptor.unallocated_inodes -= 1;
                        Some((descriptor.inode_usage_bitmap, self.block_group_inode_count))
                    }
                },
                |superblock| {
                    superblock.total_unallocated_inodes -= 1;
                },
            )
            .await?;
        let ino = ino + 1;

        log::debug!("ext2: allocated inode #{ino}");

        Ok(ino)
    }

    pub(crate) async fn free_inode(&self, ino: u32, inode: &mut Inode) -> Result<(), Error> {
        if ino == 0 {
            log::warn!("Tried to free a NULL inode");
            return Ok(());
        }

        let is_directory = inode
            .mode
            .node_type()
            .map(|m| m == FileType::Directory)
            .unwrap_or(false);

        // Free inode blocks
        inode.resize(self, 0).await?;
        inode.dtime = real_time().seconds() as u32;

        self.write_inode(ino, inode).await?;

        let inode_index = ino - 1;
        let group_index = inode_index / self.block_group_inode_count;
        self.free(
            inode_index,
            group_index,
            |descriptor| {
                if is_directory {
                    descriptor.directories = descriptor.directories.saturating_sub(1);
                }
                descriptor.unallocated_inodes += 1;
                descriptor.inode_usage_bitmap
            },
            |superblock| {
                superblock.total_unallocated_inodes += 1;
            },
        )
        .await?;

        Ok(())
    }

    pub(crate) async fn allocate_block(&self) -> Result<u32, Error> {
        let no = self
            .allocate(
                |descriptor| {
                    if descriptor.unallocated_blocks == 0 {
                        None
                    } else {
                        descriptor.unallocated_blocks -= 1;
                        Some((descriptor.block_usage_bitmap, self.block_group_block_count))
                    }
                },
                |superblock| {
                    superblock.total_unallocated_blocks -= 1;
                },
            )
            .await?;

        log::debug!("ext2: allocated block #{no}");

        Ok(no)
    }

    pub(crate) async fn free_block(&self, block_index: u32) -> Result<(), Error> {
        if block_index == 0 {
            log::warn!("Tried to free a NULL block");
            return Ok(());
        }

        let group_index = block_index / self.block_group_block_count;
        self.free(
            block_index,
            group_index,
            |descriptor| {
                descriptor.unallocated_blocks += 1;
                descriptor.block_usage_bitmap
            },
            |superblock| {
                superblock.total_unallocated_blocks += 1;
            },
        )
        .await?;

        Ok(())
    }
}
