use core::{mem::MaybeUninit, num::NonZeroU8, time::Duration};

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use libk::{task::runtime, time::monotonic_time};
use yggdrasil_abi::bitflags;

use crate::{
    class_driver::UsbDeviceDriver,
    descriptor::UsbHubDescriptorHeader,
    device::{UsbDeviceAccess, UsbSpeed},
    error::UsbError,
    pipe::control::ControlTransferSetup,
};

pub struct UsbHubDriver;

const BM_REQUEST_TYPE_GET_HUB_DESCRIPTOR: u8 = 0b10100000;
const BM_REQUEST_TYPE_GET_PORT_STATUS: u8 = 0b10100011;
const BM_REQUEST_TYPE_PORT_FEATURE: u8 = 0b00100011;

const B_REQUEST_GET_STATUS: u8 = 0x00;
const B_REQUEST_CLEAR_FEATURE: u8 = 0x01;
const B_REQUEST_SET_FEATURE: u8 = 0x03;
const B_REQUEST_GET_DESCRIPTOR: u8 = 0x06;

#[allow(unused)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u16)]
pub enum PortFeatureSelector {
    PortConnection = 0,
    PortEnable = 1,
    PortSuspend = 2,
    PortOverCurrent = 3,
    PortReset = 4,
    PortPower = 8,
    PortLowSpeed = 9,
    CPortConnection = 16,
    CPortEnable = 17,
    CPortSuspend = 18,
    CPortOverCurrent = 19,
    CPortReset = 20,
    PortTest = 21,
    PortIndicator = 22,
}

bitflags! {
    pub struct UsbHubPortStatus: u16 {
        const PORT_CONNECTION: bit 0;
        const PORT_ENABLE: bit 1;
        const PORT_SUSPEND: bit 2;
        const PORT_OVER_CURRENT: bit 3;
        const PORT_RESET: bit 4;
        const PORT_POWER: bit 8;
        const PORT_LOW_SPEED: bit 9;
        const PORT_HIGH_SPEED: bit 10;
        const PORT_TEST: bit 11;
        const PORT_INDICATOR: bit 12;
    }
}
bitflags! {
    pub struct UsbHubPortChange: u16 {
        const C_PORT_CONNECTION: bit 0;
        const C_PORT_ENABLE: bit 1;
        const C_PORT_SUSPEND: bit 2;
        const C_PORT_OVER_CURRENT: bit 3;
        const C_PORT_RESET: bit 4;
    }
}

struct UsbHub {
    device: Arc<UsbDeviceAccess>,
    hub_descriptor: UsbHubDescriptorHeader,
    children: u64,
    // TODO extra info
}

impl UsbHub {
    pub async fn setup(device: Arc<UsbDeviceAccess>) -> Result<Self, UsbError> {
        let control_pipe = device.control_pipe();
        let mut hub_descriptor = MaybeUninit::uninit();

        let len = control_pipe
            .control_transfer_in(
                ControlTransferSetup {
                    bm_request_type: BM_REQUEST_TYPE_GET_HUB_DESCRIPTOR,
                    b_request: B_REQUEST_GET_DESCRIPTOR,
                    w_value: 0,
                    w_index: 0,
                    w_length: size_of::<UsbHubDescriptorHeader>() as _,
                },
                hub_descriptor.as_bytes_mut(),
            )
            .await?;
        if len != size_of::<UsbHubDescriptorHeader>() {
            return Err(UsbError::TruncatedDescriptor(
                size_of::<UsbHubDescriptorHeader>(),
                len,
            ));
        }
        let hub_descriptor: UsbHubDescriptorHeader = unsafe { hub_descriptor.assume_init() };
        if hub_descriptor.b_nr_ports < 1 {
            log::warn!(
                "{}: ignoring hub with zero downstream ports",
                device.device.bus_address()
            );
            return Err(UsbError::InvalidDescriptorField);
        }

        log::info!(
            "{}: setting up {:?}-speed hub, {} downstream ports",
            device.device.bus_address(),
            device.device.speed(),
            hub_descriptor.b_nr_ports
        );

        device.configure_hub(&hub_descriptor).await?;

        // TODO power on the ports?
        Ok(Self {
            device,
            hub_descriptor,
            children: 0,
        })
    }

