use core::mem::{size_of, MaybeUninit};

use bytemuck::{Pod, Zeroable};
use device_api::dma::DmaAllocator;
use libk::dma::{BusAddress, DmaBuffer};
use libk_util::sync::IrqSafeSpinlock;
use ygg_driver_usb::error::UsbError;
use yggdrasil_abi::define_bitfields;

use super::{command::CommandReply, GenericRing};

pub enum Event {
    PortChange(usize),
    CommandCompletion {
        address: BusAddress,
        reply: CommandReply,
    },
    Transfer {
        address: BusAddress,
        slot_id: u8,
        endpoint_id: u8,
        status: u32,
    },
}

#[repr(C, align(16))]
pub struct EventRingSegment {
    address: BusAddress,
    // Number of TRBs supported by the ring segment. Valid values are 16 to 4096
    size: u16,
    _0: u16,
    _1: u32,
}

pub struct EventRingSegmentTable {
    entries: DmaBuffer<[EventRingSegment]>,
}

struct EventRingInner {
    trbs: DmaBuffer<[MaybeUninit<RawEventTrb>]>,
    dequeue_index: usize,
    cycle_bit: bool,
}

pub struct EventRing {
    inner: IrqSafeSpinlock<EventRingInner>,
    capacity: usize,
}

impl EventRingSegmentTable {
    pub fn for_event_rings(dma: &dyn DmaAllocator, rings: &[&EventRing]) -> Result<Self, UsbError> {
        let entries = DmaBuffer::new_slice_with(
            dma,
            |i| EventRingSegment {
                address: rings[i].bus_base(),
                size: rings[i].capacity.try_into().unwrap(),
                _0: 0,
                _1: 0,
            },
            rings.len(),
        )
        .map_err(UsbError::MemoryError)?;

        Ok(Self { entries })
    }

    pub fn bus_address(&self) -> BusAddress {
        self.entries.bus_address()
    }

    pub fn capacity(&self) -> usize {
        self.entries.len()
    }
}

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

impl EventRingInner {
    fn try_dequeue(&mut self) -> Option<Event> {
        let trb = unsafe { self.trbs[self.dequeue_index].assume_init_ref() };

        // TRB cannot be consumed -- its cycle bit not toggled
        let trb_cycle = trb.cycle_bit();
        if trb_cycle != self.cycle_bit {
            return None;
        }

        self.dequeue_index += 1;

        if self.dequeue_index == self.trbs.len() {
            self.dequeue_index = 0;
            self.cycle_bit = !self.cycle_bit;
        }

        trb.into_event()
    }
}

impl EventRing {
    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(EventRingInner {
                trbs,
                dequeue_index: 0,
                cycle_bit: true,
            }),
            capacity,
        })
    }

    pub fn try_dequeue(&self) -> Option<Event> {
        self.inner.lock().try_dequeue()
    }

    pub fn dequeue_pointer(&self) -> BusAddress {
        let i = self.inner.lock();
        i.trbs
            .bus_address()
            .add(i.dequeue_index * size_of::<RawEventTrb>())
    }
}

// TRB implementations

define_bitfields! {
    pub TransferEventStatus : u32 {
        (24..32) => completion_code,
        (0..24) => sub_length
    }
}

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

define_bitfields! {
    pub CommandCompletionEventStatus : u32 {
        (24..32) => completion_code,
        (0..24) => completion_parameter
    }
}

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

define_bitfields! {
    pub PortStatusChangeEventAddress : u32 {
        (24..32) => port_id
    }
}

define_bitfields! {
    pub RawEventFlags : u32 {
        (10..16) => ty,
        0 => cycle
    }
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct TransferEventTrb {
    pub address: BusAddress,
    pub status: TransferEventStatus,
    pub flags: TransferEventFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct CommandCompletionEventTrb {
    pub address: BusAddress,
    pub status: CommandCompletionEventStatus,
    pub flags: CommandCompletionEventFlags,
}

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

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

impl RawEventTrb {
    pub fn into_event(self) -> Option<Event> {
        match self.flags.ty() {
            32 => {
                let transfer: TransferEventTrb = bytemuck::cast(self);

                Some(Event::Transfer {
                    address: transfer.address,
                    slot_id: transfer.flags.slot_id() as _,
                    endpoint_id: transfer.flags.endpoint_id() as _,
                    status: transfer.status.into_raw(),
                })
            }
            33 => {
                let command: CommandCompletionEventTrb = bytemuck::cast(self);

                Some(Event::CommandCompletion {
                    address: command.address,
                    reply: CommandReply {
                        completion_code: command.status.completion_code() as _,
                        completion_parameter: command.status.completion_parameter(),
                        slot_id: command.flags.slot_id() as _,
                    },
                })
            }
            34 => {
                let port_status: PortStatusChangeEventTrb = bytemuck::cast(self);

                Some(Event::PortChange(port_status.address.port_id() as _))
            }
            ty => {
                log::warn!("Unhandled event TRB with type: {}", ty);
                None
            }
        }
    }

    pub fn cycle_bit(&self) -> bool {
        self.flags.cycle()
    }
}