use core::{
    ops::Range,
    sync::atomic::{AtomicU32, Ordering},
};

use alloc::vec::Vec;
use libk::error::Error;
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use ygg_driver_usb::{error::UsbError, info::UsbVersion};

pub struct ProtocolSupport {
    words: [u32; 4],
}

pub struct LegacySupport(DeviceMemoryIo<'static, [AtomicU32; 4]>);

pub enum ExtendedCapability {
    LegacySupport(IrqSafeRwLock<LegacySupport>),
    ProtocolSupport(ProtocolSupport),
}

impl ExtendedCapability {
    pub fn parse_list(
        base: PhysicalAddress,
        extended_capability_offset: usize,
    ) -> Result<Vec<Self>, Error> {
        let mut extended_capabilities = Vec::new();

        if extended_capability_offset != 0 {
            let extended_capability_base = base.add(extended_capability_offset);
            let mut current = extended_capability_base;

            loop {
                let header =
                    unsafe { DeviceMemoryIo::<[AtomicU32; 4]>::map(current, Default::default()) }?;
                let word0 = header[0].load(Ordering::Acquire);
                let id = word0 as u8;
                let next_pointer = ((word0 >> 8) & 0xFF) as usize;

                match id {
                    0x01 => extended_capabilities.push(ExtendedCapability::LegacySupport(
                        IrqSafeRwLock::new(LegacySupport(header)),
                    )),
                    0x02 => {
                        let words = [
                            word0,
                            header[1].load(Ordering::Acquire),
                            header[2].load(Ordering::Acquire),
                            header[3].load(Ordering::Acquire),
                        ];
                        extended_capabilities.push(ExtendedCapability::ProtocolSupport(
                            ProtocolSupport { words },
                        ))
                    }
                    _ => (),
                }

                if next_pointer == 0 {
                    break;
                } else {
                    current = current.add(next_pointer << 2);
                }
            }
        }

        Ok(extended_capabilities)
    }
}

impl ProtocolSupport {
    pub fn usb_revision(&self) -> Option<UsbVersion> {
        UsbVersion::from_bcd_usb((self.words[0] >> 16) as u16)
    }

    pub fn slot_type(&self) -> u8 {
        (self.words[3] & 0x1F) as u8
    }

    pub fn port_range(&self) -> Range<u8> {
        let port_offset = self.words[2] as u8;
        let port_count = (self.words[2] >> 8) as u8;

        port_offset.max(1)..port_offset.saturating_add(port_count)
    }
}

impl LegacySupport {
    fn is_bios_owned(&self) -> bool {
        self.0[0].load(Ordering::Acquire) & (1 << 16) != 0
    }

    fn set_os_owned(&mut self) {
        self.0[0].fetch_or(1 << 24, Ordering::Release);
    }

    pub fn perform_bios_handoff(&mut self, mut timeout_cycles: u64) -> Result<(), UsbError> {
        if !self.is_bios_owned() {
            return Ok(());
        }

        // Clear SMI events
        self.0[1].fetch_and(
            !((1 << 0) | (1 << 4) | (1 << 13) | (1 << 14) | (1 << 15)),
            Ordering::Release,
        );

        // Set OS owned semaphore
        self.set_os_owned();
        while timeout_cycles > 0 && self.is_bios_owned() {
            core::hint::spin_loop();
            timeout_cycles -= 1;
        }

        if timeout_cycles > 0 {
            Ok(())
        } else {
            log::error!("xhci: BIOS handoff timed out");
            Err(UsbError::DeviceBusy)
        }
    }
}
