// TODO: dedupe the code for iterating the directory records
use core::{
    any::Any,
    mem::MaybeUninit,
    ops::Deref,
    sync::atomic::{AtomicBool, Ordering},
};

use alloc::sync::Arc;
use libk::{
    block,
    error::Error,
    vfs::{
        CommonImpl, CreateFileType, CreateInfo, DirectoryImpl, DirectoryOpenPosition, Filename,
        Metadata, Node, NodeFlags, NodeRef,
    },
};
use walk::{DirentIter, DirentIterMut};
use yggdrasil_abi::io::{DirectoryEntry, FileType};

use crate::{file::RegularNode, inode::InodeAccess, symlink::SymlinkNode, Dirent, Ext2Fs, Inode};

mod walk;

pub struct DirectoryNode {
    fs: Arc<Ext2Fs>,
    locked: AtomicBool,
    pub(crate) inode: InodeAccess,
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug, Hash)]
pub struct DirentName<'a>(&'a str);

impl DirentName<'_> {
    pub const DOT: DirentName<'static> = DirentName(".");
    pub const DOT_DOT: DirentName<'static> = DirentName("..");

    pub fn is_dot_or_dotdot(&self) -> bool {
        self.0 == Self::DOT.0 || self.0 == Self::DOT_DOT.0
    }
}

impl<'a> TryFrom<&'a Filename> for DirentName<'a> {
    type Error = Error;

    fn try_from(value: &'a Filename) -> Result<Self, Self::Error> {
        if value.len() >= 255 {
            return Err(Error::InvalidArgument);
        }
        Ok(Self(value))
    }
}

impl Deref for DirentName<'_> {
    type Target = str;

    fn deref(&self) -> &Self::Target {
        self.0
    }
}

impl DirectoryNode {
    pub fn new(fs: Arc<Ext2Fs>, inode: InodeAccess) -> NodeRef {
        Node::directory(
            Self {
                fs: fs.clone(),
                locked: AtomicBool::new(false),
                inode,
            },
            NodeFlags::empty(),
            None,
            Some(fs),
        )
    }

    pub async fn create(
        fs: Arc<Ext2Fs>,
        inode: InodeAccess,
        parent_ino: Option<u32>,
    ) -> Result<NodeRef, Error> {
        let ino = inode.ino();
        let this = Self {
            fs: fs.clone(),
            locked: AtomicBool::new(false),
            inode,
        };

        // fsck wants . as first entry, .. as second
        this.create_entry(DirentName::DOT, ino).await?;
        if let Some(parent) = parent_ino {
            this.create_entry(DirentName::DOT_DOT, parent).await?;
        }

        Ok(Node::directory(this, NodeFlags::empty(), None, Some(fs)))
    }

    async fn create_entry(&self, name: DirentName<'_>, ino: u32) -> Result<(), Error> {
        async fn try_fit(
            inode: &mut Inode,
            fs: &Arc<Ext2Fs>,
            name: DirentName<'_>,
            ino: u32,
            ty: u8,
        ) -> Result<bool, Error> {
            let n = inode.blocks(fs) as u32;

            for i in 0..n {
                let fit_block = fs
                    .with_inode_block_mut(inode, i, 0, |block| {
                        let mut iter = DirentIterMut::new(fs, &mut block[..], 0);
                        if iter.try_fit(name, ino, ty) {
                            Ok(true)
                        } else {
                            Ok(false)
                        }
                    })
                    .await?;

                if fit_block {
                    return Ok(true);
                }
            }

            Ok(false)
        }

        async fn extend(
            inode: &mut Inode,
            fs: &Arc<Ext2Fs>,
            name: DirentName<'_>,
            ino: u32,
            ty: u8,
        ) -> Result<(), Error> {
            // Allocate a new block
            let block_index = inode.blocks(fs) as u32;

            // Grow the storage
            inode
                .reserve(fs, (block_index as u64 + 1) * fs.block_size as u64)
                .await?;

            fs.with_inode_block_mut(inode, block_index, fs.block_size, |block| {
                // Place a dummy dirent spanning the whole block
                block.fill(0);
                block[..size_of::<Dirent>()].copy_from_slice(bytemuck::bytes_of(&Dirent {
                    ino: 0,
                    ent_size: fs.block_size as _,
                    type_indicator: 0,
                    name_length_low: 0,
                }));

                // Insert the entry there
                let mut iter = DirentIterMut::new(fs, &mut block[..], 0);
                if iter.try_fit(name, ino, ty) {
                    Ok(())
                } else {
                    unreachable!(
                        "Possible driver bug: cannot fit a record into a just-allocated block"
                    )
                }
            })
            .await
        }

        // This directory is in the process of being unlinked and cannot be written to
        if self.locked.load(Ordering::Acquire) {
            return Err(Error::DoesNotExist);
        }

        let child = InodeAccess::new(self.inode.cache().clone(), ino);

        child.link().await?;
        let ty = child
            .map(|inode| inode.mode.dirent_indicator().ok_or(Error::InvalidArgument))
            .await?;

        // Insert the entry
        self.inode
            .amap_mut(async |inode| {
                // Try fitting the entry in existing blocks
                if try_fit(inode, &self.fs, name, ino, ty).await? {
                    return Ok(());
                }
                // Otherwise, allocate a new block and extend the directory
                extend(inode, &self.fs, name, ino, ty).await
            })
            .await?;

        Ok(())
    }

