use core::{
    mem::MaybeUninit,
    sync::atomic::{AtomicU8, Ordering},
};

use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc, vec::Vec};
use async_trait::async_trait;
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptHandler, IrqVector},
};
use libk::{
    dma::{BusAddress, DmaBuffer},
    error::Error,
    task::runtime,
};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    LocalRegisterCopy,
};
use ygg_driver_pci::{device::PciDeviceInfo, PciConfigurationSpace};
use ygg_driver_usb::{
    bus::UsbBusManager,
    device::{UsbBusAddress, UsbDeviceAccess, UsbSpeed},
    error::UsbError,
    info::UsbVersion,
    pipe::control::UsbControlPipeAccess,
    UsbHostController,
};
use yggdrasil_abi::bitflags;

use crate::{
    context::{XhciDeviceContext, XhciInputContext},
    device::XhciBusDevice,
    pipe::ControlPipe,
    regs::{
        self,
        extended::ExtendedCapability,
        operational::{PortRegs, CONFIG, PORTSC, USBCMD, USBSTS},
        PortNumber, Regs,
    },
    ring::{
        command::CommandRing,
        event::{Event, EventRing, EventRingSegmentTable},
        transfer::TransferRing,
        CommandExecutor, GenericRing,
    },
    util::EventBitmap,
};

bitflags! {
    pub struct Quirks: u32 {
        const INTEL: bit 0;
        const INTEL_PANTHER_POINT: bit 1;
    }
}

#[allow(unused)]
struct ScratchpadArray {
    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,
    array: DmaBuffer<[BusAddress]>,
}

struct RootHubPort {
    version: UsbVersion,
    slot_type: u8,
}

pub struct Xhci {
    pub(crate) regs: Regs,
    #[allow(unused)]
    pci: PciDeviceInfo,
    pub(crate) dma: Arc<dyn DmaAllocator>,
    quirks: Quirks,

    dcbaa: IrqSafeRwLock<DmaBuffer<[BusAddress]>>,
    #[allow(unused)]
    scratchpads: Option<ScratchpadArray>,
    pub(crate) command_ring: CommandRing,
    event_ring: EventRing,
    erst: EventRingSegmentTable,

    root_hub_ports: Vec<Option<RootHubPort>>,
    pub(crate) endpoints: IrqSafeRwLock<BTreeMap<(u8, u8), Arc<TransferRing>>>,
    pub(crate) slots: Vec<IrqSafeRwLock<Option<Arc<XhciBusDevice>>>>,
    pub(crate) port_slot_map: Vec<AtomicU8>,
    bus_index: OneTimeInit<u16>,
    port_event_map: EventBitmap,
}

impl ScratchpadArray {
    pub fn new(
        dma: &dyn DmaAllocator,
        capacity: usize,
        element_size: usize,
    ) -> Result<Self, UsbError> {
        let buffers = (0..capacity)
            .map(|_| DmaBuffer::new_zeroed_slice(dma, element_size))
            .collect::<Result<Vec<_>, _>>()
            .map_err(UsbError::MemoryError)?;
        let array = DmaBuffer::new_slice_with(dma, |i| buffers[i].bus_address(), capacity)
            .map_err(UsbError::MemoryError)?;
        Ok(Self { buffers, array })
    }
}