    async fn set_port_feature(
        &mut self,
        port: NonZeroU8,
        feature: PortFeatureSelector,
    ) -> Result<(), UsbError> {
        self.device
            .control_pipe()
            .control_transfer(ControlTransferSetup {
                bm_request_type: BM_REQUEST_TYPE_PORT_FEATURE,
                b_request: B_REQUEST_SET_FEATURE,
                w_value: feature as _,
                w_index: port.get() as _,
                w_length: 0,
            })
            .await
    }
    async fn clear_port_feature(
        &mut self,
        port: NonZeroU8,
        feature: PortFeatureSelector,
    ) -> Result<(), UsbError> {
        self.device
            .control_pipe()
            .control_transfer(ControlTransferSetup {
                bm_request_type: BM_REQUEST_TYPE_PORT_FEATURE,
                b_request: B_REQUEST_CLEAR_FEATURE,
                w_value: feature as _,
                w_index: port.get() as _,
                w_length: 0,
            })
            .await
    }
    async fn clear_port_changes(
        &mut self,
        port: NonZeroU8,
        change: UsbHubPortChange,
    ) -> Result<(), UsbError> {
        if change.contains(UsbHubPortChange::C_PORT_RESET) {
            self.clear_port_feature(port, PortFeatureSelector::CPortReset)
                .await?;
        }
        if change.contains(UsbHubPortChange::C_PORT_CONNECTION) {
            self.clear_port_feature(port, PortFeatureSelector::CPortConnection)
                .await?;
        }
        if change.contains(UsbHubPortChange::C_PORT_ENABLE) {
            self.clear_port_feature(port, PortFeatureSelector::CPortEnable)
                .await?;
        }
        Ok(())
    }

    async fn get_port_status(
        &mut self,
        port: NonZeroU8,
    ) -> Result<(UsbHubPortStatus, UsbHubPortChange), UsbError> {
        let mut buffer = [MaybeUninit::uninit(); 2];
        let len = self
            .device
            .control_pipe()
            .control_transfer_in(
                ControlTransferSetup {
                    bm_request_type: BM_REQUEST_TYPE_GET_PORT_STATUS,
                    b_request: B_REQUEST_GET_STATUS,
                    w_value: 0,
                    w_index: port.get() as _,
                    w_length: 4,
                },
                buffer.as_bytes_mut(),
            )
            .await?;
        if len != 4 {
            return Err(UsbError::TruncatedDescriptor(4, len));
        }
        let [w0, w1] = unsafe { MaybeUninit::array_assume_init(buffer) };
        let status = UsbHubPortStatus(w0);
        let change = UsbHubPortChange(w1);
        Ok((status, change))
    }

    async fn reset_port(&mut self, port: NonZeroU8) -> Result<UsbHubPortStatus, UsbError> {
        self.set_port_feature(port, PortFeatureSelector::PortReset)
            .await?;

        // Wait for reset to be asserted
        let deadline = monotonic_time() + Duration::from_secs(3);
        let mut status;
        let mut change;

        loop {
            (status, change) = self.get_port_status(port).await?;

            self.clear_port_changes(port, change).await?;

            if status.contains(UsbHubPortStatus::PORT_RESET)
                || status.contains(UsbHubPortStatus::PORT_ENABLE)
                || !status.contains(UsbHubPortStatus::PORT_CONNECTION)
            {
                // Port reset got asserted or port got enabled
                break;
            }

            if monotonic_time() >= deadline {
                log::warn!("Port reset did not assert in 3 sec");
                return Ok(status);
            }
        }

        if status.contains(UsbHubPortStatus::PORT_ENABLE)
            || !status.contains(UsbHubPortStatus::PORT_CONNECTION)
        {
            return Ok(status);
        }

        todo!()
    }