    async fn prepare_for_removal(&self) -> Result<(), Error> {
        // Check that the directory doesn't have any children besides . and .. and set the locked
        // flag to prevent anyone from adding entries while the directory is being unlinked
        // TODO find a better way to do this

        self.locked.store(true, Ordering::Release);

        let result = self
            .inode
            .map_blocks(|_, _, block| {
                let mut iter = DirentIter::new(&self.fs, block, 0);
                while let Some(record) = iter.next_record() {
                    if !record.is_empty()
                        && record.name_bytes() != b"."
                        && record.name_bytes() != b".."
                    {
                        return Err(Error::IsADirectory);
                    }
                }

                Ok(Option::<()>::None)
            })
            .await;

        match result {
            Ok(_) => {
                // Remove "." link
                self.inode
                    .map_mut(|inode| {
                        inode.hard_links -= 1;
                        Ok(())
                    })
                    .await?;
                Ok(())
            }
            Err(error) => {
                self.locked.store(false, Ordering::Release);
                Err(error)
            }
        }
    }

    async fn remove_entry(&self, child: &Node, name: DirentName<'_>) -> Result<(), Error> {
        assert_eq!(
            Arc::as_ptr(&child.filesystem().unwrap()).addr(),
            Arc::as_ptr(&self.fs).addr()
        );
        let (child_ino, is_directory) = match child.data_as_any() {
            data if let Some(dir) = data.downcast_ref::<DirectoryNode>() => {
                // Check that the directory is not empty
                dir.prepare_for_removal().await?;
                (dir.inode.ino(), true)
            }
            data if let Some(file) = data.downcast_ref::<RegularNode>() => {
                (file.inode.ino(), false)
            }
            data if let Some(link) = data.downcast_ref::<SymlinkNode>() => {
                (link.inode.ino(), false)
            }
            _ => return Err(Error::InvalidOperation),
        };

        if is_directory {
            // child's .. points to this node
            self.inode
                .map_mut(|inode| {
                    inode.hard_links = inode.hard_links.saturating_sub(1);
                    Ok(())
                })
                .await?;
        }

        let ino = self
            .inode
            .map_blocks_mut(|_inode, _index, block| {
                let mut iter = DirentIterMut::new(&self.fs, block, 0);
                Ok(iter.remove(name))
            })
            .await?;

        let ino = ino.ok_or(Error::DoesNotExist)?;

        assert_eq!(ino, child_ino);
        log::info!("ext2: unlinked ino #{ino}");

        let child = InodeAccess::new(self.inode.cache().clone(), ino);
        // TODO put the inode into the "free list" to be freed on cache flush/when no one is
        //      holding a ref to it any longer
        child
            .map_mut(|inode| {
                inode.hard_links = inode.hard_links.saturating_sub(1);
                log::info!("#{ino} hard links = {}", inode.hard_links);
                Ok(())
            })
            .await?;

        Ok(())
    }

    async fn lookup_entry(&self, name: DirentName<'_>) -> Result<NodeRef, Error> {
        let entry = self
            .inode
            .map_blocks(|_inode, _index, block| {
                let mut iter = DirentIter::new(&self.fs, block, 0);
                Ok(iter.find_ino(name))
            })
            .await?;

        let ino = entry.ok_or(Error::DoesNotExist)?;
        // Shouldn't happen, but just in case
        if ino == 0 {
            return Err(Error::DoesNotExist);
        }

        self.fs.load_node(ino).await
    }