impl Xhci {
    pub fn new(
        dma: Arc<dyn DmaAllocator>,
        pci: PciDeviceInfo,
        regs: Regs,
    ) -> Result<Self, UsbError> {
        let mut quirks = Quirks::empty();
        if pci.vendor_id == 0x8086 {
            quirks |= Quirks::INTEL;

            if pci.device_id == 0x1E31 {
                quirks |= Quirks::INTEL_PANTHER_POINT;
            }
        }

        let mut dcbaa = DmaBuffer::new_slice(&*dma, BusAddress::ZERO, regs.slot_count + 1)
            .map_err(UsbError::MemoryError)?;
        let command_ring = CommandRing::new(&*dma, 128)?;
        let event_ring = EventRing::new(&*dma, 64)?;
        let erst = EventRingSegmentTable::for_event_rings(&*dma, &[&event_ring])?;

        // Setup scratch buffers
        // TODO: Linux seems to just ignore the PAGESIZE, it's (1 << 0) everywhere
        let scratchpads = if regs.scratch_count != 0 {
            let array = ScratchpadArray::new(&*dma, regs.scratch_count, 0x1000)?;
            dcbaa[0] = array.array.bus_address();
            // dcbaa[0] = unsafe { array.array.as_physical_address() };
            Some(array)
        } else {
            None
        };

        let mut root_hub_ports: Vec<Option<_>> = (0..regs.port_count).map(|_| None).collect();

        for cap in regs.extended_capabilities.iter() {
            match cap {
                ExtendedCapability::ProtocolSupport(support) => {
                    let Some(version) = support.usb_revision() else {
                        continue;
                    };

                    for port in support.port_range() {
                        log::info!("* Port {port}: {version}");
                        root_hub_ports[port as usize - 1] = Some(RootHubPort {
                            version,
                            slot_type: support.slot_type(),
                        });
                    }
                }
                ExtendedCapability::LegacySupport(legsup) => {
                    legsup.write().perform_bios_handoff(10000000)?;
                }
            }
        }

        // TODO this is only needed if there's an EHCI controller as well
        if quirks.contains(Quirks::INTEL_PANTHER_POINT) {
            log::info!("quirk: intel pantherpoint xhci/ehci port switchover");
            const INTEL_USB3PSSEN: usize = 0xD8;
            const INTEL_XUSB2PR: usize = 0xD0;

            let space = &pci.config_space;

            space.write_u32(INTEL_USB3PSSEN, 0xFFFFFFFF);
            space.write_u32(INTEL_XUSB2PR, 0xFFFFFFFF);
        }

        let port_slot_map = (0..regs.port_count).map(|_| AtomicU8::new(0)).collect();
        let slots = (0..regs.slot_count)
            .map(|_| IrqSafeRwLock::new(None))
            .collect();

        Ok(Self {
            regs,
            pci,
            dma,
            quirks,

            dcbaa: IrqSafeRwLock::new(dcbaa),
            scratchpads,
            command_ring,
            event_ring,
            erst,

            root_hub_ports,
            bus_index: OneTimeInit::new(),
            endpoints: IrqSafeRwLock::new(BTreeMap::new()),
            slots,
            port_slot_map,
            port_event_map: EventBitmap::new(),
        })
    }

    fn notify_endpoint(&self, slot_id: u8, endpoint_id: u8, address: BusAddress, status: u32) {
        if let Some(endpoint) = self.endpoints.read().get(&(slot_id, endpoint_id)) {
            endpoint.notify(address, status);
        } else {
            log::warn!("Endpoint {slot_id}:{endpoint_id} does not exist");
        }
    }

    fn kill_endpoint(&self, slot_id: u8, endpoint_id: u8) {
        let endpoint = self.endpoints.write().remove(&(slot_id, endpoint_id));
        if let Some(endpoint) = endpoint {
            endpoint.shutdown();
        }
    }

    async fn allocate_device_slot(
        self: &Arc<Self>,
        port_id: PortNumber,
        slot_type: u8,
        speed: UsbSpeed,
    ) -> Result<(u8, Arc<XhciBusDevice>, Arc<TransferRing>), UsbError> {
        let device_context = XhciDeviceContext::new(&*self.dma, self.regs.context_size)?;
        let slot_id = self.command_ring.enable_slot(&**self, slot_type).await?;
        let (control_pipe, control_ring) = ControlPipe::new(self.clone(), slot_id, 1, 128)?;
        let control_pipe = UsbControlPipeAccess(Box::new(control_pipe));

        self.dcbaa.write()[slot_id as usize] = device_context.bus_address();
        self.endpoints
            .write()
            .insert((slot_id, 1), control_ring.clone());

        let slot = Arc::new(XhciBusDevice {
            port_id,
            slot_id,
            device_context,
            control_pipe,
            speed,
            detach_handler: IrqSafeSpinlock::new(None),
            endpoints: IrqSafeRwLock::new(BTreeMap::from_iter([(1, control_ring.clone())])),
            xhci: self.clone(),
            bus_address: OneTimeInit::new(),
        });

        *self.slots[slot_id as usize].write() = Some(slot.clone());
        self.port_slot_map[port_id.index()].store(slot_id, Ordering::Release);

        Ok((slot_id, slot, control_ring))
    }

