#![allow(clippy::missing_transmute_annotations)]
use core::ops::{Deref, DerefMut};

use alloc::sync::Arc;
use libk_util::{lru_hash_table::LruCache, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::error::Error;

use crate::{dma::DmaBuffer, task::sync::AsyncMutex};

use super::BlockDevice;

pub struct CachedSegment {
    data: DmaBuffer<[u8]>,
    dirty: bool,
}

pub struct UncachedCache {
    device: Arc<dyn BlockDevice>,
    block_size: usize,
}

pub enum DeviceMapper {
    Uncached(UncachedCache),
    Cached(BlockCache),
}

pub struct BlockCache {
    device: Arc<dyn BlockDevice>,
    block_size: usize,
    segment_size: usize,
    cache: AsyncMutex<LruCache<u64, Arc<IrqSafeRwLock<CachedSegment>>>>,
}

impl DeviceMapper {
    pub fn cached_with_capacity(
        device: Arc<dyn BlockDevice>,
        block_size: usize,
        segment_size: usize,
        bucket_capacity: usize,
        bucket_count: usize,
        filesystem: &str,
    ) -> Result<DeviceMapper, Error> {
        BlockCache::with_capacity(
            device,
            block_size,
            segment_size,
            bucket_capacity,
            bucket_count,
            filesystem,
        )
        .map(DeviceMapper::Cached)
    }

    pub fn uncached(
        device: Arc<dyn BlockDevice>,
        block_size: usize,
        filesystem: &str,
    ) -> Result<DeviceMapper, Error> {
        if block_size % device.block_size() != 0 {
            log::error!(
                "Couldn't create block mapper for {filesystem}: \
                 cache block size is not a multiple of device block size"
            );
            return Err(Error::InvalidArgument);
        }
        let uncache = UncachedCache { device, block_size };
        Ok(Self::Uncached(uncache))
    }

    pub fn device(&self) -> &Arc<dyn BlockDevice> {
        match self {
            Self::Uncached(uncache) => uncache.device(),
            Self::Cached(cache) => cache.device(),
        }
    }

    pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        pos: u64,
        mapper: F,
    ) -> Result<T, Error> {
        match self {
            Self::Uncached(uncache) => uncache.try_with(pos, mapper).await,
            Self::Cached(cache) => cache.try_with(pos, mapper).await,
        }
    }

    pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
        &self,
        pos: u64,
        size: usize,
        mapper: F,
    ) -> Result<T, Error> {
        match self {
            Self::Uncached(uncache) => uncache.try_with_mut(pos, size, mapper).await,
            Self::Cached(cache) => cache.try_with_mut(pos, size, mapper).await,
        }
    }

    pub async fn flush(&self) -> Result<(), Error> {
        match self {
            Self::Uncached(_) => Ok(()),
            Self::Cached(cache) => {
                cache.flush().await;
                Ok(())
            }
        }
    }
}

impl UncachedCache {
    pub fn device(&self) -> &Arc<dyn BlockDevice> {
        &self.device
    }

    pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        pos: u64,
        mapper: F,
    ) -> Result<T, Error> {
        if pos % self.block_size as u64 != 0 {
            log::warn!(
                "uncached: position {pos} is not a multiple of block size {}",
                self.block_size
            );
            return Err(Error::InvalidArgument);
        }
        let mut buffer = self.device.allocate_buffer(self.block_size)?;
        self.device
            .read_aligned(pos, buffer.slice_mut(0..self.block_size))
            .await?;
        let result = mapper(unsafe { buffer[..].assume_init_ref() })?;
        Ok(result)
    }

    pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
        &self,
        pos: u64,
        size: usize,
        mapper: F,
    ) -> Result<T, Error> {
        if pos % self.block_size as u64 != 0 {
            log::warn!(
                "uncached: position {pos} is not a multiple of block size {}",
                self.block_size
            );
            return Err(Error::InvalidArgument);
        }
        let mut buffer = self.device.allocate_buffer(self.block_size)?;
        // No need to read a block only to then fully rewrite it
        if size != self.block_size {
            self.device
                .read_aligned(pos, buffer.slice_mut(0..self.block_size))
                .await?;
        }
        let mut buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };
        let result = mapper(&mut buffer[..])?;
        self.device
            .write_aligned(pos, buffer.slice(0..self.block_size))
            .await?;
        Ok(result)
    }
}

