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

use crate::{
    device::UsbDeviceAccess,
    error::UsbError,
    info::{UsbDeviceClass, UsbDeviceProtocol},
};

pub mod hid_keyboard;
pub mod mass_storage;

#[derive(Debug)]
pub struct UsbClassInfo {
    pub class: UsbDeviceClass,
    pub subclass: u8,
    pub protocol: UsbDeviceProtocol,
    pub device_protocol_number: u8,
    pub interface_protocol_number: u8,
}

#[async_trait]
pub trait UsbDriver: Send + Sync {
    async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError>;

    fn name(&self) -> &'static str;
    fn probe(&self, class: &UsbClassInfo, device: &UsbDeviceAccess) -> bool;
}

async fn extract_class_info(device: &UsbDeviceAccess) -> Result<Option<UsbClassInfo>, UsbError> {
    if device.info.num_configurations != 1 {
        return Ok(None);
    }
    let device_info = &device.info;
    let config_info = device.query_configuration_info(0).await?;

    if !config_info.interfaces.is_empty() {
        let if_info = &config_info.interfaces[0];

        let class = if device_info.device_class == UsbDeviceClass::FromInterface {
            if_info.interface_class
        } else {
            device_info.device_class
        };
        let subclass = if device_info.device_subclass == 0 {
            if_info.interface_subclass
        } else {
            device_info.device_subclass
        };
        let protocol = if device_info.device_protocol == UsbDeviceProtocol::FromInterface {
            if_info.interface_protocol
        } else {
            device_info.device_protocol
        };

        Ok(Some(UsbClassInfo {
            class,
            subclass,
            protocol,
            interface_protocol_number: if_info.interface_protocol_number,
            device_protocol_number: device_info.device_protocol_number,
        }))
    } else {
        Ok(None)
    }
}

async fn pick_driver(
    device: &UsbDeviceAccess,
) -> Result<Option<Arc<dyn UsbDriver + 'static>>, UsbError> {
    let Some(class) = extract_class_info(device).await? else {
        return Ok(None);
    };

    for driver in USB_DEVICE_DRIVERS.read().iter() {
        if driver.probe(&class, device) {
            return Ok(Some(driver.clone()));
        }
    }
    Ok(None)
}

pub async fn spawn_driver(device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
    if let Some(driver) = pick_driver(&device).await? {
        runtime::spawn(async move {
            let name = driver.name();
            match driver.run(device).await {
                e @ Err(UsbError::DeviceDisconnected) => {
                    log::warn!(
                        "Driver {:?} did not exit cleanly: device disconnected",
                        name,
                    );

                    e
                }
                e => e,
            }
        })
        .map_err(UsbError::SystemError)?;
    }
    Ok(())
}

pub fn register_driver(driver: Arc<dyn UsbDriver + 'static>) {
    // TODO check for duplicates
    USB_DEVICE_DRIVERS.write().push(driver);
}

pub fn register_default_class_drivers() {
    register_driver(Arc::new(hid_keyboard::UsbHidKeyboardDriver));
    register_driver(Arc::new(mass_storage::UsbMassStorageDriverBulkOnly));
}

static USB_DEVICE_DRIVERS: IrqSafeRwLock<Vec<Arc<dyn UsbDriver + 'static>>> =
    IrqSafeRwLock::new(Vec::new());
