mod buf;
pub mod config;
mod cqe;
mod current;
pub mod handle;
mod provided;
mod recv_msg_out;
mod request;
mod resources;
mod ring;
mod submit;
mod wire;

use std::io;
use std::marker::PhantomData;
use std::os::fd::AsRawFd;
use std::time::Duration;

use io_uring::{opcode, types};

use dope_core::addr::SockAddr;
use dope_core::driver::DriverOps;
use dope_core::submit::SendMsgPacket;

use crate::sys::Sqe;
use crate::sys::fd::FdUringExt;
use crate::sys::token::Token;
use crate::{Fd, FixedFd};

pub use self::buf::{FixedBufGuard, FixedBufWrite, ProvidedBuf};
pub use self::cqe::{CqeKind, CqeRouter, CqeView};
pub use self::current::{CurrentSlot, current, current_driver_mut, set_current};
pub use self::provided::ProvidedRing;
pub use self::recv_msg_out::RecvMsgOut;
pub use self::request::Request;

pub type DriverRef<'a> = &'a Driver;

pub(crate) use self::request::{RequestId, Requests};
pub(crate) use self::resources::Resources;
pub use self::ring::Ring;

use self::cqe::{RawCqe, make_view};
use self::wire::Wire;

pub struct Driver {
    pub(crate) requests: Requests,
    pub(crate) ring: Ring,
    pub(crate) resources: Resources,
    pub(crate) provided: ProvidedRing,
    pub(in crate::driver) cqe_scratch: Vec<self::cqe::RawCqe>,

    _not_send: PhantomData<*const ()>,
}

impl Drop for Driver {
    fn drop(&mut self) {
        let self_ptr = self as *mut Driver;
        if let Some(slot) = self::current::current()
            && slot.driver().as_ptr() == self_ptr
        {
            self::current::set_current(None);
        }
    }
}

impl Driver {
    pub(crate) fn new(ring: Ring, cfg: crate::DriverConfig, capacity: usize) -> io::Result<Self> {
        let provided = ProvidedRing::new(&ring, cfg.provided_buf_entries, cfg.provided_buf_len)
            .map_err(|e| {
                io::Error::other(format!(
                    "dope: required capability setup failed (provided buffer ring): {e}"
                ))
            })?;
        let resources = Resources::new(&ring, &cfg)?;
        Ok(Self {
            requests: Requests::with_capacity(capacity),
            ring,
            resources,
            provided,
            cqe_scratch: Vec::with_capacity(capacity),
            _not_send: PhantomData,
        })
    }

    #[inline(always)]
    pub fn has_pending(&mut self) -> bool {
        self.ring.has_submission() || self.ring.has_completion()
    }

    #[inline(always)]
    pub fn has_completion(&mut self) -> bool {
        self.ring.has_completion()
    }

    #[inline(always)]
    pub fn park(&mut self) -> io::Result<()> {
        self.ring.submit_and_wait(1)
    }

    #[inline(always)]
    pub fn park_for(&mut self, duration: Duration) -> io::Result<()> {
        self.ring.submit_with_timeout(1, duration)
    }
}

impl Driver {
    fn submit_retry_error(e: &io::Error) -> bool {
        matches!(e.raw_os_error(), Some(libc::EAGAIN) | Some(libc::EBUSY))
    }

    pub(super) fn submit_retry(&mut self) -> io::Result<usize> {
        match self.ring.submit() {
            Ok(submitted) => Ok(submitted),
            Err(ref e) if Self::submit_retry_error(e) => {
                self.tick()?;
                Ok(0)
            }
            Err(e) => Err(e),
        }
    }

    #[inline(always)]
    fn drive_loop(
        &mut self,
        mut on_completion: impl FnMut(&mut Self) -> io::Result<()>,
    ) -> io::Result<bool> {
        let mut any = false;
        for _ in 0..256 {
            let mut progressed = false;
            if self.ring.has_submission() {
                let submitted = self.submit_retry()?;
                if submitted != 0 {
                    progressed = true;
                }
            }
            if self.ring.has_completion() {
                on_completion(self)?;
                progressed = true;
            }
            if !progressed {
                break;
            }
            any = true;
        }
        Ok(any)
    }

