#![feature(if_let_guard, impl_trait_in_assoc_type)]
#![cfg_attr(not(test), no_std)]
#![allow(clippy::new_ret_no_self)]

extern crate alloc;

use core::mem;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use bytemuck::Zeroable;
use data::{FsReadonlyFeatures, FsRequiredFeatures, DIRECT_BLOCK_COUNT};
use dir::DirectoryNode;
use file::RegularNode;
use inode::{InodeAccess, InodeCache};
use libk::{
    device::block::{cache::DeviceMapper, BlockDevice},
    error::Error,
    vfs::{Filesystem, FilesystemMountOption, NodeRef},
};
use libk_util::OneTimeInit;

pub mod access;
mod allocation;
mod data;
pub mod inode;

pub mod dir;
pub mod file;
pub mod symlink;

pub use data::{BlockGroupDescriptor, Dirent, ExtendedSuperblock, Inode};
use symlink::SymlinkNode;
use yggdrasil_abi::io::FileType;

struct Bgdt {
    base: u32,
    entry_count: usize,
    block_count: usize,
}

pub struct Ext2Fs {
    mapper: DeviceMapper,
    inode_cache: OneTimeInit<Arc<InodeCache>>,

    sb_position: u64,
    bgdt: Bgdt,

    total_inodes: u32,
    total_blocks: u32,
    block_group_inode_count: u32,
    block_group_block_count: u32,

    inode_size: usize,
    block_size: usize,
    inodes_per_block: usize,
    pointers_per_block: usize,

    max_file_size: u64,

    required_features: FsRequiredFeatures,
    write_features: FsReadonlyFeatures,

    force_readonly: bool,

    root: OneTimeInit<NodeRef>,
}

#[async_trait]
impl Filesystem for Ext2Fs {
    async fn flush(&self) -> Result<(), Error> {
        log::info!("ext2: flushing caches");
        let mut last_err = None;

        if let Err(error) = self.flush_inode_cache().await {
            log::error!("ext2: inode cache flush error {error:?}");
            last_err = Some(error);
        }
        if let Err(error) = self.mapper.flush().await {
            log::error!("ext2: block cache flush error {error:?}");
            last_err = Some(error);
        }

        match last_err {
            None => Ok(()),
            Some(error) => Err(error),
        }
    }

    fn display_name(&self) -> &'static str {
        "ext2"
    }
}

impl Ext2Fs {
    pub async fn create<'a, I: IntoIterator<Item = FilesystemMountOption<'a>>>(
        device: Arc<dyn BlockDevice>,
        options: I,
    ) -> Result<NodeRef, Error> {
        let mut cached = true;
        for option in options {
            match option {
                FilesystemMountOption::Sync => cached = false,
                FilesystemMountOption::FsSpecific(opt) => {
                    log::warn!("ext2: ignoring unrecognized option {opt:?}");
                }
            }
        }

        let fs = Self::create_fs(device, cached).await.inspect_err(|e| {
            log::error!("Ext2 init error: {:?}", e);
        })?;
        let fs = Arc::new(fs);
        let inode_cache = InodeCache::with_capacity(fs.clone(), 64, !cached);

        fs.inode_cache.init(inode_cache.into());

        let root = fs.load_node(data::ROOT_INODE).await?;

        fs.root.init(root.clone());

        Ok(root)
    }

    fn handle_readonly_features(features: FsReadonlyFeatures) -> (bool, FsReadonlyFeatures) {
        // TODO this implementation ignores backups
        let supported =
            FsReadonlyFeatures::FILE_SIZE_64_BIT | FsReadonlyFeatures::SPARSE_SUPERBLOCKS;
        let mut enabled = FsReadonlyFeatures::empty();
        let mut readonly = false;

        for feature in features.set_fields() {
            let name = match feature {
                FsReadonlyFeatures::FILE_SIZE_64_BIT => "file-size-64-bit",
                FsReadonlyFeatures::DIR_BTREE => "dir-btree",
                FsReadonlyFeatures::SPARSE_SUPERBLOCKS => "sparse-superblock",
                _ => "unknown-feature",
            };

            if !supported.contains(feature) {
                log::warn!("ext2: unsupported write feature {name:?}");
                readonly = true;
            } else {
                log::debug!("ext2: enable write feature: {name:?}");
                enabled |= feature;
            }
        }

        (readonly, enabled)
    }

