use core::{fmt, ops::Deref};

use alloc::{boxed::Box, sync::Arc, vec::Vec};
use async_trait::async_trait;
use libk_util::sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard};

use crate::{
    error::UsbError,
    info::{
        UsbConfigurationInfo, UsbDeviceInfo, UsbEndpointInfo, UsbEndpointType, UsbInterfaceInfo,
        UsbVersion,
    },
    pipe::{
        control::{ConfigurationDescriptorEntry, UsbControlPipeAccess},
        normal::{
            UsbBulkInPipeAccess, UsbBulkOutPipeAccess, UsbInterruptInPipeAccess, UsbNormalPipeIn,
            UsbNormalPipeOut,
        },
    },
    UsbHostController,
};

// High-level structures for info provided through descriptors

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub struct UsbBusAddress {
    pub bus: u16,
    pub device: u8,
}

pub struct UsbDeviceAccess {
    pub device: Arc<dyn UsbDevice>,
    pub info: UsbDeviceInfo,
    pub current_configuration: IrqSafeRwLock<Option<UsbConfigurationInfo>>,
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub enum UsbSpeed {
    Low,
    Full,
    High,
    Super,
}

pub trait UsbDeviceDetachHandler: Send + Sync {
    fn handle_device_detach(&self);
}

#[async_trait]
#[allow(unused)]
pub trait UsbDevice: Send + Sync {
    // Endpoint "0"
    fn control_pipe(&self) -> &UsbControlPipeAccess;

    async fn open_normal_in_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
        ty: UsbEndpointType,
    ) -> Result<Box<dyn UsbNormalPipeIn>, UsbError>;

    async fn open_normal_out_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
        ty: UsbEndpointType,
    ) -> Result<Box<dyn UsbNormalPipeOut>, UsbError>;

    fn port_number(&self) -> u8;
    fn bus_address(&self) -> UsbBusAddress;
    fn speed(&self) -> UsbSpeed;
    fn controller_ref(&self) -> &dyn UsbHostController;

    fn set_detach_handler(&self, handler: Arc<dyn UsbDeviceDetachHandler>);
    fn handle_detach(&self);

    fn debug(&self) {}
}

impl UsbDeviceAccess {
    /// Expected device state:
    ///
    /// * Link-layer stuff has been reset and established properly by the HCD
    /// * Device is not yet configured
    /// * Control pipe for the device has been properly set up
    /// * Device has been assigned a bus address
    pub async fn setup(raw: Arc<dyn UsbDevice>) -> Result<Self, UsbError> {
        let control = raw.control_pipe();

        let device_desc = control.query_device_descriptor().await?;

        let bcd_usb = device_desc.bcd_usb;
        let usb_version = UsbVersion::from_bcd_usb(device_desc.bcd_usb)
            .ok_or(UsbError::InvalidDescriptorField)
            .inspect_err(|_| {
                log::error!(
                    "{}: unsupported/invalid USB version: {:#x}",
                    raw.bus_address(),
                    bcd_usb
                )
            })?;

        let manufacturer = control.query_string(device_desc.manufacturer_str).await?;
        let product = control.query_string(device_desc.product_str).await?;

        let info = UsbDeviceInfo {
            manufacturer,
            product,
            usb_version,

            id_vendor: device_desc.id_vendor,
            id_product: device_desc.id_product,

            device_class: device_desc.class(),
            device_subclass: device_desc.device_subclass,
            device_protocol: device_desc.protocol(),
            device_protocol_number: device_desc.device_protocol,

            num_configurations: device_desc.num_configurations,

            max_packet_size: device_desc.max_packet_size(usb_version, raw.speed())?,
        };

        Ok(Self {
            device: raw,
            info,
            current_configuration: IrqSafeRwLock::new(None),
        })
    }

    pub async fn open_interrupt_in_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
    ) -> Result<UsbInterruptInPipeAccess, UsbError> {
        let pipe = self
            .device
            .open_normal_in_pipe(number, max_packet_size, UsbEndpointType::Interrupt)
            .await?;
        Ok(UsbInterruptInPipeAccess(pipe))
    }

    pub async fn open_bulk_in_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
    ) -> Result<UsbBulkInPipeAccess, UsbError> {
        let pipe = self
            .device
            .open_normal_in_pipe(number, max_packet_size, UsbEndpointType::Bulk)
            .await?;
        Ok(UsbBulkInPipeAccess(pipe))
    }

    pub async fn open_bulk_out_pipe(
        &self,
        number: u8,
        max_packet_size: u16,
    ) -> Result<UsbBulkOutPipeAccess, UsbError> {
        let pipe = self
            .device
            .open_normal_out_pipe(number, max_packet_size, UsbEndpointType::Bulk)
            .await?;
        Ok(UsbBulkOutPipeAccess(pipe))
    }

    pub fn read_current_configuration(
        &self,
    ) -> IrqSafeRwLockReadGuard<'_, Option<UsbConfigurationInfo>> {
        self.current_configuration.read()
    }

    pub async fn select_configuration<F: Fn(&UsbConfigurationInfo) -> bool>(
        &self,
        predicate: F,
    ) -> Result<Option<UsbConfigurationInfo>, UsbError> {
        let mut current_config = self.current_configuration.write();
        let control_pipe = self.control_pipe();

        for i in 0..self.info.num_configurations {
            let info = self.query_configuration_info(i).await?;

            if predicate(&info) {
                log::debug!("Selected configuration: {:#?}", info);
                let config = current_config.insert(info);

                control_pipe
                    .set_configuration(config.config_value as _)
                    .await?;

                return Ok(Some(config.clone()));
            }
        }

        Ok(None)
    }

    pub async fn query_configuration_info(
        &self,
        index: u8,
    ) -> Result<UsbConfigurationInfo, UsbError> {
        if index >= self.info.num_configurations {
            return Err(UsbError::InvalidConfiguration);
        }
        let control_pipe = self.control_pipe();
        let query = control_pipe.query_configuration_descriptor(index).await?;

        let configuration_name = control_pipe
            .query_string(query.configuration().config_str)
            .await?;

        let mut endpoints = Vec::new();
        let mut interfaces = Vec::new();

        for desc in query.descriptors() {
            match desc {
                ConfigurationDescriptorEntry::Endpoint(ep) => {
                    endpoints.push(UsbEndpointInfo {
                        number: ep.number(),
                        direction: ep.direction(),
                        max_packet_size: ep.max_packet_size as _,
                        ty: ep.transfer_type(),
                    });
                }
                ConfigurationDescriptorEntry::Interface(iface) => {
                    let name = control_pipe.query_string(iface.interface_str).await?;
                    interfaces.push(UsbInterfaceInfo {
                        name,
                        number: iface.interface_number,

                        interface_class: iface.class(),
                        interface_subclass: iface.interface_subclass,
                        interface_protocol: iface.protocol(),
                        interface_protocol_number: iface.interface_protocol,
                    });
                }
                _ => (),
            }
        }

        let info = UsbConfigurationInfo {
            name: configuration_name,
            config_value: query.configuration().config_val,
            interfaces,
            endpoints,
        };

        Ok(info)
    }

    pub fn set_detach_handler(&self, handler: Arc<dyn UsbDeviceDetachHandler>) {
        self.device.set_detach_handler(handler);
    }
}

impl Deref for UsbDeviceAccess {
    type Target = dyn UsbDevice;

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

impl fmt::Display for UsbBusAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}", self.bus, self.device)
    }
}
