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

use alloc::vec::Vec;
use device_api::dma::DmaAllocator;
use libk::{
    dma::{BusAddress, DmaBuffer},
    error::Error,
};
use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker};

pub trait GenericTxDescriptor: Sized {
    const EMPTY: Self;
    const EMPTY_LAST: Self = Self::EMPTY;

    fn consume(&mut self) -> Option<Result<(), Error>>;
    fn setup_tx(
        &mut self,
        buffer_address: BusAddress,
        size: usize,
        index: usize,
        capacity: usize,
    ) -> Result<(), Error>;
}

pub trait GenericRxDescriptor: Sized {
    // TODO Rx flags (short, runt, etc)
    fn consume(&mut self) -> Option<Result<usize, Error>>;
    fn setup_rx(buffer_address: BusAddress, buffer_size: usize, last: bool) -> Result<Self, Error>;
}

struct TxInFlight {
    #[allow(unused)]
    buffer: DmaBuffer<[u8]>,
}

struct TxRingInner<T: GenericTxDescriptor> {
    entries: DmaBuffer<[T]>,
    in_flight: Vec<Option<TxInFlight>>,

    wr: usize,
    rd: usize,
}

struct RxRingInner<T: GenericRxDescriptor> {
    entries: DmaBuffer<[T]>,
    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,

    rd: usize,
}

pub struct GenericTxRing<T: GenericTxDescriptor> {
    inner: IrqSafeSpinlock<TxRingInner<T>>,

    buffer_base: BusAddress,
    capacity: usize,
    free_notify: QueueWaker,
}

pub struct GenericRxRing<T: GenericRxDescriptor> {
    inner: IrqSafeSpinlock<RxRingInner<T>>,

    rx_buffer_size: usize,
    buffer_base: BusAddress,
    capacity: usize,
}

pub struct GenericQueue<Tx: GenericTxDescriptor, Rx: GenericRxDescriptor> {
    pub tx_ring: GenericTxRing<Tx>,
    pub rx_ring: GenericRxRing<Rx>,
}

impl<Tx: GenericTxDescriptor, Rx: GenericRxDescriptor> GenericQueue<Tx, Rx> {
    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        tx_capacity: usize,
        rx_capacity: usize,
        rx_buffer_size: usize,
    ) -> Result<Self, Error> {
        let tx_ring = GenericTxRing::with_capacity(dma, tx_capacity)?;
        let rx_ring = GenericRxRing::with_capacity(dma, rx_capacity, rx_buffer_size)?;
        Ok(Self { tx_ring, rx_ring })
    }

    pub fn tx_buffer_base(&self) -> BusAddress {
        self.tx_ring.buffer_base
    }
    pub fn rx_buffer_base(&self) -> BusAddress {
        self.rx_ring.buffer_base
    }

    pub fn try_push_xmit(&self, frame: DmaBuffer<[u8]>) -> Result<usize, Error> {
        self.tx_ring.try_push_xmit(frame)
    }

    pub fn drop_tx_until(&self, head: usize) {
        self.tx_ring.drop_until(head);
    }
    pub fn consume_tx(&self) -> Result<usize, Error> {
        self.tx_ring.consume()
    }
    pub fn consume_rx<F: Fn(DmaBuffer<[u8]>, usize)>(
        &self,
        dma: &dyn DmaAllocator,
        head: Option<usize>,
        handler: F,
    ) -> Result<usize, Error> {
        self.rx_ring.consume(dma, head, handler)
    }
}

