use core::{
    fmt,
    future::poll_fn,
    mem::{size_of, MaybeUninit},
    task::Poll,
};

use alloc::collections::BTreeMap;
use bytemuck::{Pod, Zeroable};
use device_api::dma::DmaAllocator;
use libk::dma::{BusAddress, DmaBuffer};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    waker::QueueWaker,
};
use ygg_driver_usb::error::UsbError;
use yggdrasil_abi::define_bitfields;

use crate::context::XhciInputContext;

use super::{CommandExecutor, GenericRing, LinkTrb};

struct CommandRingInner {
    trbs: DmaBuffer<[MaybeUninit<RawCommandTrb>]>,
    enqueue_index: usize,
    #[allow(unused)]
    dequeue_index: usize,
    cycle_bit: bool,
}

pub struct CommandRing {
    inner: IrqSafeSpinlock<CommandRingInner>,
    // TODO maybe use Vec of "slots"?
    completions: IrqSafeRwLock<BTreeMap<BusAddress, CommandReply>>,
    completion_notify: QueueWaker,
}

impl GenericRing for CommandRing {
    fn bus_base(&self) -> BusAddress {
        self.inner.lock().trbs.bus_address()
    }
}

impl CommandRingInner {
    fn enqueue<C: CommandTrb>(&mut self, trb: C) -> BusAddress {
        let mut raw: RawCommandTrb = bytemuck::cast(trb);

        raw.flags.set_ty(C::TRB_TYPE as _);
        raw.flags.set_cycle(self.cycle_bit);

        self.trbs[self.enqueue_index].write(raw);

        let address = self
            .trbs
            .bus_address()
            .add(self.enqueue_index * size_of::<RawCommandTrb>());
        // let address = unsafe { self.trbs.as_physical_address() }
        //     .add(self.enqueue_index * size_of::<RawCommandTrb>());

        // Move to the next TRB slot
        self.enqueue_index += 1;
        if self.enqueue_index >= self.trbs.len() - 1 {
            self.enqueue_link();

            // Wrap around
            self.cycle_bit = !self.cycle_bit;
            self.enqueue_index = 0;
        }

        address
    }

    fn enqueue_link(&mut self) {
        let base = self.trbs.bus_address();

        let link = LinkTrb::new(base, self.cycle_bit);
        self.trbs[self.enqueue_index].write(bytemuck::cast(link));
    }

    fn advance(&mut self) {
        self.dequeue_index += 1;

        if self.dequeue_index >= self.trbs.len() - 1 {
            self.dequeue_index = 0;
        }
    }
}

impl CommandRing {
    pub fn new(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, UsbError> {
        let trbs = DmaBuffer::new_zeroed_slice(dma, capacity).map_err(UsbError::MemoryError)?;

        Ok(Self {
            inner: IrqSafeSpinlock::new(CommandRingInner {
                trbs,
                enqueue_index: 0,
                dequeue_index: 0,
                cycle_bit: true,
            }),
            completions: IrqSafeRwLock::new(BTreeMap::new()),
            completion_notify: QueueWaker::new(),
        })
    }

    pub fn enqueue<C: CommandTrb>(&self, trb: C) -> BusAddress {
        let mut inner = self.inner.lock();
        inner.enqueue(trb)
    }

    pub async fn address_device<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_id: u8,
        input_context: &XhciInputContext,
        bsr: bool,
    ) -> Result<(), UsbError> {
        self.submit_and_wait(
            executor,
            AddressDeviceCommandTrb::new(input_context.bus_address(), slot_id, bsr),
        )
        .await?;
        Ok(())
    }

    pub async fn configure_endpoint<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_id: u8,
        input_context: &XhciInputContext,
    ) -> Result<(), UsbError> {
        self.submit_and_wait(
            executor,
            ConfigureEndpointCommandTrb::new(input_context.bus_address(), slot_id),
        )
        .await?;
        Ok(())
    }

    pub async fn reset_endpoint<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_id: u8,
        endpoint_id: u8,
        preserve_transaction_state: bool,
    ) -> Result<(), UsbError> {
        self.submit_and_wait(
            executor,
            ResetEndpointCommandTrb::new(slot_id, endpoint_id, preserve_transaction_state),
        )
        .await?;
        Ok(())
    }

    pub async fn set_tr_dequeue_pointer<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_id: u8,
        endpoint_id: u8,
        dequeue_pointer: BusAddress,
        dequeue_cycle: bool,
    ) -> Result<(), UsbError> {
        self.submit_and_wait(
            executor,
            SetTrDequeuePointerCommandTrb::new(
                slot_id,
                endpoint_id,
                dequeue_pointer,
                dequeue_cycle,
            ),
        )
        .await?;
        Ok(())
    }

    pub async fn enable_slot<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_type: u8,
    ) -> Result<u8, UsbError> {
        self.submit_and_wait(executor, EnableSlotCommandTrb::new(slot_type))
            .await
    }

    #[allow(unused)]
    pub async fn disable_slot<E: CommandExecutor>(
        &self,
        executor: &E,
        slot_id: u8,
    ) -> Result<u8, UsbError> {
        self.submit_and_wait(executor, DisableSlotCommandTrb::new(slot_id))
            .await
    }

    pub async fn submit_and_wait<E: CommandExecutor, C: CommandTrb>(
        &self,
        executor: &E,
        trb: C,
    ) -> Result<u8, UsbError> {
        let token = self.enqueue(trb);
        executor.ring_doorbell(0, 0);
        poll_fn(|cx| {
            self.completion_notify.register(cx.waker());
            if let Some(status) = self.get_completion(token) {
                self.completion_notify.remove(cx.waker());
                if status.completion_code == 1 {
                    Poll::Ready(Ok(status.slot_id))
                } else {
                    Poll::Ready(Err(UsbError::HostControllerCommandFailed(
                        status.completion_code,
                        status.completion_parameter,
                    )))
                }
            } else {
                Poll::Pending
            }
        })
        .await
    }

    pub fn get_completion(&self, address: BusAddress) -> Option<CommandReply> {
        self.completions.write().remove(&address)
    }

    pub fn notify(&self, address: BusAddress, reply: CommandReply) {
        self.inner.lock().advance();
        self.completions.write().insert(address, reply);
        self.completion_notify.wake_all();
    }
}

