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

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

pub(crate) struct RxRing {
    pub(crate) descriptors: DmaBuffer<[RxDescriptor]>,
    buffers: Vec<DmaBuffer<[MaybeUninit<u8>]>>,
    buffer_size: usize,
    tail: u16,
}

pub(crate) struct TxRing {
    pub(crate) descriptors: DmaBuffer<[TxDescriptor]>,
    buffers: Vec<Option<DmaBuffer<[u8]>>>,
    // Consumer end
    tail: u16,
    // Producer end
    head: u16,
}

#[repr(C)]
pub(crate) struct RxDescriptor {
    address: BusAddress,
    length: u16,
    checksum: u16,
    status: u8,
    errors: u8,
    special: u16,
}

#[repr(C)]
pub(crate) struct TxDescriptor {
    address: BusAddress,
    length: u16,
    cso: u8,
    cmd: u8,
    sta: u8,
    _0: u8,
    css: u8,
    special: u8,
}

impl RxRing {
    pub fn with_capacity(
        dma: &dyn DmaAllocator,
        capacity: usize,
        buffer_size: usize,
    ) -> Result<Self, Error> {
        let buffers = (0..capacity)
            .map(|_| DmaBuffer::new_uninit_slice(dma, buffer_size))
            .collect::<Result<Vec<_>, _>>()?;
        let descriptors = DmaBuffer::new_slice_with(
            dma,
            |i| RxDescriptor::new(buffers[i].bus_address(), buffer_size as u16),
            capacity,
        )?;

        Ok(Self {
            descriptors,
            buffers,
            tail: 0,
            buffer_size,
        })
    }

    // TODO move to background task/softirq to reduce amount of code run by the irq handler
    pub fn handle_rx<F: FnMut(DmaBuffer<[u8]>, usize)>(
        &mut self,
        dma: &dyn DmaAllocator,
        head: u16,
        mut handler: F,
    ) -> u16 {
        let capacity = self.descriptors.len();
        while self.tail != head {
            let index = self.tail as usize;
            // Replace the buffer

            let new_buffer = DmaBuffer::new_uninit_slice(dma, self.buffer_size).unwrap();
            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) };

            let descriptor = &mut self.descriptors[index];

            if descriptor.errors & !1 != 0 {
                log::warn!("igbe: drop erroneous packet {:#x}", descriptor.errors);
            } else {
                let len = descriptor.length as usize;
                handler(buffer, len);
            }

            // Replace the descriptor
            *descriptor = RxDescriptor::new(new_buffer_address, self.buffer_size as u16);

            self.tail = (self.tail + 1) & (capacity as u16 - 1);
        }

        (self.tail + capacity as u16 - 1) & (capacity as u16 - 1)
    }
}

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

    pub fn handle_tx(&mut self, head: u16) {
        self.tail = head;
    }

    pub fn tx_now(&mut self, buffer: DmaBuffer<[u8]>) -> Result<u16, DmaBuffer<[u8]>> {
        // Queue full
        let capacity = self.descriptors.len();
        if (self.head + 1) & (capacity as u16 - 1) == self.tail {
            log::warn!("igbe: tx queue full");
            return Err(buffer);
        }
        let index = self.head as usize;

        let descriptor = &mut self.descriptors[index];
        // Only generate interrupts for every 1/4th of the buffer
        let quarter = capacity / 4;
        descriptor.setup_tx(
            buffer.bus_address(),
            buffer.len() as u16 - 4,
            index % quarter == quarter - 1,
        );
        self.descriptors.cache_flush_element(index, true);

        self.buffers[index] = Some(buffer);

        self.head = (self.head + 1) & (capacity as u16 - 1);
        Ok(self.head)
    }
}

impl RxDescriptor {
    pub fn new(address: BusAddress, length: u16) -> Self {
        Self {
            address,
            length,
            checksum: 0,
            status: 0,
            errors: 0,
            special: 0,
        }
    }
}

impl TxDescriptor {
    // Descriptor done
    const STA_DD: u8 = 1 << 0;
    // End of packet
    const CMD_EOP: u8 = 1 << 0;
    // Insert MAC FCS
    const CMD_IFCS: u8 = 1 << 1;
    // // Insert checksum
    // const CMD_IC: u8 = 1 << 2;
    // Report status
    const CMD_RS: u8 = 1 << 3;

    pub const fn empty() -> Self {
        Self {
            address: BusAddress::ZERO,
            length: 0,
            cso: 0,
            cmd: 0,
            sta: Self::STA_DD,
            _0: 0,
            css: 0,
            special: 0,
        }
    }

    pub fn setup_tx(&mut self, address: BusAddress, length: u16, ioc: bool) {
        let mut cmd = Self::CMD_EOP | Self::CMD_IFCS;
        if ioc {
            cmd |= Self::CMD_RS;
        }
        self.address = address;
        self.length = length;
        self.css = 0;
        self.cso = 0;
        self.sta = 0;
        self.cmd = cmd;
        self.special = 0;
    }
}