    fn handle_required_features(features: FsRequiredFeatures) -> (bool, FsRequiredFeatures) {
        let supported = FsRequiredFeatures::DIRENT_TYPE_FIELD;
        let mut enabled = FsRequiredFeatures::empty();
        let mut unsupported = false;

        for feature in features.set_fields() {
            let name = match feature {
                FsRequiredFeatures::DIRENT_TYPE_FIELD => "dirent-type-field",
                FsRequiredFeatures::COMPRESSION => "compression",
                _ => "unknown-feature",
            };

            if !supported.contains(feature) {
                log::warn!("ext2: unsupported required feature: {name:?}");
                unsupported = true;
            } else {
                log::debug!("ext2: enable required feature: {name:?}");
                enabled |= feature;
            }
        }

        (!unsupported, enabled)
    }

    async fn create_fs(device: Arc<dyn BlockDevice>, cached: bool) -> Result<Self, Error> {
        let mut superblock = ExtendedSuperblock::zeroed();
        device
            .read_exact(
                data::SUPERBLOCK_OFFSET,
                bytemuck::bytes_of_mut(&mut superblock),
            )
            .await?;

        if superblock.signature != data::EXT2_SIGNATURE {
            log::warn!(
                "Invalid ext2 signature: expected {:#x}, got {:#x}",
                data::EXT2_SIGNATURE,
                superblock.signature
            );
            return Err(Error::InvalidArgument);
        }

        let required = superblock.required_features();
        let readonly = superblock.readonly_features();

        let (supported, required_features) = Self::handle_required_features(required);
        if !supported {
            log::warn!("ext2: filesystem not supported");
            return Err(Error::NotImplemented);
        }

        let (force_readonly, write_features) = Self::handle_readonly_features(readonly);
        if force_readonly {
            log::warn!("ext2: cannot mount as read-write, mounting read-only");
        }

        let block_size = 1024usize << superblock.block_size_log2;

        if data::SUPERBLOCK_OFFSET as usize + size_of::<ExtendedSuperblock>() > block_size {
            log::error!("ext2: superblock struct crosses block boundary, cannot mount");
            return Err(Error::InvalidArgument);
        }

        let bgdt_block_index = (data::SUPERBLOCK_OFFSET as usize).div_ceil(block_size) as u32;
        let bgdt_entry_count = superblock
            .total_blocks
            .div_ceil(superblock.block_group_block_count) as usize;
        let bgdt_block_count =
            (bgdt_entry_count * size_of::<BlockGroupDescriptor>()).div_ceil(block_size);

        let bgdt = Bgdt {
            base: bgdt_block_index,
            block_count: bgdt_block_count,
            entry_count: bgdt_entry_count,
        };

        log::info!(
            "ext2 v{}.{}",
            superblock.version_major,
            superblock.version_minor
        );
        log::info!("Block groups: {}", bgdt_entry_count);
        log::info!(
            "BGDT size: {} blocks ({} bytes)",
            bgdt_block_count,
            bgdt_block_count * block_size
        );

        let total_inodes = superblock.total_inodes;
        let total_blocks = superblock.total_blocks;
        let block_group_inode_count = superblock.block_group_inode_count;
        let block_group_block_count = superblock.block_group_block_count;

        log::info!("Inode size: {}", superblock.inode_size());
        log::info!("Total free inodes: {}", superblock.total_unallocated_inodes);
        log::info!("Total free blocks: {}", superblock.total_unallocated_blocks);

        let mapper = match cached {
            false => DeviceMapper::uncached(device, block_size, "ext2")?,
            true => {
                let segment_size = (block_size * 16).min(65536);
                DeviceMapper::cached_with_capacity(
                    device,
                    block_size,
                    segment_size,
                    128,
                    64,
                    "ext2",
                )?
            }
        };

        let pointers_per_block = block_size / size_of::<u32>();
        let max_inode_blocks = (DIRECT_BLOCK_COUNT
            + pointers_per_block
            + pointers_per_block * pointers_per_block
            + pointers_per_block * pointers_per_block * pointers_per_block)
            as u64;
        let max_file_size = max_inode_blocks.saturating_mul(block_size as u64);
        let max_file_size = if write_features.contains(FsReadonlyFeatures::FILE_SIZE_64_BIT) {
            max_file_size
        } else {
            max_file_size.min(u32::MAX as u64)
        };

        Ok(Self {
            block_size,
            inode_size: superblock.inode_size(),
            inodes_per_block: block_size / superblock.inode_size(),
            pointers_per_block,

            total_inodes,
            total_blocks,
            block_group_inode_count,
            block_group_block_count,

            max_file_size,

            mapper,
            inode_cache: OneTimeInit::new(),

            sb_position: data::SUPERBLOCK_OFFSET,
            bgdt,

            required_features,
            write_features,
            force_readonly,

            root: OneTimeInit::new(),
        })
    }

