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

use alloc::vec::Vec;
use device_api::dma::DmaAllocator;
use libk::{
    dma::{BusAddress, DmaBuffer},
    error::Error,
};

pub struct TxRing {
    entries: DmaBuffer<[TxDescriptor]>,
    buffers: Vec<Option<DmaBuffer<[u8]>>>,

    wr: usize,
    rd: usize,
}

pub struct RxRing {
    entries: DmaBuffer<[RxDescriptor]>,
    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,

    rd: usize,
}

#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct TxDescriptor {
    tdes0: u32,
    tdes1: u32,
    tdes2: u32,
    tdes3: u32,
}

#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct RxDescriptor {
    rdes0: u32,
    rdes1: u32,
    rdes2: u32,
    rdes3: u32,
}

impl TxRing {
    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
        let entries = DmaBuffer::new_slice(dma, TxDescriptor::empty(), capacity)?;
        let buffers = (0..capacity).map(|_| None).collect();
        Ok(Self {
            entries,
            buffers,

            wr: 0,
            rd: 0,
        })
    }

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

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

    pub fn can_xmit(&self) -> bool {
        self.wr.wrapping_add(1) != self.rd
    }

    pub fn outstanding_tx(&self) -> bool {
        self.wr != self.rd
    }

    pub fn push_xmit(&mut self, frame: DmaBuffer<[u8]>) -> Result<usize, Error> {
        if !self.can_xmit() {
            return Err(Error::WouldBlock);
        }

        let address = frame.bus_address();
        let frame_len = frame.len();

        let capacity = self.entries.len();
        let index = self.wr % capacity;
        assert!(self.buffers[index].is_none());

        frame.cache_flush_all(true);
        self.entries[index].setup_tx(address, frame_len, true)?;
        self.entries.cache_flush_element(index, true);

        self.buffers[index] = Some(frame);
        self.wr = self.wr.wrapping_add(1);

        Ok(self.wr % self.capacity())
    }

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

        loop {
            let index = self.rd % self.entries.len();
            let entry = &self.entries[index];

            if self.rd == self.wr {
                break;
            }

            if let Some(status) = entry.tx_status() {
                if status != 0 {
                    log::warn!("tx_ring[{index}] error: {status:#x}");
                }
                let _ = self.buffers[index].take().unwrap();
                self.rd = self.rd.wrapping_add(1);
                count += 1;
            } else {
                break;
            }
        }

        Ok(count)
    }
}

impl TxDescriptor {
    const TDES3_OWN: u32 = 1 << 31;
    const TDES3_FD: u32 = 1 << 29;
    const TDES3_LD: u32 = 1 << 28;
    const TDES2_IOC: u32 = 1 << 31;

    pub const fn empty() -> Self {
        Self {
            tdes0: 0,
            tdes1: 0,
            tdes2: 0,
            tdes3: 0,
        }
    }

    pub fn tx_status(&self) -> Option<u32> {
        if self.tdes3 & Self::TDES3_OWN == 0 {
            Some(self.tdes3 & !(0xFFFF << 16))
        } else {
            None
        }
    }

    pub fn setup_tx(
        &mut self,
        frame: BusAddress,
        frame_len: usize,
        ioc: bool,
    ) -> Result<(), Error> {
        let tdes0 = frame.try_into_u32().map_err(|_| Error::InvalidArgument)?;
        if frame_len & !0x3FFF != 0 {
            return Err(Error::InvalidArgument);
        }
        let mut tdes2 = frame_len as u32;
        if ioc {
            tdes2 |= Self::TDES2_IOC;
        }
        let tdes3 = Self::TDES3_OWN | Self::TDES3_FD | Self::TDES3_LD;

        self.tdes0 = tdes0;
        self.tdes1 = 0;
        self.tdes2 = tdes2;
        self.tdes3 = tdes3;
        Ok(())
    }
}

impl RxRing {
    pub fn with_capacity(dma: &dyn DmaAllocator, capacity: usize) -> Result<Self, Error> {
        let mut entries = DmaBuffer::new_slice(dma, RxDescriptor::empty(), capacity)?;
        let buffers = (0..capacity)
            .map(|_| DmaBuffer::new_uninit_slice(dma, 4096))
            .collect::<Result<Vec<_>, _>>()?;
        for i in 0..capacity {
            entries[i].setup_rx(buffers[i].bus_address(), true)?;
        }
        Ok(Self {
            buffers,
            entries,

            rd: 0,
        })
    }

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

    pub fn consume<F: Fn(DmaBuffer<[u8]>)>(
        &mut self,
        dma: &dyn DmaAllocator,
        packet_handler: F,
    ) -> Result<usize, Error> {
        let mut count = 0;

        loop {
            let index = self.rd % self.entries.len();
            let entry = &mut self.entries[index];

            if entry.rx_completed().is_some() {
                // 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, 4096)?;
                let new_buffer_address = new_buffer.bus_address();
                let buffer = mem::replace(&mut self.buffers[index], new_buffer);
                let buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };
                // TODO packet size hint
                packet_handler(buffer);
                entry.setup_rx(new_buffer_address, true)?;
                self.rd = self.rd.wrapping_add(1);
                count += 1;
            } else {
                break;
            }
        }

        Ok(count)
    }
}

impl RxDescriptor {
    const RDES3_OWN: u32 = 1 << 31;
    const RDES3_IOC: u32 = 1 << 30;
    const RDES3_BUF1V: u32 = 1 << 24;

    pub const fn empty() -> Self {
        Self {
            rdes0: 0,
            rdes1: 0,
            rdes2: 0,
            rdes3: 0,
        }
    }

    pub fn rx_completed(&self) -> Option<usize> {
        if self.rdes3 & Self::RDES3_OWN == 0 {
            Some((self.rdes3 & 0x7FFF) as usize)
        } else {
            None
        }
    }

    pub fn setup_rx(&mut self, buffer: BusAddress, ioc: bool) -> Result<(), Error> {
        self.rdes0 = buffer.try_into_u32().map_err(|_| Error::InvalidArgument)?;
        self.rdes1 = 0;
        self.rdes2 = 0;
        self.rdes3 = Self::RDES3_BUF1V;
        if ioc {
            self.rdes3 |= Self::RDES3_IOC;
        }
        self.rdes3 |= Self::RDES3_OWN;
        Ok(())
    }
}
