#![no_std]

use alloc::sync::Arc;
use async_trait::async_trait;
use data::{Bpb, ClusterNumber, Fat32Ebpb, Fat32FsInfo};
use directory::DirectoryNode;
use libk::{
    device::block::{cache::DeviceMapper, BlockDevice},
    error::Error,
    vfs::{Filesystem, FilesystemMountOption, Metadata, Node, NodeFlags, NodeRef},
};
use libk_util::get_le_u32;
use yggdrasil_abi::{
    io::{FileMode, GroupId, UserId},
    time::SystemTime,
};

extern crate alloc;

mod data;

mod directory;
mod file;

pub struct FsLayout {
    // From BPB/EBPB
    bytes_per_sector: usize,
    total_sectors: u64,
    sectors_per_fat: usize,
    sectors_per_cluster: usize,
    root_directory_cluster: ClusterNumber,
    first_fat_sector: u64,

    // Computed
    first_data_sector: u64,
}

pub struct Fat32Fs {
    mapper: DeviceMapper,
    pub layout: FsLayout,
}

#[async_trait]
impl Filesystem for Fat32Fs {
    fn display_name(&self) -> &'static str {
        "fat32"
    }
}

impl Fat32Fs {
    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 {
            #[allow(clippy::single_match)]
            match option {
                FilesystemMountOption::Sync => cached = false,
                _ => (),
            }
        }

        let fs = Arc::new(Self::create_fs(device, cached).await?);

        // Setup root node
        let root_metadata = Metadata {
            uid: UserId::root(),
            gid: GroupId::root(),
            mode: FileMode::default_dir(),
            // TODO
            ctime: SystemTime::ZERO,
            mtime: SystemTime::ZERO,
            atime: SystemTime::ZERO,
            inode: Some(fs.layout.root_directory_cluster.0),
            block_size: fs.layout.bytes_per_sector as u64,
            block_count: 0,
        };
        let root_directory = DirectoryNode {
            fs: fs.clone(),
            cluster: fs.layout.root_directory_cluster,
            parent: None,
            size_bytes: None,
            metadata: root_metadata,
        };
        let root_node = Node::directory(root_directory, NodeFlags::empty(), None, Some(fs.clone()));

        Ok(root_node)
    }

    async fn create_fs(device: Arc<dyn BlockDevice>, cached: bool) -> Result<Fat32Fs, Error> {
        let mut boot_sector = [0; 512];
        let mut fsinfo = [0; 512];
        device.read_exact(0, &mut boot_sector).await?;

        let bpb = Bpb::from_bytes(&boot_sector);
        let ebpb = Fat32Ebpb::from_bytes(&boot_sector);
        let fsinfo_sector = ebpb.fsinfo_sector();
        let bytes_per_sector = bpb.bytes_per_sector();
        let total_sectors = bpb.total_sectors();

        if bpb.fat_size_16() != 0 {
            log::warn!("fat32: not a FAT32 filesystem");
            return Err(Error::InvalidArgument);
        }

        if fsinfo_sector >= total_sectors {
            log::warn!("fat32: FSInfo sector is beyond filesystem end");
            return Err(Error::InvalidArgument);
        }

        device
            .read_exact(bytes_per_sector as u64 * fsinfo_sector, &mut fsinfo)
            .await?;

        let fsinfo = Fat32FsInfo::from_bytes(fsinfo);
        if !fsinfo.is_valid_fat32() {
            log::warn!("fat32: not a FAT32 filesystem");
            return Err(Error::InvalidArgument);
        }

        let sectors_per_fat = ebpb.sectors_per_fat();
        let sectors_per_cluster = bpb.sectors_per_cluster();
        let root_directory_cluster = ebpb.root_directory_cluster();

        // + RootDirSectors, but RootDirSectors = 0 on FAT32
        let first_data_sector = bpb.reserved_sectors() + sectors_per_fat * bpb.fat_count();
        let first_fat_sector = bpb.reserved_sectors();

        log::info!("fat32: mounted {:?}", ebpb.volume_label());
        log::info!(
            "fat32: sector {}B, cluster {}B",
            bytes_per_sector,
            sectors_per_cluster * bytes_per_sector
        );

        let layout = FsLayout {
            bytes_per_sector,
            total_sectors,
            sectors_per_fat,
            sectors_per_cluster,
            root_directory_cluster,

            first_fat_sector: first_fat_sector as u64,
            first_data_sector: first_data_sector as u64,
        };

        let mapper = if cached {
            DeviceMapper::cached_with_capacity(
                device,
                // Most often 512
                layout.bytes_per_sector,
                // 512 * 16 * 4 = 32768
                (layout.bytes_per_sector * layout.sectors_per_cluster) * 4,
                128,
                64,
                "fat32",
            )?
        } else {
            DeviceMapper::uncached(device, layout.bytes_per_sector, "fat32")?
        };

        Ok(Self { layout, mapper })
    }

    async fn with_cluster_sector<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        cluster: ClusterNumber,
        sector_in_cluster: u64,
        mapper: F,
    ) -> Result<T, Error> {
        let sector = cluster.first_sector(&self.layout) + sector_in_cluster;
        if sector >= self.layout.total_sectors {
            log::warn!("fat32: sector {sector} beyond filesystem end");
            return Err(Error::InvalidArgument);
        }
        self.mapper
            .try_with(sector * self.layout.bytes_per_sector as u64, mapper)
            .await
    }

    async fn with_fat_sector<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        fat_sector: u64,
        mapper: F,
    ) -> Result<T, Error> {
        if fat_sector >= self.layout.sectors_per_fat as u64 {
            log::warn!("fat32: FAT sector {fat_sector} outside of FAT size");
            return Err(Error::InvalidArgument);
        }
        let sector = self.layout.first_fat_sector + fat_sector;
        self.mapper
            .try_with(sector * self.layout.bytes_per_sector as u64, mapper)
            .await
    }

    async fn iterate_clusters<F: FnMut(ClusterNumber)>(
        &self,
        start: ClusterNumber,
        count: usize,
        mut mapper: F,
    ) -> Result<(), Error> {
        let mut current = start;

        for _ in 0..count {
            let offset_in_fat = current.0 as usize * size_of::<u32>();
            let sector_in_fat = (offset_in_fat / self.layout.bytes_per_sector) as u64;
            let offset_in_sector = offset_in_fat % self.layout.bytes_per_sector;

            let cluster = self
                .with_fat_sector(sector_in_fat, |buffer| {
                    let number = get_le_u32(&buffer[offset_in_sector..]) & 0x0FFFFFFF;
                    Ok(number)
                })
                .await?;

            if cluster >= 0x0FFFFFF7 {
                return Err(Error::InvalidFile);
            }

            let cluster = ClusterNumber(cluster);
            mapper(cluster);
            current = cluster;
        }

        Ok(())
    }
}
