#![no_std]

use core::hint::spin_loop;
use core::sync::atomic::fence;
use core::sync::atomic::AtomicU32;
use core::sync::atomic::Ordering;

use io_uring_dev::mmap_sq_cq;
use io_uring_dev::mmap_sqe_array;
use io_uring_dev::setup;
use io_uring_dev::wakeup;
use io_uring_dev::CqPtr;
use io_uring_dev::SqPtr;
use io_uring_sys::io_uring_cqe;
use io_uring_sys::io_uring_sqe;
use io_uring_sys::IORING_FEAT_FAST_POLL;
use io_uring_sys::IORING_FEAT_NODROP;
use io_uring_sys::IORING_FEAT_RW_CUR_POS;
use io_uring_sys::IORING_FEAT_SINGLE_MMAP;
use io_uring_sys::IORING_FEAT_SQPOLL_NONFIXED;
use io_uring_sys::IORING_FEAT_SUBMIT_STABLE;
use io_uring_sys::IORING_SQ_CQ_OVERFLOW;
use io_uring_sys::IORING_SQ_NEED_WAKEUP;
use log::info;
use log::trace;

/// TODO
/// # Safety
/// TODO
pub unsafe fn new(sq_entries: u32, cq_entries: u32) -> (Sq, Cq) {
    let (ring, params) = setup(sq_entries, cq_entries).unwrap();
    let needed = IORING_FEAT_SINGLE_MMAP
        | IORING_FEAT_NODROP
        | IORING_FEAT_SUBMIT_STABLE
        | IORING_FEAT_RW_CUR_POS
        | IORING_FEAT_FAST_POLL
        | IORING_FEAT_SQPOLL_NONFIXED;
    if params.features & needed != needed {
        panic!("Kernel version too low.");
    }
    let (sq, cq, cqe_array) = mmap_sq_cq(ring, &params).unwrap();
    let sqe_array = mmap_sqe_array(ring, &params).unwrap();
    let sq = Sq::new(sq, sqe_array, ring);
    let cq = Cq::new(cq, cqe_array);
    (sq, cq)
}

pub struct Sq {
    ring: u32,
    entries: u32,
    mask: u32,
    tail: u32,
    k_head: &'static AtomicU32,
    k_tail: &'static AtomicU32,
    k_flags: &'static AtomicU32,
    array: *mut io_uring_sqe,
}

impl Sq {
    /// TODO
    /// # Safety
    /// TODO
    pub unsafe fn push(&mut self, producer: impl FnOnce(&mut io_uring_sqe)) {
        let tail = self.get_tail();

        while tail.wrapping_sub(self.get_head()) == self.entries {
            trace!("Waiting for kernel to consume some sqe.");
            spin_loop();
        }

        let sqe = self.get_sqe((tail & self.mask) as usize);
        producer(&mut *sqe);
        self.set_tail(tail + 1);

        fence(Ordering::SeqCst);
        let flags = self.get_flags();

        if flags & IORING_SQ_CQ_OVERFLOW != 0 {
            panic!("Completion queue overflowed.");
        }

        if flags & IORING_SQ_NEED_WAKEUP != 0 {
            info!("IORING_SQ_NEED_WAKEUP is set.");
            wakeup(self.ring).unwrap();

            while self.get_flags() & IORING_SQ_NEED_WAKEUP != 0 {
                trace!("Waiting for kernel thread to wakeup.");
                spin_loop();
            }
        }
    }

    unsafe fn new(ptr: SqPtr, array: *mut io_uring_sqe, ring: u32) -> Self {
        Self {
            ring,
            entries: ptr.entries,
            mask: ptr.mask,
            tail: 0,
            k_head: &*(ptr.head.cast()),
            k_tail: &*(ptr.tail.cast()),
            k_flags: &*(ptr.flags.cast()),
            array,
        }
    }

    unsafe fn get_head(&self) -> u32 {
        self.k_head.load(Ordering::Acquire)
    }

    unsafe fn get_tail(&self) -> u32 {
        self.tail
    }

    unsafe fn set_tail(&mut self, val: u32) {
        self.k_tail.store(val, Ordering::Release);
        self.tail = val;
    }

    unsafe fn get_flags(&self) -> u32 {
        self.k_flags.load(Ordering::Relaxed)
    }

    unsafe fn get_sqe(&self, index: usize) -> *mut io_uring_sqe {
        self.array.add(index)
    }
}

pub struct Cq {
    mask: u32,
    head: u32,
    k_head: &'static AtomicU32,
    k_tail: &'static AtomicU32,
    array: *mut io_uring_cqe,
}

impl Cq {
    /// TODO
    /// # Safety
    /// TODO
    pub unsafe fn pop(&mut self, consumer: impl FnOnce(&mut io_uring_cqe)) {
        let head = self.get_head();

        while self.get_tail().wrapping_sub(head) == 0 {
            trace!("Waiting for kernel to produce some cqe.");
            spin_loop();
        }

        let cqe = self.get_cqe((head & self.mask) as usize);
        consumer(&mut *cqe);
        self.set_head(head + 1);
    }

    unsafe fn new(ptr: CqPtr, array: *mut io_uring_cqe) -> Self {
        Self {
            mask: ptr.mask,
            head: 0,
            k_head: &*(ptr.head.cast()),
            k_tail: &*(ptr.tail.cast()),
            array,
        }
    }

    unsafe fn get_head(&self) -> u32 {
        self.head
    }

    unsafe fn set_head(&mut self, val: u32) {
        self.k_head.store(val, Ordering::Release);
        self.head = val;
    }

    unsafe fn get_tail(&self) -> u32 {
        self.k_tail.load(Ordering::Acquire)
    }

    unsafe fn get_cqe(&self, index: usize) -> *mut io_uring_cqe {
        self.array.add(index)
    }
}