// TRB implementations

define_bitfields! {
    pub EnableSlotCommandFlags : u32 {
        (16..20) => slot_type
    }
}

define_bitfields! {
    pub DisableSlotCommandFlags : u32 {
        (24..32) => slot_id
    }
}

define_bitfields! {
    pub AddressDeviceCommandFlags : u32 {
        9 => bsr + set_bsr,
        (24..32) => slot_id
    }
}

define_bitfields! {
    pub ConfigureEndpointCommandFlags : u32 {
        (24..32) => slot_id
    }
}

define_bitfields! {
    pub ResetEndpointCommandFlags: u32 {
        (24..32) => slot_id,
        (16..21) => endpoint_id,
        9 => tsp,
    }
}

define_bitfields! {
    pub SetTrDequeuePointerCommandFlags: u32 {
        (24..32) => slot_id,
        (16..21) => endpoint_id,
    }
}

define_bitfields! {
    pub RawCommandFlags : u32 {
        (10..16) => ty + set_ty,
        0 => cycle + set_cycle
    }
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct EnableSlotCommandTrb {
    _0: [u32; 3],
    pub flags: EnableSlotCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct DisableSlotCommandTrb {
    _0: [u32; 3],
    pub flags: DisableSlotCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct AddressDeviceCommandTrb {
    pub input_context_address: BusAddress,
    _0: u32,
    pub flags: AddressDeviceCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct ConfigureEndpointCommandTrb {
    pub input_context_address: BusAddress,
    _0: u32,
    pub flags: ConfigureEndpointCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct ResetEndpointCommandTrb {
    _0: [u32; 3],
    pub flags: ResetEndpointCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct SetTrDequeuePointerCommandTrb {
    pub value: u64,
    _0: u32,
    pub flags: SetTrDequeuePointerCommandFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct RawCommandTrb {
    _0: [u32; 3],
    flags: RawCommandFlags,
}

#[derive(Debug, Clone, Copy)]
pub struct CommandReply {
    pub completion_code: u8,
    pub completion_parameter: u32,
    pub slot_id: u8,
}

pub trait CommandTrb: Pod + fmt::Debug {
    const TRB_TYPE: u8;
}

impl EnableSlotCommandTrb {
    pub fn new(slot_type: u8) -> Self {
        Self {
            _0: [0; 3],
            flags: EnableSlotCommandFlags::new(slot_type as u32),
        }
    }
}

impl DisableSlotCommandTrb {
    pub fn new(slot_id: u8) -> Self {
        Self {
            _0: [0; 3],
            flags: DisableSlotCommandFlags::new(slot_id as u32),
        }
    }
}

impl AddressDeviceCommandTrb {
    pub fn new(input_context_address: BusAddress, slot_id: u8, bsr: bool) -> Self {
        Self {
            input_context_address,
            _0: 0,
            flags: AddressDeviceCommandFlags::new(bsr, slot_id as _),
        }
    }
}

impl ConfigureEndpointCommandTrb {
    pub fn new(input_context_address: BusAddress, slot_id: u8) -> Self {
        Self {
            input_context_address,
            _0: 0,
            flags: ConfigureEndpointCommandFlags::new(slot_id as _),
        }
    }
}

impl ResetEndpointCommandTrb {
    pub fn new(slot_id: u8, endpoint_id: u8, preserve_transaction_state: bool) -> Self {
        Self {
            _0: [0; 3],
            flags: ResetEndpointCommandFlags::new(
                slot_id as _,
                endpoint_id as _,
                preserve_transaction_state,
            ),
        }
    }
}

impl SetTrDequeuePointerCommandTrb {
    pub fn new(
        slot_id: u8,
        endpoint_id: u8,
        dequeue_pointer: BusAddress,
        dequeue_cycle: bool,
    ) -> Self {
        let mut value = dequeue_pointer.into_u64();
        if dequeue_cycle {
            value |= 1;
        }
        Self {
            value,
            _0: 0,
            flags: SetTrDequeuePointerCommandFlags::new(slot_id as _, endpoint_id as _),
        }
    }
}

impl CommandTrb for EnableSlotCommandTrb {
    const TRB_TYPE: u8 = 9;
}

impl CommandTrb for DisableSlotCommandTrb {
    const TRB_TYPE: u8 = 10;
}

impl CommandTrb for AddressDeviceCommandTrb {
    const TRB_TYPE: u8 = 11;
}

impl CommandTrb for ConfigureEndpointCommandTrb {
    const TRB_TYPE: u8 = 12;
}

impl CommandTrb for ResetEndpointCommandTrb {
    const TRB_TYPE: u8 = 14;
}

impl CommandTrb for SetTrDequeuePointerCommandTrb {
    const TRB_TYPE: u8 = 16;
}
