// TODO Split drivers into device and interface drivers
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, UsbInterfaceAccess},
    error::UsbError,
};

mod hid;
mod hub;
mod mass_storage;

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

    fn name(&self) -> &'static str;
    fn probe(&self, class: u8, subclass: u8, protocol: u8) -> bool;
}

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

    fn name(&self) -> &'static str;
    fn probe(&self, class: u8, subclass: u8, protocol: u8) -> bool;
}

async fn spawn_device_driver(device: Arc<UsbDeviceAccess>) -> Result<bool, UsbError> {
    let class = device.device_descriptor.device_class;
    let subclass = device.device_descriptor.device_subclass;
    let protocol = device.device_descriptor.device_protocol;
    let Some(driver) = USB_DEVICE_DRIVERS.read().iter().find_map(|driver| {
        driver
            .probe(class, subclass, protocol)
            .then(|| driver.clone())
    }) else {
        return Ok(false);
    };

    // if let Some(driver) = pick_driver(&device)? {
    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(true)
}

async fn spawn_interface_driver(interface: UsbInterfaceAccess) -> Result<(), UsbError> {
    let (class, subclass, protocol) = interface.class();
    let Some(driver) = USB_INTERFACE_DRIVERS.read().iter().find_map(|driver| {
        driver
            .probe(class, subclass, protocol)
            .then(|| driver.clone())
    }) else {
        return Ok(());
    };

    runtime::spawn(async move {
        let name = driver.name();
        match driver.run(interface).await {
            e @ Err(UsbError::DeviceDisconnected) => {
                log::warn!(
                    "Driver {:?} did not exit cleanly: device disconnected",
                    name,
                );

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

    Ok(())
}

pub async fn spawn_driver(device: Arc<UsbDeviceAccess>) {
    match spawn_device_driver(device.clone()).await {
        Ok(true) => return,
        Ok(false) => (),
        Err(error) => {
            log::error!(
                "{}: device driver probe failed: {:?}",
                device.bus_address(),
                error
            );
        }
    }
    // Enumerate interfaces
    for index in 0..device.configuration0_info.interfaces.len() as u8 {
        let access = device.interface(index);
        if let Err(error) = spawn_interface_driver(access).await {
            log::error!(
                "{}: interface {} driver probe failed: {:?}",
                device.bus_address(),
                index,
                error
            );
        }
    }
}

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

pub fn register_interface_driver(driver: Arc<dyn UsbInterfaceDriver + 'static>) {
    USB_INTERFACE_DRIVERS.write().push(driver);
}

pub fn register_default_class_drivers() {
    register_device_driver(Arc::new(hub::UsbHubDriver));
    register_interface_driver(Arc::new(hid::UsbHidKeyboardDriver));
    register_interface_driver(Arc::new(hid::UsbHidMouseDriver));
    register_interface_driver(Arc::new(mass_storage::UsbMassStorageDriverBulkOnly));
    // register_driver(Arc::new(hub::UsbHubDriver));
    // register_driver(Arc::new(hid_keyboard::UsbHidKeyboardDriver));
    // register_driver(Arc::new(mass_storage::UsbMassStorageDriverBulkOnly));
}

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