use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use libk::dma::{DmaBuffer, DmaSlice, DmaSliceMut};
use ygg_driver_usb::{
    error::{TransferError, UsbError},
    pipe::{
        control::{ControlTransferSetup, UsbControlPipe},
        normal::{UsbNormalPipeIn, UsbNormalPipeOut},
        UsbGenericPipe,
    },
};

use crate::{
    controller::Xhci,
    ring::transfer::{ControlDataStage, TransferRing},
};

pub struct ControlPipe {
    xhci: Arc<Xhci>,
    pub(crate) ring: Arc<TransferRing>,
}

pub struct NormalInPipe {
    xhci: Arc<Xhci>,
    pub(crate) ring: Arc<TransferRing>,
}

pub struct NormalOutPipe {
    xhci: Arc<Xhci>,
    pub(crate) ring: Arc<TransferRing>,
}

impl ControlPipe {
    pub fn new(
        xhci: Arc<Xhci>,
        slot_id: u8,
        endpoint_id: u8,
        capacity: usize,
    ) -> Result<(Self, Arc<TransferRing>), UsbError> {
        let ring = Arc::new(TransferRing::new(
            &*xhci.dma,
            slot_id,
            endpoint_id,
            capacity,
        )?);
        Ok((
            Self {
                xhci,
                ring: ring.clone(),
            },
            ring,
        ))
    }
}

impl NormalInPipe {
    pub fn new(
        xhci: Arc<Xhci>,
        slot_id: u8,
        endpoint_id: u8,
        capacity: usize,
    ) -> Result<(Self, Arc<TransferRing>), UsbError> {
        let ring = Arc::new(TransferRing::new(
            &*xhci.dma,
            slot_id,
            endpoint_id,
            capacity,
        )?);
        Ok((
            Self {
                xhci,
                ring: ring.clone(),
            },
            ring,
        ))
    }
}

impl NormalOutPipe {
    pub fn new(
        xhci: Arc<Xhci>,
        slot_id: u8,
        endpoint_id: u8,
        capacity: usize,
    ) -> Result<(Self, Arc<TransferRing>), UsbError> {
        let ring = Arc::new(TransferRing::new(
            &*xhci.dma,
            slot_id,
            endpoint_id,
            capacity,
        )?);
        Ok((
            Self {
                xhci,
                ring: ring.clone(),
            },
            ring,
        ))
    }
}

#[async_trait]
impl UsbControlPipe for ControlPipe {
    async fn control_transfer(&self, setup: ControlTransferSetup) -> Result<(), UsbError> {
        self.ring
            .control_transfer(self.xhci.as_ref(), setup, ControlDataStage::None)
            .await?;
        Ok(())
    }

    async fn control_transfer_in(
        &self,
        setup: ControlTransferSetup,
        buffer: &mut [MaybeUninit<u8>],
    ) -> Result<usize, UsbError> {
        let data_len = buffer.len();
        let mut dma_buffer = DmaBuffer::new_uninit_slice(&*self.xhci.dma, data_len)
            .map_err(UsbError::MemoryError)?;

        let result = self
            .ring
            .control_transfer(
                self.xhci.as_ref(),
                setup,
                ControlDataStage::In(&mut dma_buffer),
            )
            .await;
        let result = allow_short_packet(data_len, result);

        match result {
            Ok(len) => {
                buffer[..len].copy_from_slice(&dma_buffer[..len]);
                Ok(len)
            }
            Err(error) => Err(error),
        }
    }

    async fn control_transfer_out(
        &self,
        setup: ControlTransferSetup,
        buffer: &[u8],
    ) -> Result<usize, UsbError> {
        let dma_buffer =
            DmaBuffer::from_slice(&*self.xhci.dma, buffer).map_err(UsbError::MemoryError)?;

        self.ring
            .control_transfer(
                self.xhci.as_ref(),
                setup,
                ControlDataStage::Out(&dma_buffer),
            )
            .await
    }
}

impl UsbGenericPipe for NormalInPipe {
    fn allocate_dma_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, UsbError> {
        DmaBuffer::new_uninit_slice(&*self.xhci.dma, len).map_err(UsbError::MemoryError)
    }
}

#[async_trait]
impl UsbNormalPipeIn for NormalInPipe {
    async fn read_dma(&self, buffer: DmaSliceMut<'_, MaybeUninit<u8>>) -> Result<usize, UsbError> {
        let result = self
            .ring
            .normal_transfer(self.xhci.as_ref(), buffer.bus_address(), buffer.len())
            .await;
        allow_short_packet(buffer.len(), result)
    }
}

impl UsbGenericPipe for NormalOutPipe {
    fn allocate_dma_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, UsbError> {
        DmaBuffer::new_uninit_slice(&*self.xhci.dma, len).map_err(UsbError::MemoryError)
    }
}

#[async_trait]
impl UsbNormalPipeOut for NormalOutPipe {
    async fn write_dma(&self, buffer: DmaSlice<'_, u8>) -> Result<usize, UsbError> {
        self.ring
            .normal_transfer(self.xhci.as_ref(), buffer.bus_address(), buffer.len())
            .await
    }
}

fn allow_short_packet(data_len: usize, result: Result<usize, UsbError>) -> Result<usize, UsbError> {
    match result {
        // Short packets are okay for control transfers
        Err(UsbError::TransferFailed(TransferError::ShortPacket(residual))) => {
            Ok(data_len.saturating_sub(residual))
        }
        result => result,
    }
}