impl<T: GenericTxDescriptor> GenericTxRing<T> {
    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
        let entries = DmaBuffer::new_slice_with(
            dma,
            |i| {
                if i == capacity - 1 {
                    T::EMPTY_LAST
                } else {
                    T::EMPTY
                }
            },
            capacity,
        )?;
        let in_flight = (0..capacity).map(|_| None).collect();
        let buffer_base = entries.bus_address();
        let inner = TxRingInner {
            entries,
            in_flight,
            wr: 0,
            rd: 0,
        };
        Ok(Self {
            inner: IrqSafeSpinlock::new(inner),
            buffer_base,
            capacity,
            free_notify: QueueWaker::new(),
        })
    }

    pub fn base(&self) -> BusAddress {
        self.buffer_base
    }
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    pub async fn push_xmit(&self, _frame: DmaBuffer<[u8]>) -> Result<usize, Error> {
        todo!()
    }

    pub fn try_push_xmit(&self, frame: DmaBuffer<[u8]>) -> Result<usize, Error> {
        let mut inner = self.inner.lock();
        if inner.wr.wrapping_add(1) % self.capacity == inner.rd % self.capacity {
            return Err(Error::WouldBlock);
        }

        let buffer_address = frame.bus_address();
        let buffer_len = frame.len();

        let index = inner.wr % self.capacity;
        assert!(inner.in_flight[index].is_none());

        frame.cache_flush_all(true);
        inner.entries[index].setup_tx(buffer_address, buffer_len, index, self.capacity)?;
        inner.entries.cache_flush_element(index, true);

        inner.in_flight[index] = Some(TxInFlight { buffer: frame });
        inner.wr = inner.wr.wrapping_add(1);

        Ok(inner.wr % self.capacity)
    }

    pub fn drop_until(&self, head: usize) {
        self.inner.lock().rd = head % self.capacity;
    }

    pub fn consume(&self) -> Result<usize, Error> {
        let mut count = 0;

        {
            let mut inner = self.inner.lock();

            while inner.rd != inner.wr {
                let index = inner.rd % self.capacity;
                inner.entries.cache_flush_element(index, false);
                let entry = &mut inner.entries[index];

                if let Some(_status) = entry.consume() {
                    let _ = inner.in_flight[index].take().unwrap();
                    inner.rd = inner.rd.wrapping_add(1);
                    count += 1;
                } else {
                    break;
                }
            }
        }

        if count != 0 {
            self.free_notify.wake_all();
        }

        Ok(count)
    }
}

impl<T: GenericRxDescriptor> GenericRxRing<T> {
    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        capacity: usize,
        rx_buffer_size: usize,
    ) -> Result<Self, Error> {
        let buffers = (0..capacity)
            .map(|_| DmaBuffer::new_uninit_slice(dma, rx_buffer_size))
            .collect::<Result<Vec<_>, _>>()?;
        let entries = DmaBuffer::new_slice_with(
            dma,
            |i| {
                let buffer = buffers[i].bus_address();
                let last = i == capacity - 1;
                T::setup_rx(buffer, rx_buffer_size, last).expect("Rx buffer descriptor setup error")
            },
            capacity,
        )?;
        let buffer_base = entries.bus_address();
        let inner = RxRingInner {
            entries,
            buffers,

            rd: 0,
        };
        Ok(Self {
            inner: IrqSafeSpinlock::new(inner),
            rx_buffer_size,
            buffer_base,
            capacity,
        })
    }

    pub fn base(&self) -> BusAddress {
        self.buffer_base
    }
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    pub fn consume<F: Fn(DmaBuffer<[u8]>, usize)>(
        &self,
        dma: &dyn DmaAllocator,
        head: Option<usize>,
        handler: F,
    ) -> Result<usize, Error> {
        let mut inner = self.inner.lock();

        while head.map_or(true, |tail| tail != inner.rd) {
            let index = inner.rd % self.capacity;

            inner.entries.cache_flush_element(index, false);

            if let Some(status) = inner.entries[index].consume() {
                // Grab the current buffer (the one just written to by the DMA), replace it
                // with the newly allocated one, and mark the descriptor as DMA-owned again
                let new_buffer = DmaBuffer::new_uninit_slice(dma, self.rx_buffer_size)?;
                let new_buffer_address = new_buffer.bus_address();
                let buffer = mem::replace(&mut inner.buffers[index], new_buffer);
                let old_buffer_address = buffer.bus_address();
                let buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };

                match status {
                    Ok(len) => {
                        handler(buffer, len);
                    }
                    Err(_error) => {
                        log::warn!("Drop error packet {old_buffer_address:#x}");
                    }
                }

                inner.entries[index] = T::setup_rx(
                    new_buffer_address,
                    self.rx_buffer_size,
                    index == self.capacity - 1,
                )?;

                inner.rd = inner.rd.wrapping_add(1);
            } else {
                break;
            }
        }

        Ok(inner.rd)
    }
}