    // TODO proper timeout
    async fn reset_port(&self, regs: &PortRegs) -> Result<(), UsbError> {
        // Reset the port
        regs.modify_portsc_preserving(PORTSC::PR::SET);

        let mut timeout = 10000000;
        while timeout > 0 {
            let status = regs.PORTSC.extract();

            if status.matches_all(PORTSC::PR::CLEAR) {
                break;
            }

            timeout -= 1;
        }

        // Clear PRC
        regs.modify_portsc_preserving(PORTSC::PRC::SET);
        Ok(())
    }

    // TODO clean up resources if init fails
    async fn setup_port(
        self: &Arc<Self>,
        regs: &PortRegs,
        number: PortNumber,
    ) -> Result<(), UsbError> {
        let root_hub_port = self.root_hub_ports[number.index()]
            .as_ref()
            .ok_or(UsbError::PortInitFailed)?;

        let need_reset = !root_hub_port.version.is_version_3();

        if need_reset {
            self.reset_port(regs).await?;
        }

        let status = regs.PORTSC.extract();

        let speed = status.read(PORTSC::PS) as u8;
        let pls = status.read(PORTSC::PLS);

        let (usb_speed, max_packet_size) = match speed {
            1 => (UsbSpeed::Full, 8),
            2 => (UsbSpeed::Low, 8),
            3 => (UsbSpeed::High, 64),
            4 => (UsbSpeed::Super, 512),
            _ => {
                log::error!("Port {number} invalid speed value {speed}");
                return Err(UsbError::DeviceDisconnected);
            }
        };

        // PLS != 0 after reset
        if pls != 0 {
            log::error!("Port {number} pls!=0 after reset");
            return Err(UsbError::DeviceDisconnected);
        }

        let (slot_id, slot, control_ring) = self
            .allocate_device_slot(number, root_hub_port.slot_type, usb_speed)
            .await?;

        let input_cx = XhciInputContext::new_address_device(
            &*self.dma,
            self.regs.context_size,
            number,
            max_packet_size,
            speed,
            control_ring.bus_address(),
        )?;

        self.command_ring
            .address_device(&**self, slot_id, &input_cx, false)
            .await
            .inspect_err(|error| {
                log::error!("Port {number} Address Device TRB (BSR=0) failed: {error:?}")
            })?;

        // TODO update max_packet_size for default control endpoint after fetching the device
        //      descriptor

        let bus_address = slot
            .device_context
            .map(|device| device.slot().usb_device_address());
        slot.bus_address.init(UsbBusAddress {
            bus: *self.bus_index.get(),
            device: bus_address,
        });

        let device = UsbDeviceAccess::setup(slot).await?;
        UsbBusManager::register_device(device.into());

        Ok(())
    }

    async fn handle_disconnect(&self, number: PortNumber) -> Result<(), UsbError> {
        let slot_id = self.port_slot_map[number.index()].swap(0, Ordering::Acquire);
        if slot_id == 0 {
            return Ok(());
        }
        log::info!("Port {number} disconnected");

        let slot = self.slots[slot_id as usize].write().take();
        if let Some(slot) = slot {
            UsbBusManager::detach_device(*slot.bus_address.get());

            // Clean up xHC resources
            for (&endpoint_id, endpoint) in slot.endpoints.read().iter() {
                endpoint.shutdown();
                self.kill_endpoint(slot_id, endpoint_id);
            }
        } else {
            log::warn!("No device in slot {slot_id}");
        }

        // TODO stop endpoints on the xHC side?
        self.command_ring.disable_slot(self, slot_id).await?;

        Ok(())
    }