impl BlockCache {
    pub fn with_capacity(
        device: Arc<dyn BlockDevice>,
        block_size: usize,
        segment_size: usize,
        bucket_capacity: usize,
        bucket_count: usize,
        filesystem: &str,
    ) -> Result<Self, Error> {
        if block_size % device.block_size() != 0 {
            log::error!(
                "Couldn't create block cache for {filesystem}: \
                cache block size is not a multiple of device block size"
            );
            return Err(Error::InvalidArgument);
        }
        if segment_size % block_size != 0 {
            log::error!(
                "Couldn't create block cache for {filesystem}: \
                segment size is not a multiple of block size"
            );
            return Err(Error::InvalidArgument);
        }
        if segment_size < block_size {
            log::error!(
                "Couldn't create block cache for {filesystem}: \
                segment size is smaller than block size"
            );
            return Err(Error::InvalidArgument);
        }

        log::info!("New block cache: block: {block_size}B, segment: {segment_size}B, geometry: {bucket_capacity}x{bucket_count}");
        log::info!(
            "Worst-case memory usage: {}K",
            (segment_size * bucket_capacity * bucket_count) / 1024
        );

        Ok(BlockCache {
            device,
            block_size,
            segment_size,
            cache: AsyncMutex::new(LruCache::with_capacity(bucket_capacity, bucket_count)),
        })
    }

    pub fn device(&self) -> &Arc<dyn BlockDevice> {
        &self.device
    }

    async fn evict_block(&self, segment_position: u64, block: Arc<IrqSafeRwLock<CachedSegment>>) {
        let read = block.read();
        if read.dirty {
            assert_eq!(segment_position % self.segment_size as u64, 0);
            if let Err(err) = self
                .device
                .write_aligned(segment_position, read.data.slice(0..self.segment_size))
                .await
            {
                log::error!("Disk error: flushing block {}: {:?}", segment_position, err);
            }
        }
    }

    async fn fetch_block(
        &self,
        segment_position: u64,
    ) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
        let mut buffer = self.device.allocate_buffer(self.segment_size)?;
        self.device
            .read_aligned(segment_position, buffer.slice_mut(0..self.segment_size))
            .await?;
        let data = unsafe { DmaBuffer::assume_init_slice(buffer) };
        Ok(Arc::new(IrqSafeRwLock::new(CachedSegment {
            data,
            dirty: false,
        })))
    }

    async fn entry(
        &self,
        segment_position: u64,
    ) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
        assert_eq!(segment_position % self.segment_size as u64, 0);
        let mut lock = self.cache.lock().await;
        let (value, evicted) = lock
            .try_get_or_insert_with_async(segment_position, || self.fetch_block(segment_position))
            .await?;

        if let Some((segment_position, block)) = evicted {
            self.evict_block(segment_position, block).await;
        }

        Ok(value.clone())
    }

    pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
        &self,
        block_position: u64,
        mapper: F,
    ) -> Result<T, Error> {
        if block_position % self.block_size as u64 != 0 {
            log::warn!(
                "mapper: position {block_position} is not a multiple of block size {}",
                self.block_size
            );
            return Err(Error::InvalidArgument);
        }
        let segment_position = block_position & !(self.segment_size as u64 - 1);
        let segment_offset = (block_position & (self.segment_size as u64 - 1)) as usize;
        let block = self.entry(segment_position).await?;
        let result = mapper(&block.read()[segment_offset..segment_offset + self.block_size])?;
        Ok(result)
    }

    pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
        &self,
        block_position: u64,
        _size: usize,
        mapper: F,
    ) -> Result<T, Error> {
        if block_position % self.block_size as u64 != 0 {
            log::warn!(
                "mapper: position {block_position} is not a multiple of block size {}",
                self.block_size
            );
            return Err(Error::InvalidArgument);
        }
        let segment_position = block_position & !(self.segment_size as u64 - 1);
        let segment_offset = (block_position & (self.segment_size as u64 - 1)) as usize;
        let block = self.entry(segment_position).await?;
        let mut block = block.write();
        block.dirty = true;
        let result = mapper(&mut block[segment_offset..segment_offset + self.block_size])?;
        Ok(result)
    }

    pub async fn flush(&self) {
        for (pos, block) in self.cache.lock().await.flush() {
            self.evict_block(pos, block).await;
        }
    }
}

impl CachedSegment {
    pub fn set_dirty(&mut self) {
        self.dirty = true;
    }
}

impl Deref for CachedSegment {
    type Target = DmaBuffer<[u8]>;

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

impl DerefMut for CachedSegment {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.dirty = true;
        &mut self.data
    }
}
