use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::device::Device;
use libk_mm::{
    address::PhysicalAddress, table::MapAttributes, OnDemandPage, PageProvider, VirtualPage,
};
use yggdrasil_abi::error::Error;

use crate::dma::{DmaBuffer, DmaSlice, DmaSliceMut};

use super::BlockDevice;

pub mod gpt;
pub mod mbr;

pub(crate) use gpt::probe_gpt;
pub(crate) use mbr::probe_mbr;

pub struct Partition {
    device: Arc<dyn BlockDevice>,
    lba_start: u64,
    lba_end: u64,
}

impl Partition {
    pub fn new(device: Arc<dyn BlockDevice>, lba_start: u64, lba_end: u64) -> Self {
        assert!(lba_start < lba_end);
        Self {
            device,
            lba_start,
            lba_end,
        }
    }
}

impl PageProvider for Partition {
    fn ondemand_fetch(&self, _opaque: u64) -> Result<OnDemandPage, Error> {
        unimplemented!()
    }

    fn get_page(&self, _offset: u64) -> Result<VirtualPage, Error> {
        unimplemented!()
    }

    fn release_page(
        &self,
        _offset: u64,
        _phys: PhysicalAddress,
        _dirty: bool,
    ) -> Result<(), Error> {
        unimplemented!()
    }

    fn clone_page(
        &self,
        _offset: u64,
        _src_phys: PhysicalAddress,
        _src_attrs: MapAttributes,
    ) -> Result<PhysicalAddress, Error> {
        unimplemented!()
    }
}

impl Device for Partition {
    fn display_name(&self) -> &str {
        "Partition"
    }
}

#[async_trait]
impl BlockDevice for Partition {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        self.device.allocate_buffer(size)
    }

    async fn read_aligned(
        &self,
        position: u64,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        // TODO check against partition bounds
        let lba = self.lba_start * self.block_size() as u64 + position;
        self.device.read_aligned(lba, buffer).await
    }

    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
        // TODO check against partition bounds
        let lba = self.lba_start * self.block_size() as u64 + position;
        self.device.write_aligned(lba, buffer).await
    }

    fn block_size(&self) -> usize {
        self.device.block_size()
    }

    fn block_count(&self) -> u64 {
        self.lba_end - self.lba_start
    }

    fn is_readable(&self) -> bool {
        self.device.is_readable()
    }

    fn is_writeable(&self) -> bool {
        self.device.is_writeable()
    }

    fn device_request(&self, option: u32, buffer: &mut [u8], len: usize) -> Result<usize, Error> {
        self.device.device_request(option, buffer, len)
    }

    fn max_blocks_per_request(&self) -> usize {
        self.device.max_blocks_per_request()
    }
}
