use alloc::{sync::Arc, vec::Vec};
use device_api::{device::Device, interrupt::InterruptHandler};
use libk::error::Error;
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    device::DeviceMemoryIo,
    PageBox,
};
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use tock_registers::interfaces::{ReadWriteable, Readable, Writeable};
use ygg_driver_pci::device::PciDeviceInfo;
use ygg_driver_usb::error::UsbError;

use crate::{
    regs::{ExtendedCapability, Regs, CONFIG, HCSPARAMS1, PORTLI, PORTSC, USBCMD, USBSTS},
    ring::{
        command::CommandRing,
        event::{Event, EventRing, EventRingSegmentTable},
        GenericRing,
    },
};

struct ScratchpadArray {
    buffers: Vec<PageBox<[u8]>>,
    array: PageBox<[PhysicalAddress]>,
}

pub struct Xhci {
    regs: Regs,
    pci: PciDeviceInfo,

    dcbaa: IrqSafeRwLock<PageBox<[PhysicalAddress]>>,
    scratchpads: Option<ScratchpadArray>,
    command_ring: CommandRing,
    event_ring: EventRing,
    erst: EventRingSegmentTable,
}

impl ScratchpadArray {
    pub fn new(capacity: usize, element_size: usize) -> Result<Self, UsbError> {
        let buffers = (0..capacity)
            .map(|_| PageBox::new_slice(0, element_size))
            .collect::<Result<Vec<_>, _>>()
            .map_err(UsbError::MemoryError)?;
        let array =
            PageBox::new_slice_with(|i| unsafe { buffers[i].as_physical_address() }, capacity)
                .map_err(UsbError::MemoryError)?;
        Ok(Self { buffers, array })
    }
}

impl Xhci {
    pub fn new(pci: PciDeviceInfo, regs: Regs) -> Result<Self, UsbError> {
        let mut dcbaa = PageBox::new_slice(PhysicalAddress::ZERO, regs.slot_count + 1)
            .map_err(UsbError::MemoryError)?;
        let command_ring = CommandRing::new(128)?;
        let event_ring = EventRing::new(128)?;
        let erst = EventRingSegmentTable::for_event_rings(&[&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(regs.scratch_count, 0x1000)?;
            dcbaa[0] = unsafe { array.array.as_physical_address() };
            Some(array)
        } else {
            None
        };

        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}");
                    }
                }
                ExtendedCapability::LegacySupport(legsup) => {
                    legsup.write().perform_bios_handoff(10000000)?;
                }
            }
        }

        Ok(Self {
            regs,
            pci,

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

    fn handle_event(&self) {
        while let Some(event) = self.event_ring.try_dequeue() {
            match event {
                Event::PortChange(index) => {
                    log::info!("Event Ring: port {index} changed");
                    let ports = self.regs.ports.write();
                    let port = &ports[index - 1];
                    let status = port.PORTSC.extract();
                    let portli = port.PORTLI.extract();

                    if status.matches_any(PORTSC::CSC::SET) {
                        let connected = status.matches_any(PORTSC::CCS::SET);

                        port.PORTSC.set(status.get());

                        let pp = status.read(PORTSC::PP);
                        let ccs = status.read(PORTSC::CCS);
                        let ped = status.read(PORTSC::PED);
                        let pls = status.read(PORTSC::PLS);
                        let ps = status.read(PORTSC::PS);

                        log::info!("port {index} PORTSC={:#x}", status.get());
                        log::info!("* xhci {}", self.pci.address);
                        log::info!("* PP={pp} CCS={ccs} PED={ped} PLS={pls} PS={ps}");
                    }
                }
                _ => {
                    log::info!("other event");
                }
            }
        }

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

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

        // Configure the HC
        let dcbaap = unsafe { self.dcbaa.read().as_physical_address() };
        let cr_base = self.command_ring.base();

        let op = self.regs.operational.write();
        let rt = self.regs.runtime.write();
        let ports = self.regs.ports.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);
        op.wait_usbsts_bit(USBSTS::CNR::CLEAR, 100000000)?;

        Ok(())
    }

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

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

        log::info!("IRQ started: {:#x}", status.get());

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

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

        log::info!("IRQ finished");

        true
    }
}