    pub async fn load_node(self: &Arc<Self>, ino: u32) -> Result<NodeRef, Error> {
        let cache = self.inode_cache.get();
        let inode = InodeAccess::new(cache.clone(), ino);
        let mode = inode.map(|inode| Ok(inode.mode)).await?;
        match mode.node_type() {
            Some(FileType::Directory) => Ok(DirectoryNode::new(self.clone(), inode)),
            Some(FileType::File) => Ok(RegularNode::new(self.clone(), inode)),
            Some(FileType::Symlink) => Ok(SymlinkNode::new(self.clone(), inode)),
            e => {
                log::error!("Unhandled inode type: {e:?} ({mode:#x?})");
                Err(Error::InvalidArgument)
            }
        }
    }

    #[inline]
    fn block_address(&self, index: u32) -> u64 {
        index as u64 * self.block_size as u64
    }

    async fn inode(&self, ino: u32) -> Result<(u32, usize), Error> {
        if ino < 1 || ino >= self.total_inodes {
            return Err(Error::InvalidFile);
        }
        let ino = ino - 1;

        let ino_group = ino / self.block_group_inode_count;
        let ino_in_group = ino % self.block_group_inode_count;
        let offset_in_block = (ino_in_group as usize % self.inodes_per_block) * self.inode_size;

        let inode_table_block = self
            .with_bgdt_entry(ino_group, |entry| Ok(entry.inode_table))
            .await?;
        let inode_block = inode_table_block + ino_in_group / self.inodes_per_block as u32;

        assert!(offset_in_block < self.block_size);

        Ok((inode_block, offset_in_block))
    }

    pub async fn read_inode(&self, ino: u32) -> Result<Inode, Error> {
        let (ino_block, offset_in_block) = self.inode(ino).await?;
        self.with_block(ino_block, |block| {
            let inode =
                bytemuck::from_bytes(&block[offset_in_block..offset_in_block + size_of::<Inode>()]);
            Ok(*inode)
        })
        .await
    }

    pub async fn write_inode(&self, ino: u32, inode: &Inode) -> Result<(), Error> {
        let (ino_block, offset_in_block) = self.inode(ino).await?;
        self.with_block_mut(ino_block, size_of::<Inode>(), |block| {
            block[offset_in_block..offset_in_block + size_of::<Inode>()]
                .copy_from_slice(bytemuck::bytes_of(inode));
            Ok(())
        })
        .await?;
        Ok(())
    }

    pub async fn write_or_free_inode(&self, ino: u32, inode: &mut Inode) -> Result<(), Error> {
        if inode.is_deleted() {
            self.free_inode(ino, inode).await
        } else {
            self.write_inode(ino, inode).await
        }
    }

