use alloc::vec::Vec;

use crate::xhci::data_structures::{TransferRequestBlock, TrbLink, TrbRingSegment, TypedTrb};

#[derive(Default, Copy, Clone)]
struct TrbPointer {
    /// Index into the vector of trb segments.
    segment_index: usize,
    /// Index into the specific segment.
    /// This is a TransferRequestBlock index,
    /// to get the physical_offset use segment_physical_offset()
    segment_offset: usize,
}

impl TrbPointer {
    fn segment_physical_offset(&self) -> usize {
        self.segment_offset * size_of::<TransferRequestBlock>()
    }
}

pub struct TrbRing {
    segments: Vec<TrbRingSegment>,
    enqueue_pointer: TrbPointer,
    dequeue_pointer: TrbPointer,
    cycle_bit: bool,
}

impl TrbRing {
    pub fn new() -> Self {
        Self {
            // TODO: What size and count should this be.
            segments: alloc::vec![TrbRingSegment::new(100)],
            enqueue_pointer: TrbPointer::default(),
            dequeue_pointer: TrbPointer::default(),
            // Start with this as true so we are flipping bits from 0 (default) to 1
            // to mark the enqueue pointer.
            cycle_bit: true,
        }
    }

    pub fn physical_base_address(&self) -> usize {
        self.segments[0].physical_address()
    }

    pub fn enqueue_trb(&mut self, trb: TransferRequestBlock) {
        *self.next_trb_ref() = trb.with_cycle(self.cycle_bit);
        self.advance_enqueue_pointer();
    }

    fn next_trb_ref(&mut self) -> &mut TransferRequestBlock {
        &mut self.segments[self.enqueue_pointer.segment_index][self.enqueue_pointer.segment_offset]
    }

    fn advance_enqueue_pointer(&mut self) {
        self.enqueue_pointer.segment_offset += 1;

        if self.enqueue_pointer.segment_offset
            == self.segments[self.enqueue_pointer.segment_index].len() - 1
        {
            // We have reached the end of the segment, insert a link trb.

            // Increment the segment index with wrapping.
            let next_segment_index =
                if self.enqueue_pointer.segment_index + 1 == self.segments.len() {
                    0
                } else {
                    self.enqueue_pointer.segment_index + 1
                };

            let next_segment_pointer = self.segments[next_segment_index].physical_address();
            let toggle_cycle = next_segment_index == 0;

            *self.next_trb_ref() = TrbLink::new()
                .with_ring_segment_pointer(next_segment_pointer as u64)
                .with_cycle(self.cycle_bit)
                .with_toggle_cycle(toggle_cycle)
                .to_trb();

            // Flip toggle cycle bit if necessary.
            self.cycle_bit ^= toggle_cycle;

            self.enqueue_pointer = TrbPointer {
                segment_index: next_segment_index,
                segment_offset: 0,
            };
        }
    }
}