use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use libk_mm::{address::AsPhysicalAddress, PageSlice};
use ygg_driver_usb::{
    communication::UsbDirection,
    error::{TransferError, UsbError},
    pipe::{
        control::{ControlTransferSetup, UsbControlPipe},
        normal::{UsbNormalPipeIn, UsbNormalPipeOut},
    },
};

use crate::{controller::Xhci, ring::transfer::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(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(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(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,
        data: Option<(&mut PageSlice<MaybeUninit<u8>>, UsbDirection)>,
    ) -> Result<usize, UsbError> {
        let data_len = data.as_ref().map_or(0, |(data, _)| data.len());
        let result = self
            .ring
            .control_transfer(self.xhci.as_ref(), setup, data)
            .await;
        allow_short_packet(data_len, result)
    }
}

#[async_trait]
impl UsbNormalPipeIn for NormalInPipe {
    async fn read(&self, buffer: &mut PageSlice<u8>) -> Result<usize, UsbError> {
        let data_len = buffer.len();
        let result = self
            .ring
            .normal_transfer(
                self.xhci.as_ref(),
                unsafe { buffer.as_physical_address() },
                buffer.len(),
            )
            .await;
        allow_short_packet(data_len, result)
    }
}

#[async_trait]
impl UsbNormalPipeOut for NormalOutPipe {
    async fn write(&self, buffer: &PageSlice<u8>) -> Result<usize, UsbError> {
        self.ring
            .normal_transfer(
                self.xhci.as_ref(),
                unsafe { buffer.as_physical_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,
    }
}