    pub fn drive(&mut self) -> io::Result<bool> {
        self.drive_loop(Self::tick)
    }

    pub fn drive_with<H: CqeRouter>(&mut self, handler: &mut H) -> io::Result<bool> {
        self.drive_loop(|s| {
            s.dispatch_cqes(handler);
            Ok(())
        })
    }

    pub(crate) fn tick(&mut self) -> io::Result<()> {
        let Self {
            requests,
            ring,
            provided,
            ..
        } = self;
        let mut cq = ring.completion();
        for item in &mut cq {
            let result = item.result();
            let flags = item.flags();
            if let Wire::Request(id) = wire::decode_wire(item.user_data()) {
                requests.on_request(id, result, flags);
            }
        }
        cq.sync();
        let _ = provided.flush();
        Ok(())
    }

    #[inline(always)]
    pub fn cancel_write(&mut self, token: Token) {
        self.ring.submit_cancel(wire::write(token));
    }

    #[inline(always)]
    pub fn cancel_accept(&mut self, token: Token) {
        self.ring.submit_cancel(wire::accept(token));
    }

    fn drain_to_scratch(&mut self) {
        let Self {
            ring, cqe_scratch, ..
        } = self;
        cqe_scratch.clear();
        let mut cq = ring.completion();
        for item in &mut cq {
            cqe_scratch.push(RawCqe {
                user_data: item.user_data(),
                result: item.result(),
                flags: item.flags(),
            });
        }
        cq.sync();
    }

    pub fn dispatch_cqes<H: CqeRouter>(&mut self, handler: &mut H) {
        self.drain_to_scratch();
        let mut scratch = std::mem::take(&mut self.cqe_scratch);
        for raw in scratch.drain(..) {
            let wire = wire::decode_wire(raw.user_data);
            match wire {
                Wire::Request(id) => self.requests.on_request(id, raw.result, raw.flags),
                Wire::Ignore => {}
                Wire::Accept(_) | Wire::Recv(_) | Wire::Write(_) => {
                    handler.on_cqe(self, make_view(wire, raw.result, raw.flags));
                }
            }
        }
        self.cqe_scratch = scratch;
        let _ = self.provided.flush();
    }
}

impl Driver {
    pub fn submit_only(&mut self) -> io::Result<()> {
        if self.ring.submission_len() == 0 {
            return Ok(());
        }
        for _ in 0..256 {
            if self.ring.submission_len() == 0 {
                return Ok(());
            }
            self.submit_retry()?;
        }
        Ok(())
    }

    pub(crate) fn submit_sqe_op(&mut self, sqe: crate::Sqe) -> io::Result<RequestId> {
        if self.ring.is_full() {
            self.submit_only()?;
        }
        let index = self.requests.insert();
        let sqe = sqe.user_data(wire::request(index));
        if !self.ring.push(&sqe) {
            self.requests.rollback(index);
            return Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "SQ full after flush",
            ));
        }
        Ok(index)
    }
}

impl Driver {
    #[inline(always)]
    pub(crate) fn try_submit_op<O, F>(
        &mut self,
        op: O,
        sqe_builder: F,
    ) -> Result<Request<O>, (io::Error, O)>
    where
        O: Unpin + 'static,
        F: FnOnce(&mut O) -> Sqe,
    {
        let mut boxed = Box::new(op);
        let sqe = sqe_builder(&mut *boxed);
        Request::submit_boxed(self, boxed, sqe).map_err(|(e, b)| (e, *b))
    }

    #[inline(always)]
    pub(crate) fn try_submit_op_boxed<O, F>(
        &mut self,
        mut boxed: Box<O>,
        sqe_builder: F,
    ) -> Result<Request<O>, (io::Error, Box<O>)>
    where
        O: Unpin + 'static,
        F: FnOnce(&mut O) -> Sqe,
    {
        let sqe = sqe_builder(&mut *boxed);
        Request::submit_boxed(self, boxed, sqe)
    }

    pub fn try_submit_sendmsg_raw_boxed<P>(
        &mut self,
        fd: Fd,
        packet: Box<P>,
    ) -> Result<Request<P>, (io::Error, Box<P>)>
    where
        P: SendMsgPacket + Unpin + 'static,
    {
        let raw = fd.as_raw_fd();
        self.try_submit_op_boxed(packet, move |packet| {
            io_uring::opcode::SendMsg::new(io_uring::types::Fd(raw), packet.msg_ptr()).build()
        })
    }
}

