use core::{any::Any, mem::MaybeUninit, ops::Deref};

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::device::Device;
use libk_mm::PageProvider;
use yggdrasil_abi::error::Error;

use crate::{
    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
    vfs::{CommonImpl, NodeRef},
};

pub mod cache;
pub mod partition;

struct Chunked {
    bs: usize,
    bpr: usize,

    position: u64,
    offset: usize,

    lba: u64,
    remaining_lba: usize,
    remaining: usize,
}

impl Chunked {
    pub fn begin<D: BlockDevice + ?Sized>(
        device: &D,
        position: u64,
        len: usize,
    ) -> Option<(Self, usize)> {
        let bc = device.block_count();
        let bs = device.block_size();
        let bpr = device.max_blocks_per_request();

        let start_lba = position / bs as u64;
        let end_lba = (position + len as u64).div_ceil(bs as u64).min(bc);

        if start_lba >= end_lba {
            return None;
        }

        let remaining_lba = (end_lba - start_lba) as usize;
        let remaining = len;
        let max_lba_count = remaining_lba.min(bpr);

        Some((
            Self {
                bs,
                bpr,
                position,
                lba: start_lba,
                remaining_lba,
                remaining,
                offset: 0,
            },
            max_lba_count,
        ))
    }
}

impl Iterator for Chunked {
    // 0: lba
    // 1: lba_count
    // 2: offset in block
    // 3: offset in buffer
    // 4: amount
    type Item = (u64, usize, usize, usize, usize);

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            assert_eq!(self.remaining_lba, 0);
            return None;
        }

        let block_count = self.remaining_lba.min(self.bpr);
        let block_offset = (self.position % self.bs as u64) as usize;
        let amount = (block_count * self.bs - block_offset).min(self.remaining);

        let item = (self.lba, block_count, block_offset, self.offset, amount);

        self.lba += block_count as u64;
        self.remaining_lba -= block_count;
        self.remaining -= amount;
        self.offset += amount;

        Some(item)
    }
}

#[async_trait]
pub trait BlockDevice: Device + PageProvider {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;

    async fn read_aligned(
        &self,
        position: u64,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        let _ = (position, buffer);
        Err(Error::NotImplemented)
    }
    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
        let _ = (position, buffer);
        Err(Error::NotImplemented)
    }

    async fn read(&self, position: u64, buffer: &mut [u8]) -> Result<usize, Error> {
        let bs = self.block_size();
        let Some((iter, max_lba_count)) = Chunked::begin(self, position, buffer.len()) else {
            return Ok(0);
        };
        let mut read_buffer = self.allocate_buffer(max_lba_count * bs)?;
        // let mut read_buffer = PageBox::new_uninit_slice(max_lba_count * bs)?;
        let mut total = 0;

        for (lba, block_count, block_offset, offset, amount) in iter {
            let slice = read_buffer.slice_mut(0..block_count * bs);
            self.read_aligned(lba * bs as u64, slice).await?;

            let src = unsafe {
                MaybeUninit::slice_assume_init_ref(
                    &read_buffer[block_offset..block_offset + amount],
                )
            };
            let dst = &mut buffer[offset..offset + amount];

            dst.copy_from_slice(src);

            total += amount;
        }

        Ok(total)
    }

    async fn write(&self, mut pos: u64, mut buf: &[u8]) -> Result<usize, Error> {
        let bs = self.block_size();

        let len = buf.len();
        let mut remaining = buf.len();

        while remaining != 0 {
            let block_offset = (pos % bs as u64) as usize;
            let lba = pos / bs as u64;

            let amount = core::cmp::min(bs - block_offset, buf.len());

            let mut dma_buffer = self.allocate_buffer(bs)?;
            if amount != bs {
                // Need to read the block first -- it's modified partially
                self.read_aligned(lba * bs as u64, dma_buffer.slice_mut(0..bs))
                    .await?;
            }

            let mut block = unsafe { DmaBuffer::assume_init_slice(dma_buffer) };

            block[block_offset..block_offset + amount].copy_from_slice(&buf[..amount]);

            // Write the block back
            self.write_aligned(lba * bs as u64, block.slice(0..bs))
                .await?;

            buf = &buf[amount..];
            remaining -= amount;
            pos += amount as u64;
        }

        Ok(len)
    }

    async fn read_exact(&self, position: u64, buffer: &mut [u8]) -> Result<(), Error> {
        match self.read(position, buffer).await {
            Ok(len) if len == buffer.len() => Ok(()),
            Ok(_) => Err(Error::InvalidArgument),
            Err(error) => Err(error),
        }
    }

    async fn write_exact(&self, position: u64, buffer: &[u8]) -> Result<(), Error> {
        match self.write(position, buffer).await {
            Ok(len) if len == buffer.len() => Ok(()),
            Ok(_) => Err(Error::InvalidArgument),
            Err(error) => Err(error),
        }
    }

    fn is_readable(&self) -> bool {
        true
    }
    fn is_writeable(&self) -> bool {
        true
    }

    fn block_size(&self) -> usize;
    fn block_count(&self) -> u64;
    fn max_blocks_per_request(&self) -> usize;

    fn device_request(&self, option: u32, buffer: &mut [u8], len: usize) -> Result<usize, Error> {
        let _ = option;
        let _ = buffer;
        let _ = len;
        Err(Error::InvalidOperation)
    }
}

#[derive(Clone)]
pub struct BlockDeviceFile(pub(crate) Arc<dyn BlockDevice>);

impl CommonImpl for BlockDeviceFile {
    fn size(&self, node: &NodeRef) -> Result<u64, Error> {
        let _ = node;
        Ok(self.0.block_count() * self.0.block_size() as u64)
    }

    fn as_any(&self) -> &dyn Any {
        self as _
    }
}

impl Deref for BlockDeviceFile {
    type Target = dyn BlockDevice;

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