use core::{
    fmt,
    mem::MaybeUninit,
    num::NonZeroU8,
    sync::atomic::{AtomicUsize, Ordering},
};

use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, vec, vec::Vec};
use device_api::{device::Device, interrupt::InterruptHandler};
use libk::task::runtime;
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    PageBox,
};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use xhci_lib::ExtendedCapability;
use ygg_driver_usb::{
    bus::UsbBusManager,
    device::{UsbBusAddress, UsbDeviceAccess},
    error::UsbError,
    info::UsbVersion,
    pipe::control::UsbControlPipeAccess,
    util::UsbAddressAllocator,
    UsbHostController,
};
use yggdrasil_abi::error::Error;

use crate::{
    context::{XhciDeviceContext, XhciInputContext},
    device::XhciBusDevice,
    pipe::ControlPipe,
    regs::{Mapper, PortSpeed, Regs},
    ring::{
        CommandExecutor, CommandRing, ControlTransferRing, Event, EventRing, EventRingSegmentTable,
        GenericTransferRing,
    },
    util::EventBitmap,
};

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ContextSize {
    Context32,
    Context64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct PortNumber(NonZeroU8);

impl PortNumber {
    pub fn new(port: u8) -> Option<Self> {
        NonZeroU8::new(port).map(Self)
    }

    pub fn from_port_index(index: usize) -> Option<Self> {
        let as_u8 = (index + 1).try_into().ok()?;
        Some(Self(unsafe { NonZeroU8::new_unchecked(as_u8) }))
    }

    pub fn port_index(&self) -> usize {
        u8::from(self.0) as usize - 1
    }
}

#[allow(unused)]
struct Scratchpads {
    buffers: Vec<PageBox<[MaybeUninit<u8>]>>,
    array: PageBox<[PhysicalAddress]>,
}

#[allow(unused)]
struct RootHubPort {
    version: UsbVersion,
    slot_type: u8,
    max_hub_depth: Option<usize>,
}

pub struct Xhci {
    regs: Regs,

    bus_address: OneTimeInit<u16>,
    address_allocator: UsbAddressAllocator,

    port_count: usize,
    pub(crate) context_size: ContextSize,

    dcbaa: IrqSafeRwLock<PageBox<[PhysicalAddress]>>,
    endpoints: IrqSafeRwLock<BTreeMap<(u8, u8), Arc<dyn GenericTransferRing>>>,
    event_ring: EventRing,
    pub(crate) command_ring: CommandRing,
    #[allow(unused)]
    scratchpads: Option<Scratchpads>,

    root_hub_map: BTreeMap<PortNumber, RootHubPort>,

    port_event_map: EventBitmap,
}

impl Scratchpads {
    pub fn new(count: usize) -> Result<Option<Self>, UsbError> {
        if count == 0 {
            return Ok(None);
        }

        let buffers = (0..count)
            .map(|_| PageBox::new_uninit_slice(4096))
            .collect::<Result<Vec<_>, _>>()
            .map_err(UsbError::MemoryError)?;
        let array = PageBox::new_slice_with(|i| unsafe { buffers[i].as_physical_address() }, count)
            .map_err(UsbError::MemoryError)?;

        Ok(Some(Self { buffers, array }))
    }
}

impl Xhci {
    // Extract all info about the xHC, but don't do any init besides performing a BIOS->OS handoff
    pub fn new(
        regs: xhci_lib::Registers<Mapper>,
        mut extended: xhci_lib::extended_capabilities::List<Mapper>,
    ) -> Result<Self, UsbError> {
        let event_ring = EventRing::new(128)?;
        let command_ring = CommandRing::new(128)?;

        let regs = Regs::from(regs);

        let scratchpad_count = regs.scratchpad_count();
        let port_count = regs.port_count();
        let slot_count = regs.max_slot_count();
        let context_size = regs.context_size();
        let context_size = match context_size {
            32 => ContextSize::Context32,
            64 => ContextSize::Context64,
            _ => {
                log::error!("Unhandled context size: {context_size}");
                return Err(UsbError::InvalidConfiguration);
            }
        };

        let scratchpads = Scratchpads::new(scratchpad_count)?;
        let mut dcbaa = PageBox::new_slice(PhysicalAddress::ZERO, slot_count + 1)
            .map_err(UsbError::MemoryError)?;

        if let Some(scratchpads) = scratchpads.as_ref() {
            dcbaa[0] = unsafe { scratchpads.array.as_physical_address() };
        }

        let mut root_hub_map = BTreeMap::new();
        for cap in extended.into_iter() {
            let Ok(cap) = cap else {
                continue;
            };

            match cap {
                ExtendedCapability::UsbLegacySupport(mut legsup) => {
                    let mut handoff = false;
                    legsup.usblegsup.update_volatile(|legsup| {
                        if legsup.hc_bios_owned_semaphore() {
                            handoff = true;
                            legsup.set_hc_os_owned_semaphore();
                        }
                    });

                    if !handoff {
                        continue;
                    }
                    log::info!("xhci: BIOS->OS handoff started");

                    let mut timeout = 10000000;
                    while timeout > 0 {
                        let status = legsup.usblegsup.read_volatile();
                        if !status.hc_bios_owned_semaphore() && status.hc_os_owned_semaphore() {
                            break;
                        }

                        core::hint::spin_loop();
                        timeout -= 1;
                    }
                    if timeout == 0 {
                        log::error!("xhci: BIOS->OS handoff failed");
                        return Err(UsbError::DeviceBusy);
                    }
                    log::info!("xhci: BIOS->OS handoff finished");
                }
                ExtendedCapability::XhciSupportedProtocol(proto) => {
                    let header = proto.header.read_volatile();

                    if header.name_string() != 0x20425355 {
                        log::warn!("Skip unknown xHCI supported protocol capability");
                        continue;
                    }
                    let version_major = header.major_revision();
                    let version_minor = header.minor_revision();
                    let slot_type = header.protocol_slot_type();
                    let max_hub_depth = header.hub_depth();
                    let version = ((version_major as u16) << 8) | (version_minor as u16);
                    let Some(version) = UsbVersion::from_bcd_usb(version) else {
                        log::warn!("Skip unknown xHCI supported protocol revision: {version_major:x}.{version_minor:x}");
                        continue;
                    };

                    let port_range = header.compatible_port_offset()
                        ..header.compatible_port_offset() + header.compatible_port_count();
                    log::info!("Ports {port_range:?}: USB {version_major:x}.{version_minor:x}, slot type {slot_type}, max hub depth {max_hub_depth}");

                    for port in port_range {
                        let Some(number) = PortNumber::new(port) else {
                            continue;
                        };
                        root_hub_map.insert(
                            number,
                            RootHubPort {
                                version,
                                slot_type,
                                max_hub_depth: (max_hub_depth != 0)
                                    .then_some(max_hub_depth as usize),
                            },
                        );
                    }
                }
                _ => (),
            }
        }

        Ok(Self {
            regs,

            bus_address: OneTimeInit::new(),
            address_allocator: UsbAddressAllocator::new(),

            port_count,
            context_size,

            event_ring,
            command_ring,

            dcbaa: IrqSafeRwLock::new(dcbaa),
            endpoints: IrqSafeRwLock::new(BTreeMap::new()),
            scratchpads,
            root_hub_map,

            port_event_map: EventBitmap::new(),
        })
    }

    pub fn register_device_context(&self, slot_id: u8, context: PhysicalAddress) {
        self.dcbaa.write()[slot_id as usize] = context;
    }

    pub fn register_endpoint(
        &self,
        slot_id: u8,
        endpoint_id: u8,
        ring: Arc<dyn GenericTransferRing>,
    ) {
        self.endpoints.write().insert((slot_id, endpoint_id), ring);
    }

    pub fn shutdown_endpoint(&self, slot_id: u8, endpoint_id: u8) {
        if let Some(endpoint) = self.endpoints.write().remove(&(slot_id, endpoint_id)) {
            endpoint.shutdown();
        } else {
            log::warn!(
                "Endpoint {}:{} does not exist, maybe already shut down?",
                slot_id,
                endpoint_id
            );
        }
    }

    pub fn notify_transfer(
        &self,
        slot_id: u8,
        endpoint_id: u8,
        address: PhysicalAddress,
        status: u32,
    ) {
        if let Some(ep) = self.endpoints.read().get(&(slot_id, endpoint_id)) {
            ep.notify(address, status);
        } else {
            log::warn!("No endpoint slot={slot_id}, ep={endpoint_id}");
        }
    }

    async fn reset_port(&self, port: PortNumber) -> Result<(), UsbError> {
        let index = port.port_index();

        // Set port reset and wait for it to clear
        self.regs.ports.update(index, |regs| {
            regs.portsc.set_port_reset();
        });

        // TODO timeout
        loop {
            self.port_event_map.wait_specific(index).await;

            let mut status = None;
            self.regs.ports.update(index, |regs| {
                // Port became disconnected
                if !regs.portsc.port_enabled_disabled() || !regs.portsc.current_connect_status() {
                    log::warn!("xhci: port {port} disconnected during reset");
                    status = Some(Err(UsbError::DeviceDisconnected));
                    return;
                }

                if !regs.portsc.port_reset() {
                    regs.portsc.clear_port_reset_change();
                    status = Some(Ok(()));
                }
            });

            if let Some(status) = status {
                return status;
            }
        }
    }

    async fn setup_connected_port(self: Arc<Self>, port: PortNumber) -> Result<(), UsbError> {
        // TODO cleanup after a failed device init:
        //  * Deallocate the bus address
        //  * Issue a Disable Slot TRB
        //  * Remove Device Context from DCBAA
        //  * Deregister the device's Default Control Endpoint

        let index = port.port_index();
        let root_hub_port = self
            .root_hub_map
            .get(&port)
            .ok_or(UsbError::PortInitFailed)?;

        log::info!(
            "xhci: setup {} device at port {port}",
            root_hub_port.version
        );

        if root_hub_port.version == UsbVersion::Usb20 {
            // Port needs a reset first
            log::info!("xhci: reset port {port}");
            self.reset_port(port).await?;
        }

        let speed = PortSpeed::try_from(self.regs.ports.read(index).portsc.port_speed())
            .map_err(|_| UsbError::PortInitFailed)?;

        log::info!("xhci: port {port} effective speed {speed:?}");

        // Allocate a device slot
        let slot_id = self
            .command_ring
            .enable_slot(&*self, root_hub_port.slot_type)
            .await?;

        // Allocate some address for the device
        let bus_address = self.address_allocator.allocate()?;

        // Allocate a default endpoint ring
        let control_ring = Arc::new(ControlTransferRing::new(slot_id, 1, 128)?);

        // Register control endpoint
        self.register_endpoint(slot_id, 1, control_ring.clone());

        let control_pipe = UsbControlPipeAccess(Box::new(ControlPipe::new(
            self.clone(),
            slot_id,
            control_ring.clone(),
        )));

        // Setup Device Context
        let device_cx = XhciDeviceContext::new(self.context_size)?;
        self.register_device_context(slot_id, device_cx.physical_address());

        // Issue an Address Device TRB with BSR=1 first
        let input_cx = XhciInputContext::new_address_device(
            self.context_size,
            bus_address,
            port,
            speed,
            None,
            control_ring.dequeue_pointer(),
        )?;
        self.command_ring
            .address_device(&*self, slot_id, input_cx.physical_address(), true)
            .await?;

        let state = device_cx.map(|cx| cx.slot().slot_state());
        log::info!("xhci: port {port} slot {slot_id} state after BSR=1: {state:?}");

        // After an Address Device with BSR=1, retrieving a Device Descriptor is possible
        // This is needed for Full-speed devices, where the max_packet_size is 8 at first, but
        // is determined by the Device Descriptor when it's available
        let device_descriptor = control_pipe.query_device_descriptor().await?;
        let max_packet_size =
            device_descriptor.max_packet_size(root_hub_port.version, speed.into())?;

        log::info!("xhci: port {port}:");
        log::info!(" * max_packet_size = {max_packet_size}");
        let max_packet_size = Some(max_packet_size);

        // Reset the control endpoint for a proper dequeue pointer and a proper DCS
        control_ring.reset();

        // Issue an Address Device TRB with BSR=0
        let input_cx = XhciInputContext::new_address_device(
            self.context_size,
            bus_address,
            port,
            speed,
            max_packet_size,
            control_ring.dequeue_pointer(),
        )?;
        self.command_ring
            .address_device(&*self, slot_id, input_cx.physical_address(), false)
            .await
            .inspect_err(|error| {
                log::error!("xhci: port {port} Address Device (BSR=0) error: {error:?}")
            })?;

        let state = device_cx.map(|cx| cx.slot().slot_state());
        log::info!("xhci: port {port} slot {slot_id} state after BSR=0: {state:?}");

        log::info!("xhci: port {port} device addressed ({bus_address})");

        // Hand off the device to the general USB stack

        let bus_address = UsbBusAddress {
            bus: *self.bus_address.get(),
            device: bus_address,
        };

        let device = Box::new(XhciBusDevice {
            xhci: self.clone(),
            slot_id,
            port_id: port.into(),
            bus_address,
            speed,
            control_pipe,
            device_context: device_cx,
            rings: IrqSafeRwLock::new(vec![control_ring]),
            detach_handler: IrqSafeSpinlock::new(None),
        });
        let device = UsbDeviceAccess::setup(device).await?;

        UsbBusManager::register_device(Arc::new(device));

        Ok(())
    }

    async fn port_manager_task(self: Arc<Self>) -> Result<(), UsbError> {
        // Inject events for the root hub ports
        for (&port, _) in self.root_hub_map.iter() {
            self.port_event_map.signal(port.port_index());
        }

        let mut show_disconnect_mask = 0;
        loop {
            let events = self.port_event_map.wait_any(self.port_count).await;

            for port_index in events {
                let port = PortNumber::from_port_index(port_index).unwrap();

                let mut connected = None;

                self.regs.ports.update(port_index, |regs| {
                    if regs.portsc.port_enabled_disabled_change() {
                        regs.portsc.clear_port_enabled_disabled_change();
                    }
                    if regs.portsc.connect_status_change() {
                        regs.portsc.clear_connect_status_change();
                        connected = Some(regs.portsc.current_connect_status());
                    }
                });

                let Some(state) = connected else {
                    continue;
                };

                if state {
                    if show_disconnect_mask & (1 << port_index) == 0 {
                        log::info!("xhci: port {port} connected");
                        if let Err(error) = self.clone().setup_connected_port(port).await {
                            show_disconnect_mask &= !(1 << port_index);
                            log::error!("xhci: port {port} setup failed: {error:?}");
                        } else {
                            show_disconnect_mask |= 1 << port_index;
                        }
                    }
                } else {
                    if show_disconnect_mask & (1 << port_index) != 0 {
                        log::warn!("xhci: port {port} disconnected");
                        show_disconnect_mask &= !(1 << port_index);
                    }
                }
            }
        }
    }

    fn handle_event(self: Arc<Self>) {
        while let Some(event) = self.event_ring.try_dequeue() {
            match event {
                Event::PortChange(port) => {
                    if port > 0 {
                        self.port_event_map.signal(port - 1);
                    }
                }
                Event::CommandCompletion { address, reply } => {
                    self.command_ring.notify(address, reply);
                }
                Event::Transfer {
                    address,
                    slot_id,
                    endpoint_id,
                    status,
                } => {
                    self.notify_transfer(slot_id, endpoint_id, address, status);
                }
            }
        }

        self.regs
            .set_interrupter_0_dequeue_pointer(self.event_ring.dequeue_pointer());
    }
}

impl UsbHostController for Xhci {}

impl CommandExecutor for Xhci {
    fn ring_doorbell(&self, index: usize, target: u8) {
        self.regs.ring_doorbell(index, target);
    }
}

impl Device for Xhci {
    unsafe fn init(self: Arc<Self>) -> Result<(), Error> {
        static XHCI_COUNT: AtomicUsize = AtomicUsize::new(0);
        log::info!("Init USB xHCI");

        if XHCI_COUNT.fetch_add(1, Ordering::Release) != 0 {
            log::warn!("Skip second xhci init");
            return Ok(());
        }

        self.regs.reset();
        self.regs.set_max_slot_count();

        let erst = EventRingSegmentTable::for_event_rings(&[&self.event_ring])?;
        let dcbaa = self.dcbaa.read();

        self.regs
            .configure(&dcbaa, &self.command_ring, &self.event_ring, &erst);

        let bus = UsbBusManager::register_bus(self.clone());
        self.bus_address.init(bus);

        // Start the port manager task
        runtime::spawn(self.clone().port_manager_task())?;

        Ok(())
    }

    fn display_name(&self) -> &str {
        "USB xHCI"
    }
}

impl InterruptHandler for Xhci {
    fn handle_irq(self: Arc<Self>, _vector: Option<usize>) -> bool {
        if let Some(status) = self.regs.handle_interrupt() {
            if status.event_interrupt() {
                self.handle_event();
            }

            true
        } else {
            false
        }
    }
}

impl fmt::Display for PortNumber {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

impl From<PortNumber> for u8 {
    fn from(value: PortNumber) -> Self {
        value.0.into()
    }
}