use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::device::Device;
use libk::{
    device::block::BlockDevice,
    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
    error::Error,
};
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    table::MapAttributes,
    OnDemandPage, PageProvider, PageSlice, VirtualPage,
};

use crate::{command::IdentifyNamespaceRequest, register_nvme_namespace, IoDirection};

use super::{error::NvmeError, NvmeController};

#[allow(unused)]
pub struct NvmeNamespace {
    controller: Arc<NvmeController>,
    nsid: u32,
    total_lba_count: u64,
    lba_size: u64,
    max_lba_per_request: usize,
}

impl NvmeNamespace {
    pub async fn create(
        controller: Arc<NvmeController>,
        nsid: u32,
        max_transfer_size: usize,
    ) -> Result<Arc<NvmeNamespace>, NvmeError> {
        let admin_q = controller.admin_q.get();
        let identify = admin_q
            .request(&*controller.dma, IdentifyNamespaceRequest { nsid })
            .await?;

        let current_lba_format_idx = identify.current_lba_fmt_idx();
        let current_lba_format = identify.lba_fmt(current_lba_format_idx).unwrap();
        let lba_size = current_lba_format.lba_data_size().unwrap();
        let total_lba_count = identify.total_lba_count();
        let max_lba_per_request = (max_transfer_size / lba_size as usize).min(512);

        log::debug!(
            "ns = {}, lba = {}B, size = {}M, max lba/req = {}",
            nsid,
            lba_size,
            (total_lba_count * lba_size) / (1024 * 1024),
            max_lba_per_request,
        );
        let dev = NvmeNamespace {
            controller,
            nsid,
            total_lba_count,
            lba_size,
            max_lba_per_request,
        };
        let dev = Arc::new(dev);

        register_nvme_namespace(dev.clone(), true);

        Ok(dev)
    }

    pub fn controller_id(&self) -> u32 {
        *self.controller.controller_id.get()
    }

    pub fn id(&self) -> u32 {
        self.nsid
    }
}

impl Device for NvmeNamespace {
    fn display_name(&self) -> &str {
        "NVMe Namespace"
    }
}

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

    // TODO read directly to cache
    async fn read_aligned(
        &self,
        position: u64,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        if position % self.block_size() as u64 != 0 {
            return Err(Error::InvalidOperation);
        }
        if buffer.len() % self.block_size() != 0 || buffer.is_empty() {
            return Err(Error::InvalidOperation);
        }
        let lba = position / self.block_size() as u64;
        let lba_count = buffer.len().div_ceil(self.block_size());
        if lba + lba_count as u64 > self.block_count() {
            return Err(Error::InvalidOperation);
        }

        let result = self
            .controller
            .perform_read(self.nsid, lba, lba_count, buffer)
            .await;

        log::trace!(target: "io", "read #{lba}, {lba_count} blocks -> {result:?}");

        result.map_err(NvmeError::into)
    }

    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
        if position % self.block_size() as u64 != 0 {
            return Err(Error::InvalidOperation);
        }
        if buffer.len() % self.block_size() != 0 || buffer.is_empty() {
            return Err(Error::InvalidOperation);
        }
        let lba = position / self.block_size() as u64;
        let lba_count = buffer.len().div_ceil(self.block_size());
        if lba + lba_count as u64 > self.block_count() {
            return Err(Error::InvalidOperation);
        }

        let result = self
            .controller
            .perform_write(self.nsid, lba, lba_count, buffer)
            .await;

        log::trace!(target: "io", "write -> #{lba}, {lba_count} blocks -> {result:?}");

        result.map_err(NvmeError::into)
    }

    fn block_size(&self) -> usize {
        self.lba_size as _
    }

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

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

impl PageProvider for NvmeNamespace {
    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!()
    }
}
