use core::sync::atomic::{AtomicU16, Ordering};

use alloc::{collections::BTreeMap, sync::Arc};
use libk_util::{queue::UnboundedMpmcQueue, sync::spin_rwlock::IrqSafeRwLock};

use crate::{
    class_driver,
    device::{UsbBusAddress, UsbDeviceAccess},
    UsbHostController,
};

pub struct UsbBusManager {
    busses: IrqSafeRwLock<BTreeMap<u16, Arc<dyn UsbHostController>>>,
    devices: IrqSafeRwLock<BTreeMap<UsbBusAddress, Arc<UsbDeviceAccess>>>,

    last_bus_address: AtomicU16,
}

impl UsbBusManager {
    pub fn register_bus(hc: Arc<dyn UsbHostController>) -> u16 {
        let i = BUS_MANAGER.last_bus_address.fetch_add(1, Ordering::AcqRel);
        BUS_MANAGER.busses.write().insert(i, hc);
        i
    }

    pub fn register_device(device: Arc<UsbDeviceAccess>) {
        log::info!("usb: register device {}", device.bus_address());
        BUS_MANAGER
            .devices
            .write()
            .insert(device.bus_address(), device.clone());

        QUEUE.push_back(device);
    }

    pub fn detach_device(address: UsbBusAddress) {
        if let Some(device) = BUS_MANAGER.devices.write().remove(&address) {
            device.handle_detach();
        }
    }
}

pub async fn bus_handler() {
    class_driver::register_default_class_drivers();

    loop {
        let new_device = QUEUE.pop_front().await;
        log::info!(
            "New {:?}-speed USB device connected: {}",
            new_device.speed(),
            new_device.bus_address()
        );

        class_driver::spawn_driver(new_device).await.ok();
    }
}

static BUS_MANAGER: UsbBusManager = UsbBusManager {
    busses: IrqSafeRwLock::new(BTreeMap::new()),
    devices: IrqSafeRwLock::new(BTreeMap::new()),

    last_bus_address: AtomicU16::new(0),
};
static QUEUE: UnboundedMpmcQueue<Arc<UsbDeviceAccess>> = UnboundedMpmcQueue::new();