    async fn read_entries(
        &self,
        mut pos: u64,
        entries: &mut [MaybeUninit<DirectoryEntry>],
    ) -> Result<(usize, u64), Error> {
        self.inode
            .amap(async |inode| {
                let size = inode.size(&self.fs);
                if pos >= inode.size(&self.fs) || entries.is_empty() {
                    return Ok((0, pos));
                }

                loop {
                    if pos >= size {
                        return Ok((0, pos));
                    }

                    let index = pos / self.fs.block_size as u64;
                    let offset = (pos % self.fs.block_size as u64) as usize;

                    let (entry_count, new_pos) = self
                        .fs
                        .with_inode_block(inode, index as u32, |block| {
                            let mut iter = DirentIter::new(&self.fs, block, offset);
                            let entry_count = iter.list_entries(entries);
                            Ok((entry_count, pos + iter.offset() as u64))
                        })
                        .await?;

                    pos = new_pos;

                    // If read any entries from the block, return
                    if entry_count != 0 {
                        return Ok((entry_count, pos));
                    }

                    // Otherwise go to next block
                    pos = (index + 1) * self.fs.block_size as u64;
                }
            })
            .await
    }
}

impl CommonImpl for DirectoryNode {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn size(&self, _node: &NodeRef) -> Result<u64, Error> {
        block!(self.inode.size().await)?
    }

    fn metadata(&self, _node: &NodeRef) -> Result<Metadata, Error> {
        block!(self.inode.metadata().await)?
    }

    fn set_metadata(&self, _node: &NodeRef, metadata: &Metadata) -> Result<(), Error> {
        block!(self.inode.update_metadata(metadata).await)?
    }
}

impl DirectoryImpl for DirectoryNode {
    fn open(&self, _node: &NodeRef) -> Result<DirectoryOpenPosition, Error> {
        Ok(DirectoryOpenPosition::FromPhysical(0))
    }

    fn len(&self, _node: &NodeRef) -> Result<usize, Error> {
        Err(Error::NotImplemented)
    }

    fn lookup(&self, _node: &NodeRef, name: &Filename) -> Result<NodeRef, Error> {
        let name = DirentName::try_from(name)?;
        block!(self.lookup_entry(name).await)?
    }

    fn create_node(&self, _parent: &NodeRef, info: &CreateInfo) -> Result<NodeRef, Error> {
        if self.fs.force_readonly {
            return Err(Error::ReadOnly);
        }

        let (ty, symlink) = match info.ty {
            CreateFileType::File => (FileType::File, None),
            CreateFileType::Directory => (FileType::Directory, None),
            CreateFileType::Symlink(symlink) => (FileType::Symlink, Some(symlink)),
        };

        block!(
            InodeAccess::allocate(&self.fs, ty, info.mode, Some(self.inode.ino()), symlink).await
        )?
    }

    fn attach_node(
        &self,
        _parent: &NodeRef,
        child: &NodeRef,
        name: &Filename,
    ) -> Result<(), Error> {
        if self.fs.force_readonly {
            return Err(Error::ReadOnly);
        }
        let name = DirentName::try_from(name)?;
        // Check that child is ext2
        let child_ino = match child.data_as_any() {
            data if let Some(dir) = data.downcast_ref::<DirectoryNode>() => dir.inode.ino(),
            data if let Some(file) = data.downcast_ref::<RegularNode>() => file.inode.ino(),
            data if let Some(link) = data.downcast_ref::<SymlinkNode>() => link.inode.ino(),
            _ => return Err(Error::InvalidOperation),
        };
        block!(self.create_entry(name, child_ino).await)??;

        Ok(())
    }

    fn unlink_node(
        &self,
        _parent: &NodeRef,
        child: &NodeRef,
        name: &Filename,
    ) -> Result<(), Error> {
        if self.fs.force_readonly {
            return Err(Error::ReadOnly);
        }
        let name = DirentName::try_from(name)?;
        block!(self.remove_entry(child, name).await)?
    }

    fn read_entries(
        &self,
        _node: &NodeRef,
        pos: u64,
        entries: &mut [MaybeUninit<DirectoryEntry>],
    ) -> Result<(usize, u64), Error> {
        block!(self.read_entries(pos, entries).await)?
    }
}