    async fn setup_connected_port(&mut self, port: NonZeroU8) -> Result<(), UsbError> {
        let status = self.reset_port(port).await?;
        if !status.contains(UsbHubPortStatus::PORT_CONNECTION) {
            log::warn!(
                "{}: port {} disconnected during reset",
                self.device.bus_address(),
                port
            );
            return Ok(());
        }
        if !status.contains(UsbHubPortStatus::PORT_ENABLE) {
            log::warn!("{}: port {} did not reset", self.device.bus_address(), port);
            return Ok(());
        }

        let usb_speed = match (
            status.contains(UsbHubPortStatus::PORT_LOW_SPEED),
            status.contains(UsbHubPortStatus::PORT_HIGH_SPEED),
        ) {
            (false, false) => UsbSpeed::Full,
            (true, false) => UsbSpeed::Low,
            (false, true) => UsbSpeed::High,
            (true, true) => todo!(),
        };
        log::info!(
            "{}: hub port {}: {:?}-speed device connected",
            self.device.bus_address(),
            port,
            usb_speed
        );
        let port_string = self.device.port_string().append(port);
        if let Err(error) = self
            .device
            .host_controller()
            .setup_hub_device(self.device.clone(), port_string, usb_speed)
            .await
        {
            log::error!(
                "{}: hub port {} ({}) setup failed: {:?}",
                self.device.bus_address(),
                port,
                port_string,
                error
            );
            return Ok(());
        }
        self.children |= 1 << (port.get() - 1);

        Ok(())
    }

    async fn disconnect_port(&mut self, port: NonZeroU8) -> Result<(), UsbError> {
        let port_string = self.device.port_string().append(port);
        self.children &= !(1 << (port.get() - 1));
        self.device
            .host_controller()
            .disconnect_device(port_string)
            .await
    }

    async fn poll_port(&mut self, port: NonZeroU8) -> Result<(), UsbError> {
        let (status, change) = self.get_port_status(port).await?;
        if change.contains(UsbHubPortChange::C_PORT_CONNECTION) {
            // Clear feature: C_PORT_CONNECTION
            self.clear_port_feature(port, PortFeatureSelector::CPortConnection)
                .await?;

            if status.contains(UsbHubPortStatus::PORT_CONNECTION) {
                self.setup_connected_port(port).await?;
            } else if let Err(error) = self.disconnect_port(port).await {
                log::warn!(
                    "{}: hub port {} did not disconnect cleanly: {:?}",
                    self.device.bus_address(),
                    port,
                    error
                );
            }
        }

        Ok(())
    }

    async fn poll(&mut self) -> bool {
        for i in 0..self.hub_descriptor.b_nr_ports {
            let port = NonZeroU8::new(i + 1).unwrap();
            if let Err(error) = self.poll_port(port).await {
                log::error!(
                    "{}: hub port {} poll error: {:?}",
                    self.device.bus_address(),
                    port,
                    error
                );
                return false;
            }
        }
        true
    }

    async fn cleanup(&mut self) {
        for i in 0..self.hub_descriptor.b_nr_ports {
            let port = NonZeroU8::new(i + 1).unwrap();
            if self.children & (1 << i) != 0 {
                if let Err(error) = self.disconnect_port(port).await {
                    log::warn!(
                        "{}: downstream {} disconnect error: {:?}",
                        self.device.bus_address(),
                        port,
                        error
                    );
                }
            }
        }
    }

    async fn run(mut self) -> Result<(), UsbError> {
        loop {
            if !self.poll().await {
                self.cleanup().await;
                return Ok(());
            }
            runtime::sleep(Duration::from_millis(100)).await;
        }
    }
}

#[async_trait]
impl UsbDeviceDriver for UsbHubDriver {
    async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
        let hub = UsbHub::setup(device).await?;
        hub.run().await
    }

    fn name(&self) -> &'static str {
        "USB Hub"
    }

    fn probe(&self, class: u8, subclass: u8, protocol: u8, vid: u16, pid: u16) -> bool {
        let _ = (protocol, vid, pid);
        class == 0x09 && subclass == 0x00
    }
}