    async fn port_handler_task(self: Arc<Self>) -> Result<(), Error> {
        // Inject notify for ports that are already connected
        {
            let ports = self.regs.ports.write();
            for (index, port) in self.root_hub_ports.iter().enumerate() {
                if port.is_none() {
                    continue;
                }

                if ports[index].PORTSC.matches_all(PORTSC::CCS::SET) {
                    self.port_event_map.signal(index);
                }
            }
        }

        loop {
            let events = self.port_event_map.wait_any(self.regs.port_count).await;

            for index in events {
                let number = PortNumber::from_index(index);
                let ports = self.regs.ports.write();
                let regs = &ports[index];

                let status = regs.PORTSC.extract();
                let neutral = regs::portsc_to_neutral(status);
                let mut clear = LocalRegisterCopy::<u32, PORTSC::Register>::new(0);

                // Clear affected change bits
                let connected = if status.matches_all(PORTSC::CSC::SET) {
                    clear.modify(PORTSC::CSC::SET);
                    Some(status.matches_all(PORTSC::CCS::SET))
                } else {
                    None
                };
                if status.matches_all(PORTSC::PEC::SET) {
                    clear.modify(PORTSC::PEC::SET);
                }
                if status.matches_all(PORTSC::PLC::SET) {
                    clear.modify(PORTSC::PLC::SET);
                }

                regs.PORTSC.set(neutral.get() | clear.get());

                if let Some(connected) = connected {
                    if connected {
                        log::info!("Port {number} connected");
                        if let Err(error) = self.setup_port(regs, number).await {
                            log::error!("Port {number} setup error: {error:?}");
                        }
                    }

                    if !connected {
                        if let Err(error) = self.handle_disconnect(number).await {
                            log::error!("Port {number} disconnect error: {error:?}");
                        }
                    }
                }
            }
        }
    }

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

        self.regs
            .runtime
            .write()
            .set_interrupter_dequeue_pointer(0, self.event_ring.dequeue_pointer());
    }
}

#[async_trait]
impl CommandExecutor for Xhci {
    fn ring_doorbell(&self, index: usize, target: u8) {
        self.regs.doorbells[index].store(target as u32, Ordering::Release);
    }

    async fn reset_endpoint(
        &self,
        slot_id: u8,
        endpoint_id: u8,
        dequeue_pointer: BusAddress,
        dequeue_cycle: bool,
    ) -> Result<(), UsbError> {
        log::warn!("xhci: reset stalled endpoint {slot_id}:{endpoint_id}");

        self.command_ring
            .reset_endpoint(&*self, slot_id, endpoint_id, false)
            .await?;
        self.command_ring
            .set_tr_dequeue_pointer(&*self, slot_id, endpoint_id, dequeue_pointer, dequeue_cycle)
            .await?;

        Ok(())
    }
}

impl Device for Xhci {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        self.regs.hc_reset(10000000)?;
        log::info!("xHC reset complete");

        // Configure the HC
        let dcbaap = self.dcbaa.read().bus_address();
        let cr_base = self.command_ring.bus_base();

        let op = self.regs.operational.write();
        let rt = self.regs.runtime.write();

        log::info!("xhci: configure HC");
        op.CONFIG
            .modify(CONFIG::MaxSlotsEn.val(self.regs.slot_count as u32));
        op.set_dcbaap(dcbaap);
        op.set_crcr(cr_base, true);

        log::info!("xhci: configure interrupter");
        rt.configure_interrupter(0, &self.event_ring, &self.erst);

        // Enable interrupts and start the HC
        log::info!("xhci: start HC");
        op.USBCMD.modify(USBCMD::INTE::SET + USBCMD::HSEE::SET);
        op.USBCMD.modify(USBCMD::RS::SET);

        if self.quirks.contains(Quirks::INTEL) {
            for _ in 0..100000000 {
                core::hint::spin_loop();
            }
        }

        op.wait_usbsts_bit(USBSTS::CNR::CLEAR, 100000000)?;

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

        runtime::spawn(self.clone().port_handler_task()).ok();

        Ok(())
    }

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

impl InterruptHandler for Xhci {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let status = self.regs.handle_interrupt(0);

        if status.matches_all(USBSTS::HSE::SET) {
            log::error!("xhci: Host System Error occurred");
        }

        if status.matches_all(USBSTS::EINT::SET) {
            self.handle_event();
        }

        true
    }
}

impl UsbHostController for Xhci {}