    pub async fn read_inode_data(
        &self,
        inode: &Inode,
        mut pos: u64,
        buffer: &mut [u8],
    ) -> Result<usize, Error> {
        let size = inode.size(self);

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

        let mut offset = 0;
        let mut remaining = buffer.len().min((size - pos) as usize);

        while remaining != 0 {
            let block_index = pos / self.block_size as u64;
            let block_offset = (pos % self.block_size as u64) as usize;
            let amount = remaining.min(self.block_size - block_offset);

            let block = self.inode_block_index(inode, block_index as u32).await?;
            if block == 0 {
                buffer[offset..offset + amount].fill(0);
            } else {
                self.with_block(block, |block| {
                    buffer[offset..offset + amount]
                        .copy_from_slice(&block[block_offset..block_offset + amount]);
                    Ok(())
                })
                .await?;
            }

            pos += amount as u64;
            offset += amount;
            remaining -= amount;
        }

        Ok(offset)
    }

    pub async fn flush_inode_cache(&self) -> Result<(), Error> {
        log::info!("Flushing inode cache");
        self.inode_cache.get().flush().await
    }

    async fn read_index(&self, block_index: u32, index: usize) -> Result<u32, Error> {
        self.with_block(block_index, |block| {
            let indirect: &[u32] = unsafe {
                core::slice::from_raw_parts(block.as_ptr() as *const u32, self.pointers_per_block)
            };

            Ok(indirect[index])
        })
        .await
    }

    async fn write_index(&self, block_index: u32, index: usize, value: u32) -> Result<(), Error> {
        self.with_block_mut(block_index, size_of::<u32>(), |block| {
            let indirect: &mut [u32] = unsafe {
                core::slice::from_raw_parts_mut(block.as_mut_ptr().cast(), self.pointers_per_block)
            };

            indirect[index] = value;
            Ok(())
        })
        .await
    }

    async fn remove_index(&self, block_index: u32, index: usize) -> Result<u32, Error> {
        self.with_block_mut(block_index, size_of::<u32>(), |block| {
            let indirect: &mut [u32] = unsafe {
                core::slice::from_raw_parts_mut(block.as_mut_ptr().cast(), self.pointers_per_block)
            };

            Ok(mem::replace(&mut indirect[index], 0))
        })
        .await
    }

    async fn inode_block_index(&self, inode: &Inode, index: u32) -> Result<u32, Error> {
        let mut index = index as usize;
        // L0
        if index < data::DIRECT_BLOCK_COUNT {
            return Ok(inode.blocks.direct_blocks[index]);
        }

        // L1
        index -= data::DIRECT_BLOCK_COUNT;
        if index < self.pointers_per_block {
            return self.read_index(inode.blocks.indirect_block_l1, index).await;
        }

        // L2
        index -= self.pointers_per_block;
        if index < self.pointers_per_block * self.pointers_per_block {
            let l1_index = self
                .read_index(
                    inode.blocks.indirect_block_l2,
                    index / self.pointers_per_block,
                )
                .await?;
            return self
                .read_index(l1_index, index % self.pointers_per_block)
                .await;
        }

        // L3
        index -= self.pointers_per_block * self.pointers_per_block;
        if index < self.pointers_per_block * self.pointers_per_block * self.pointers_per_block {
            let l2_pointer_index = index / self.pointers_per_block;
            let l2_index = self
                .read_index(
                    inode.blocks.indirect_block_l3,
                    l2_pointer_index / self.pointers_per_block,
                )
                .await?;
            let l1_index = self
                .read_index(l2_index, l2_pointer_index % self.pointers_per_block)
                .await?;
            return self
                .read_index(l1_index, index % self.pointers_per_block)
                .await;
        }

        log::error!("Invalid inode block index (over L3 indirect)");
        Err(Error::InvalidFile)
    }
}