impl DriverOps for Driver {
    type Token = Token;
    type ProvidedRing = ProvidedRing;
    type IoRequest<O: Unpin + 'static> = Request<O>;

    #[inline(always)]
    fn provided_ring(&mut self) -> &mut Self::ProvidedRing {
        &mut self.provided
    }

    #[inline(always)]
    fn buf_group(&self) -> u16 {
        self.provided.group().0
    }

    #[inline(always)]
    fn arm_recv_multi(&mut self, token: Self::Token, fd: FixedFd, buf_group: u16) -> bool {
        let sqe = opcode::RecvMulti::new(types::Fixed(fd.fixed_index()), buf_group)
            .build()
            .user_data(wire::recv(token));
        self.ring.try_push_arm(sqe)
    }

    #[inline(always)]
    unsafe fn arm_recv_msg_multi(
        &mut self,
        token: Self::Token,
        fd: FixedFd,
        buf_group: u16,
        msghdr: *const libc::msghdr,
    ) -> bool {
        let sqe = opcode::RecvMsgMulti::new(types::Fixed(fd.fixed_index()), msghdr, buf_group)
            .build()
            .user_data(wire::recv(token));
        self.ring.try_push_arm(sqe)
    }

    #[inline(always)]
    fn arm_accept_multi(&mut self, token: Self::Token, fd: Fd) -> bool {
        let sqe = opcode::AcceptMulti::new(fd.uring())
            .flags(libc::SOCK_CLOEXEC)
            .build()
            .user_data(wire::accept(token));
        self.ring.try_push_arm(sqe)
    }

    #[inline(always)]
    fn submit_send_tagged(
        &mut self,
        token: Self::Token,
        fd: FixedFd,
        ptr: *const u8,
        len: u32,
    ) -> bool {
        let sqe = opcode::Send::new(types::Fixed(fd.fixed_index()), ptr, len)
            .build()
            .user_data(wire::write(token));
        self.ring.push(&sqe)
    }

    #[inline(always)]
    unsafe fn submit_send_msg_tagged(
        &mut self,
        token: Self::Token,
        fd: FixedFd,
        msg: *const libc::msghdr,
    ) -> bool {
        let sqe = opcode::SendMsg::new(types::Fixed(fd.fixed_index()), msg)
            .build()
            .user_data(wire::write(token));
        self.ring.push(&sqe)
    }

    #[inline(always)]
    fn cancel_recv(&mut self, token: Self::Token) {
        self.ring.submit_cancel(wire::recv(token));
    }

    fn try_submit_connect_raw(
        &mut self,
        fd: Fd,
        addr: SockAddr,
    ) -> Result<Self::IoRequest<SockAddr>, (io::Error, SockAddr)> {
        let raw = fd.as_raw_fd();
        self.try_submit_op(addr, move |addr| {
            io_uring::opcode::Connect::new(
                io_uring::types::Fd(raw),
                addr.addr_ptr(),
                addr.addr_len(),
            )
            .build()
        })
    }

    fn try_submit_sendmsg_raw<P: SendMsgPacket + Unpin + 'static>(
        &mut self,
        fd: Fd,
        packet: P,
    ) -> Result<Self::IoRequest<P>, (io::Error, P)> {
        let raw = fd.as_raw_fd();
        self.try_submit_op(packet, move |packet| {
            io_uring::opcode::SendMsg::new(io_uring::types::Fd(raw), packet.msg_ptr()).build()
        })
    }

    fn try_submit_send_raw<B: AsRef<[u8]> + Unpin + 'static>(
        &mut self,
        fd: Fd,
        buf: B,
    ) -> Result<Self::IoRequest<B>, (io::Error, B)> {
        let raw = fd.as_raw_fd();
        self.try_submit_op(buf, move |buf| {
            let slice = buf.as_ref();
            io_uring::opcode::Send::new(
                io_uring::types::Fd(raw),
                slice.as_ptr(),
                slice.len() as u32,
            )
            .build()
        })
    }
}
