use core::{fmt, sync::atomic::AtomicU32};

use alloc::vec::Vec;
use capability::{CapabilityRegs, HCCPARAMS1, HCSPARAMS1, HCSPARAMS2};
use extended::ExtendedCapability;
use libk::error::Error;
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use operational::{OperationalRegs, PortRegs, PORTSC, USBCMD, USBSTS};
use runtime::{RuntimeRegs, IMAN};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    LocalRegisterCopy,
};

pub mod capability;
pub mod extended;
pub mod operational;
pub mod runtime;

pub struct Regs {
    #[allow(unused)]
    pub(crate) capability: DeviceMemoryIo<'static, CapabilityRegs>,
    pub(crate) operational: IrqSafeRwLock<DeviceMemoryIo<'static, OperationalRegs>>,
    pub(crate) runtime: IrqSafeRwLock<DeviceMemoryIo<'static, RuntimeRegs>>,
    pub(crate) doorbells: DeviceMemoryIo<'static, [AtomicU32]>,
    pub(crate) extended_capabilities: Vec<ExtendedCapability>,
    // TODO per-port locks
    pub(crate) ports: IrqSafeRwLock<DeviceMemoryIo<'static, [PortRegs]>>,
    pub(crate) port_count: usize,
    pub(crate) slot_count: usize,
    pub(crate) scratch_count: usize,
    pub(crate) context_size: ContextSize,
}

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

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[repr(transparent)]
pub struct PortNumber(u8);

impl Regs {
    const PORT_REGS_OFFSET: usize = 0x400;

    pub fn map(base: PhysicalAddress) -> Result<Self, Error> {
        let capability =
            unsafe { DeviceMemoryIo::<CapabilityRegs>::map(base, Default::default()) }?;

        let operational_offset = capability.CAPLENGTH.get() as usize;
        let doorbell_offset = capability.DBOFF.get() & !0x3;
        let runtime_offset = capability.RTSOFF.get() & !0x1F;
        let extended_capability_offset =
            (capability.HCCPARAMS1.read(HCCPARAMS1::XECP) as usize) << 2;

        let port_count = capability.HCSPARAMS1.read(HCSPARAMS1::MaxPorts) as usize;
        let slot_count = capability.HCSPARAMS1.read(HCSPARAMS1::MaxSlots) as usize;

        let operational_base = base.add(operational_offset);
        let doorbell_base = base.add(doorbell_offset as usize);
        let runtime_base = base.add(runtime_offset as usize);

        let scratch_lo = capability.HCSPARAMS2.read(HCSPARAMS2::MaxScratchpadBufsLo);
        let scratch_hi = capability.HCSPARAMS2.read(HCSPARAMS2::MaxScratchpadBufsHi);
        let scratch_count = ((scratch_hi << 5) | scratch_lo) as usize;

        let context_size = if capability.HCCPARAMS1.matches_all(HCCPARAMS1::CSZ::SET) {
            ContextSize::Context64
        } else {
            ContextSize::Context32
        };

        let operational = unsafe {
            DeviceMemoryIo::<OperationalRegs>::map(operational_base, Default::default())
        }?;
        let runtime =
            unsafe { DeviceMemoryIo::<RuntimeRegs>::map(runtime_base, Default::default()) }?;

        let ports = unsafe {
            DeviceMemoryIo::map_slice(
                operational_base.add(Self::PORT_REGS_OFFSET),
                port_count,
                Default::default(),
            )
        }?;
        let doorbells =
            unsafe { DeviceMemoryIo::map_slice(doorbell_base, 256, Default::default()) }?;
        let extended_capabilities =
            ExtendedCapability::parse_list(base, extended_capability_offset)?;

        Ok(Self {
            capability,
            operational: IrqSafeRwLock::new(operational),
            runtime: IrqSafeRwLock::new(runtime),
            doorbells,
            ports: IrqSafeRwLock::new(ports),
            extended_capabilities,

            port_count,
            slot_count,
            scratch_count,
            context_size,
        })
    }

    pub fn hc_reset(&self, timeout_cycles: u64) -> Result<(), Error> {
        let op = self.operational.write();

        // Wait for CNR to get cleared
        op.wait_usbsts_bit(USBSTS::CNR::CLEAR, timeout_cycles)?;

        // Clear run/stop and wait for the HC to halt
        op.USBCMD.modify(USBCMD::RS::CLEAR);
        op.wait_usbsts_bit(USBSTS::HCH::SET, timeout_cycles)?;

        // Reset the HC, wait for CNR to get cleared again
        op.USBCMD.modify(USBCMD::HCRST::SET);
        op.wait_usbsts_bit(USBSTS::CNR::CLEAR, timeout_cycles)?;

        // Halt the HC again
        op.USBCMD.modify(USBCMD::RS::CLEAR);
        op.wait_usbsts_bit(USBSTS::HCH::SET, timeout_cycles)?;

        Ok(())
    }

    pub fn handle_interrupt(&self, vector: usize) -> LocalRegisterCopy<u32, USBSTS::Register> {
        let op = self.operational.write();
        let rt = self.runtime.write();

        let status = op.USBSTS.extract();

        // Clear the RW1C bits
        op.USBSTS.set(status.get());

        // Acknowledge interrupts for given vector (vector == interrupter index)
        let interrupter = &rt.IRn[vector];

        interrupter.IMAN.modify(IMAN::IP::SET);

        status
    }
}

pub fn portsc_to_neutral(
    portsc: LocalRegisterCopy<u32, PORTSC::Register>,
) -> LocalRegisterCopy<u32, PORTSC::Register> {
    const RO_MASK: u32 = (1 << 0) | (1 << 3) | (0xF << 10) | (1 << 30);
    const RWS_MASK: u32 = (0xF << 5)
        | (1 << 9)
        | (0x3 << 14)
        | (1 << 25)
        | (1 << 26)
        | (1 << 27)
        | (1 << 30)
        | (1 << 31);

    portsc.bitand(RO_MASK | RWS_MASK)
}

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

impl PortNumber {
    pub fn from_index(index: usize) -> Self {
        Self(index as u8 + 1)
    }

    pub fn index(&self) -> usize {
        self.0 as usize - 1
    }
}

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