use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc};
use async_trait::async_trait;
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use xhci_lib::context;
use ygg_driver_usb::{
    communication::UsbDirection,
    device::{UsbBusAddress, UsbDevice, UsbDeviceDetachHandler, UsbSpeed},
    error::UsbError,
    info::UsbEndpointType,
    pipe::{
        control::UsbControlPipeAccess,
        normal::{UsbNormalPipeIn, UsbNormalPipeOut},
    },
    UsbHostController,
};

use crate::{
    context::{XhciDeviceContext, XhciInputContext},
    pipe::{NormalInPipe, NormalOutPipe},
    regs::PortNumber,
    ring::transfer::TransferRing,
    Xhci,
};

pub struct XhciBusDevice {
    pub(crate) port_id: PortNumber,
    pub(crate) slot_id: u8,
    pub(crate) speed: UsbSpeed,
    pub(crate) bus_address: OneTimeInit<UsbBusAddress>,

    pub(crate) xhci: Arc<Xhci>,

    pub(crate) device_context: XhciDeviceContext,
    pub(crate) endpoints: IrqSafeRwLock<BTreeMap<u8, Arc<TransferRing>>>,
    pub(crate) control_pipe: UsbControlPipeAccess,
    pub(crate) detach_handler: IrqSafeSpinlock<Option<Arc<dyn UsbDeviceDetachHandler>>>,
}

#[async_trait]
impl UsbDevice for XhciBusDevice {
    fn port_number(&self) -> u8 {
        self.port_id.into()
    }

    fn bus_address(&self) -> UsbBusAddress {
        *self.bus_address.get()
    }

    fn control_pipe(&self) -> &UsbControlPipeAccess {
        &self.control_pipe
    }

    fn handle_detach(&self) {
        if let Some(handler) = self.detach_handler.lock().as_ref() {
            handler.handle_device_detach();
        }
    }

    fn set_detach_handler(&self, handler: Arc<dyn UsbDeviceDetachHandler>) {
        *self.detach_handler.lock() = Some(handler);
    }

    fn controller_ref(&self) -> &dyn UsbHostController {
        self.xhci.as_ref()
    }

    fn debug(&self) {}

    fn speed(&self) -> UsbSpeed {
        self.speed
    }

    async fn open_normal_in_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
        ty: UsbEndpointType,
    ) -> Result<Box<dyn UsbNormalPipeIn>, UsbError> {
        let dci = Self::dci(ty, UsbDirection::In, number) as u8;
        let (pipe, ring) = NormalInPipe::new(self.xhci.clone(), self.slot_id, dci, 128)?;
        self.add_endpoint(number, ring.clone(), ty, UsbDirection::In, max_packet_size)
            .await?;
        Ok(Box::new(pipe))
    }

    async fn open_normal_out_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
        ty: UsbEndpointType,
    ) -> Result<Box<dyn UsbNormalPipeOut>, UsbError> {
        let dci = Self::dci(ty, UsbDirection::Out, number) as u8;
        let (pipe, ring) = NormalOutPipe::new(self.xhci.clone(), self.slot_id, dci, 128)?;
        self.add_endpoint(number, ring.clone(), ty, UsbDirection::Out, max_packet_size)
            .await?;
        Ok(Box::new(pipe))
    }
}

impl XhciBusDevice {
    pub async fn add_endpoint(
        &self,
        endpoint_number: u8,
        ring: Arc<TransferRing>,
        ty: UsbEndpointType,
        direction: UsbDirection,
        max_packet_size: u16,
    ) -> Result<(), UsbError> {
        let dci = Self::dci(ty, direction, endpoint_number) as u8;

        log::debug!(
            "Setup endpoint dci #{}: {} {:?} {:?}",
            dci,
            endpoint_number,
            ty,
            direction
        );

        let ep_type = match (ty, direction) {
            (UsbEndpointType::Interrupt, UsbDirection::In) => context::EndpointType::InterruptIn,
            (UsbEndpointType::Interrupt, UsbDirection::Out) => context::EndpointType::InterruptOut,
            (UsbEndpointType::Bulk, UsbDirection::In) => context::EndpointType::BulkIn,
            (UsbEndpointType::Bulk, UsbDirection::Out) => context::EndpointType::BulkOut,
            (UsbEndpointType::Isochronous, UsbDirection::In) => context::EndpointType::IsochIn,
            (UsbEndpointType::Isochronous, UsbDirection::Out) => context::EndpointType::IsochOut,
            (UsbEndpointType::Control, _) => context::EndpointType::Control,
        };

        let mut input_cx = XhciInputContext::new(&*self.xhci.dma, self.xhci.regs.context_size)?;

        {
            let control = input_cx.control_mut();

            control.set_add_context_flag(0);
            control.set_add_context_flag(dci as usize);
        }

        {
            let slot = input_cx.device_mut().slot_mut();

            slot.set_context_entries(dci + 1);
        }

        {
            let endpoint = input_cx.device_mut().endpoint_mut(dci as usize);

            endpoint.set_endpoint_type(ep_type);
            endpoint.set_error_count(3);
            // TODO Pick a better value here
            endpoint.set_interval(3);
            endpoint.set_max_packet_size(max_packet_size);
            endpoint.set_tr_dequeue_pointer(ring.bus_address().into_u64());
            endpoint.set_dequeue_cycle_state();
        }

        self.xhci
            .command_ring
            .configure_endpoint(self.xhci.as_ref(), self.slot_id, &input_cx)
            .await?;

        self.endpoints.write().insert(dci, ring.clone());
        self.xhci
            .endpoints
            .write()
            .insert((self.slot_id, dci), ring);

        Ok(())
    }

    fn dci(ty: UsbEndpointType, dir: UsbDirection, number: u8) -> usize {
        match ty {
            UsbEndpointType::Control => (number as usize) * 2 + 1,
            _ => {
                let off = match dir {
                    UsbDirection::Out => 0,
                    UsbDirection::In => 1,
                };

                (number as usize * 2) + off
            }
        }
    }
}
