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

use alloc::boxed::Box;
use async_trait::async_trait;
use libk::dma::{DmaBuffer, DmaSlice, DmaSliceMut};

use crate::error::{TransferError, UsbError};

use super::UsbGenericPipe;

#[async_trait]
pub trait UsbNormalPipeIn: UsbGenericPipe {
    async fn read_dma(&self, buffer: DmaSliceMut<'_, MaybeUninit<u8>>) -> Result<usize, UsbError>;

    async fn read(&self, buffer: &mut [u8]) -> Result<usize, UsbError> {
        let mut dma_buffer = self.allocate_dma_buffer(buffer.len())?;
        let len = self.read_dma(dma_buffer.slice_mut(0..buffer.len())).await?;
        let dma_slice = unsafe { dma_buffer[..len].assume_init_ref() };
        buffer[..len].copy_from_slice(dma_slice);
        Ok(len)
    }

    async fn read_exact(&self, buffer: &mut [u8]) -> Result<(), UsbError> {
        match self.read(buffer).await {
            Ok(len) if len == buffer.len() => Ok(()),
            Ok(len) => Err(UsbError::TransferFailed(TransferError::ShortPacket(len))),
            Err(error) => Err(error),
        }
    }
}

#[async_trait]
pub trait UsbNormalPipeOut: UsbGenericPipe {
    async fn write_dma(&self, buffer: DmaSlice<'_, u8>) -> Result<usize, UsbError>;

    async fn write(&self, buffer: &[u8]) -> Result<usize, UsbError> {
        let mut dma_buffer = self.allocate_dma_buffer(buffer.len())?;
        dma_buffer[..].write_copy_of_slice(buffer);
        let dma_buffer = unsafe { DmaBuffer::assume_init_slice(dma_buffer) };

        self.write_dma(dma_buffer.slice(0..buffer.len())).await
    }
}

pub struct UsbBulkInPipeAccess(pub Box<dyn UsbNormalPipeIn>);
pub struct UsbBulkOutPipeAccess(pub Box<dyn UsbNormalPipeOut>);

pub struct UsbInterruptInPipeAccess(pub Box<dyn UsbNormalPipeIn>);
pub struct UsbInterruptOutPipeAccess(pub Box<dyn UsbNormalPipeOut>);

impl Deref for UsbBulkInPipeAccess {
    type Target = dyn UsbNormalPipeIn;

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

impl Deref for UsbBulkOutPipeAccess {
    type Target = dyn UsbNormalPipeOut;

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

impl Deref for UsbInterruptInPipeAccess {
    type Target = dyn UsbNormalPipeIn;

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

impl Deref for UsbInterruptOutPipeAccess {
    type Target = dyn UsbNormalPipeOut;

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