use std::cell::Cell;
use std::io;
use std::os::fd::RawFd;
use std::ptr;
use std::sync::atomic::{Ordering, compiler_fence};
use std::sync::{Mutex, OnceLock};
use std::time::Duration;

const IORING_OFF_SQ_RING: libc::off_t = 0;
const IORING_OFF_CQ_RING: libc::off_t = 0x0800_0000;
const IORING_OFF_SQES: libc::off_t = 0x1000_0000;

const IORING_ENTER_GETEVENTS: u32 = 1 << 0;
const IORING_SETUP_CLAMP: u32 = 1 << 4;

const IORING_FEAT_SINGLE_MMAP: u32 = 1 << 0;

pub(crate) const IORING_OP_FSYNC: u8 = 3;
pub(crate) const IORING_OP_POLL_ADD: u8 = 6;
pub(crate) const IORING_OP_SENDMSG: u8 = 9;
pub(crate) const IORING_OP_RECVMSG: u8 = 10;
pub(crate) const IORING_OP_TIMEOUT: u8 = 11;
pub(crate) const IORING_OP_LINK_TIMEOUT: u8 = 15;
pub(crate) const IORING_OP_TIMEOUT_REMOVE: u8 = 12;
pub(crate) const IORING_OP_ACCEPT: u8 = 13;
pub(crate) const IORING_OP_ASYNC_CANCEL: u8 = 14;
pub(crate) const IORING_OP_CONNECT: u8 = 16;
pub(crate) const IORING_OP_OPENAT: u8 = 18;
pub(crate) const IORING_OP_CLOSE: u8 = 19;
pub(crate) const IORING_OP_STATX: u8 = 21;
pub(crate) const IORING_OP_READ: u8 = 22;
pub(crate) const IORING_OP_WRITE: u8 = 23;
pub(crate) const IORING_OP_SEND: u8 = 26;
pub(crate) const IORING_OP_RECV: u8 = 27;
pub(crate) const IORING_OP_SHUTDOWN: u8 = 34;
pub(crate) const IORING_OP_RENAMEAT: u8 = 35;
pub(crate) const IORING_OP_UNLINKAT: u8 = 36;
pub(crate) const IORING_OP_MKDIRAT: u8 = 37;
pub(crate) const IORING_OP_MSG_RING: u8 = 40;
pub(crate) const IORING_OP_SOCKET: u8 = 45;
pub(crate) const IORING_OP_FTRUNCATE: u8 = 55;
pub(crate) const IORING_OP_BIND: u8 = 56;
pub(crate) const IORING_OP_LISTEN: u8 = 57;

const IORING_MSG_DATA: u64 = 0;
pub(crate) const IORING_FSYNC_DATASYNC: u32 = 1 << 0;
pub(crate) const IORING_TIMEOUT_ABS: u32 = 1 << 0;
pub(crate) const IORING_TIMEOUT_UPDATE: u32 = 1 << 1;
pub(crate) const IOSQE_IO_LINK: u8 = 1 << 2;
pub(crate) const IOSQE_CQE_SKIP_SUCCESS: u8 = 1 << 6;

thread_local! {
    static CURRENT_SUBMITTER: Cell<*const IoUring> = const { Cell::new(ptr::null()) };
}

static GLOBAL_SUBMITTER: OnceLock<Mutex<Option<IoUring>>> = OnceLock::new();

#[repr(C)]
#[derive(Default, Clone, Copy)]
struct IoSqringOffsets {
    head: u32,
    tail: u32,
    ring_mask: u32,
    ring_entries: u32,
    flags: u32,
    dropped: u32,
    array: u32,
    resv1: u32,
    user_addr: u64,
}

#[repr(C)]
#[derive(Default, Clone, Copy)]
struct IoCqringOffsets {
    head: u32,
    tail: u32,
    ring_mask: u32,
    ring_entries: u32,
    overflow: u32,
    cqes: u32,
    flags: u32,
    resv1: u32,
    user_addr: u64,
}

#[repr(C)]
#[derive(Default, Clone, Copy)]
struct IoUringParams {
    sq_entries: u32,
    cq_entries: u32,
    flags: u32,
    sq_thread_cpu: u32,
    sq_thread_idle: u32,
    features: u32,
    wq_fd: u32,
    resv: [u32; 3],
    sq_off: IoSqringOffsets,
    cq_off: IoCqringOffsets,
}

#[repr(C)]
#[derive(Default, Clone, Copy)]
pub(crate) struct IoUringSqe {
    pub(crate) opcode: u8,
    pub(crate) flags: u8,
    pub(crate) ioprio: u16,
    pub(crate) fd: i32,
    pub(crate) off: u64,
    pub(crate) addr: u64,
    pub(crate) len: u32,
    pub(crate) op_flags: u32,
    pub(crate) user_data: u64,
    pub(crate) buf_index: u16,
    pub(crate) personality: u16,
    pub(crate) file_index: i32,
    pub(crate) pad2: [u64; 2],
}

#[repr(C)]
#[derive(Default, Clone, Copy)]
pub(crate) struct IoUringCqe {
    pub(crate) user_data: u64,
    pub(crate) res: i32,
    pub(crate) flags: u32,
}

#[repr(C)]
#[derive(Default, Clone, Copy)]
struct KernelTimespec {
    tv_sec: i64,
    tv_nsec: i64,
}

pub(crate) struct IoUring {
    ring_fd: RawFd,
    sq_ring_ptr: *mut u8,
    cq_ring_ptr: *mut u8,
    sqes_ptr: *mut IoUringSqe,
    sq_ring_size: usize,
    cq_ring_size: usize,
    sqes_size: usize,
    single_mmap: bool,
    sq_head: *mut u32,
    sq_tail: *mut u32,
    sq_ring_mask: *mut u32,
    sq_ring_entries: *mut u32,
    sq_array: *mut u32,
    cq_head: *mut u32,
    cq_tail: *mut u32,
    cq_ring_mask: *mut u32,
    cqes: *mut IoUringCqe,
}

impl IoUring {
    pub(crate) fn new(entries: u32) -> io::Result<Self> {
        let mut params = IoUringParams {
            flags: IORING_SETUP_CLAMP,
            ..IoUringParams::default()
        };

        let ring_fd = cvt_long(unsafe {
            libc::syscall(
                libc::SYS_io_uring_setup,
                entries as libc::c_uint,
                &mut params as *mut IoUringParams,
            )
        })? as RawFd;

        let sq_ring_size =
            params.sq_off.array as usize + params.sq_entries as usize * std::mem::size_of::<u32>();
        let cq_ring_size = params.cq_off.cqes as usize
            + params.cq_entries as usize * std::mem::size_of::<IoUringCqe>();
        let single_mmap = params.features & IORING_FEAT_SINGLE_MMAP != 0;

        let sq_ring_ptr = mmap_ring(
            if single_mmap {
                sq_ring_size.max(cq_ring_size)
            } else {
                sq_ring_size
            },
            ring_fd,
            IORING_OFF_SQ_RING,
        )?;
        let cq_ring_ptr = if single_mmap {
            sq_ring_ptr
        } else {
            mmap_ring(cq_ring_size, ring_fd, IORING_OFF_CQ_RING)?
        };
        let sqes_size = params.sq_entries as usize * std::mem::size_of::<IoUringSqe>();
        let sqes_ptr = mmap_ring(sqes_size, ring_fd, IORING_OFF_SQES)? as *mut IoUringSqe;

        Ok(Self {
            ring_fd,
            sq_ring_ptr,
            cq_ring_ptr,
            sqes_ptr,
            sq_ring_size,
            cq_ring_size,
            sqes_size,
            single_mmap,
            sq_head: offset_ptr(sq_ring_ptr, params.sq_off.head),
            sq_tail: offset_ptr(sq_ring_ptr, params.sq_off.tail),
            sq_ring_mask: offset_ptr(sq_ring_ptr, params.sq_off.ring_mask),
            sq_ring_entries: offset_ptr(sq_ring_ptr, params.sq_off.ring_entries),
            sq_array: offset_ptr(sq_ring_ptr, params.sq_off.array),
            cq_head: offset_ptr(cq_ring_ptr, params.cq_off.head),
            cq_tail: offset_ptr(cq_ring_ptr, params.cq_off.tail),
            cq_ring_mask: offset_ptr(cq_ring_ptr, params.cq_off.ring_mask),
            cqes: offset_ptr(cq_ring_ptr, params.cq_off.cqes),
        })
    }

    pub(crate) fn ring_fd(&self) -> RawFd {
        self.ring_fd
    }

    pub(crate) fn bind_current_thread(&self) {
        CURRENT_SUBMITTER.with(|submitter| submitter.set(self as *const Self));
    }

    pub(crate) fn unbind_current_thread(&self) {
        CURRENT_SUBMITTER.with(|submitter| {
            if ptr::eq(submitter.get(), self) {
                submitter.set(ptr::null());
            }
        });
    }

    pub(crate) fn with_submitter<T>(f: impl FnOnce(&IoUring) -> io::Result<T>) -> io::Result<T> {
        CURRENT_SUBMITTER.with(|submitter| {
            let ptr = submitter.get();
            if !ptr.is_null() {
                let ring = unsafe { &*ptr };
                return f(ring);
            }

            let mut ring = global_submitter()
                .lock()
                .expect("global io_uring submitter should not be poisoned");
            if ring.is_none() {
                *ring = Some(IoUring::new(64)?);
            }

            f(ring
                .as_ref()
                .expect("global submitter ring should initialize"))
        })
    }

    pub(crate) fn submit_timeout(&self, token: u64, deadline: Duration) -> io::Result<()> {
        // SAFETY: `timespec` is stack-allocated and its pointer is passed to
        // the kernel via the SQE's `addr` field. This is safe because
        // `submit_pending` calls `io_uring_enter` synchronously within this
        // stack frame, before `timespec` is dropped. The kernel copies the
        // timespec value during submission of IORING_OP_TIMEOUT, so no
        // lifetime extension beyond this call is required.
        let timespec = duration_to_kernel_timespec(deadline);
        self.push_sqe(|sqe| {
            sqe.opcode = IORING_OP_TIMEOUT;
            sqe.fd = -1;
            sqe.off = 0;
            sqe.user_data = token;
            sqe.addr = (&timespec as *const KernelTimespec) as u64;
            sqe.len = 1;
            sqe.op_flags = IORING_TIMEOUT_ABS;
        })?;
        self.submit_pending().map(|_| ())
    }

    pub(crate) fn submit_timeout_remove(
        &self,
        token_to_remove: u64,
        completion: u64,
    ) -> io::Result<()> {
        self.push_sqe(|sqe| {
            sqe.opcode = IORING_OP_TIMEOUT_REMOVE;
            sqe.fd = -1;
            sqe.flags = IOSQE_CQE_SKIP_SUCCESS;
            sqe.user_data = completion;
            sqe.addr = token_to_remove;
        })?;
        self.submit_pending().map(|_| ())
    }

    pub(crate) fn submit_timeout_update(
        &self,
        token_to_update: u64,
        deadline: Duration,
    ) -> io::Result<()> {
        // SAFETY: Same stack-pointer contract as submit_timeout — the kernel
        // copies the timespec during the synchronous io_uring_enter call made
        // by submit_pending before this function returns.
        let timespec = duration_to_kernel_timespec(deadline);
        self.push_sqe(|sqe| {
            sqe.opcode = IORING_OP_TIMEOUT_REMOVE;
            sqe.fd = -1;
            sqe.off = (&timespec as *const KernelTimespec) as u64;
            sqe.addr = token_to_update;
            sqe.op_flags = IORING_TIMEOUT_UPDATE | IORING_TIMEOUT_ABS;
        })?;
        self.submit_pending().map(|_| ())
    }

    pub(crate) fn submit_msg_ring(
        &self,
        target_ring_fd: RawFd,
        target_user_data: u64,
        value: u32,
        completion: u64,
    ) -> io::Result<()> {
        self.push_sqe(|sqe| {
            sqe.opcode = IORING_OP_MSG_RING;
            sqe.flags = IOSQE_CQE_SKIP_SUCCESS;
            sqe.fd = target_ring_fd;
            sqe.off = target_user_data;
            sqe.addr = IORING_MSG_DATA;
            sqe.len = value;
            sqe.user_data = completion;
        })?;
        self.submit_pending().map(|_| ())
    }

    pub(crate) fn submit_with_token(
        &self,
        token: u64,
        fill: impl FnOnce(&mut IoUringSqe),
    ) -> io::Result<()> {
        self.push_sqe(|sqe| {
            fill(sqe);
            sqe.user_data = token;
        })?;
        self.submit_pending().map(|_| ())
    }

    /// Submits a main SQE linked to an `IORING_OP_LINK_TIMEOUT` SQE.
    ///
    /// The main SQE is submitted with `IOSQE_IO_LINK` set so that if `timeout`
    /// elapses before the main op completes, the kernel cancels the main op
    /// (`-ECANCELED`) and the timeout CQE arrives with `-ETIME`.  If the main op
    /// completes first, the timeout CQE arrives with `-ECANCELED`.
    ///
    /// # Safety of the stack-allocated timespec
    ///
    /// `KernelTimespec` is stack-allocated and its address is passed to the kernel
    /// in the timeout SQE.  This is safe for the same reason as `submit_timeout`:
    /// both SQEs are flushed via the synchronous `io_uring_enter` call inside
    /// `submit_pending` before this function returns.  The kernel copies the
    /// timespec value at submission time; no lifetime extension beyond this call
    /// is required.
    pub(crate) fn submit_linked_with_timeout(
        &self,
        main_token: u64,
        fill: impl FnOnce(&mut IoUringSqe),
        timeout_token: u64,
        timeout: Duration,
    ) -> io::Result<()> {
        let timespec = duration_to_kernel_timespec(timeout);
        self.push_sqe(|sqe| {
            fill(sqe);
            sqe.flags |= IOSQE_IO_LINK;
            sqe.user_data = main_token;
        })?;
        self.push_sqe(|sqe| {
            sqe.opcode = IORING_OP_LINK_TIMEOUT;
            sqe.fd = -1;
            sqe.addr = (&timespec as *const KernelTimespec) as u64;
            sqe.len = 1;
            sqe.user_data = timeout_token;
        })?;
        self.submit_pending().map(|_| ())
    }

    pub(crate) fn drain_completions(&self, mut f: impl FnMut(IoUringCqe)) -> bool {
        let mut head = load_u32(self.cq_head);
        let tail = load_u32(self.cq_tail);
        if head == tail {
            return false;
        }
        let mask = load_u32(self.cq_ring_mask);

        while head != tail {
            let index = (head & mask) as usize;
            let cqe = unsafe { ptr::read_volatile(self.cqes.add(index)) };
            f(cqe);
            head = head.wrapping_add(1);
        }

        store_u32(self.cq_head, head);
        true
    }

    pub(crate) fn wait_for_cqe(&self) -> io::Result<()> {
        loop {
            match self.enter(0, 1, IORING_ENTER_GETEVENTS) {
                Ok(_) => return Ok(()),
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) => return Err(error),
            }
        }
    }

    fn push_sqe(&self, fill: impl FnOnce(&mut IoUringSqe)) -> io::Result<()> {
        let head = load_u32(self.sq_head);
        let tail = load_u32(self.sq_tail);
        let entries = load_u32(self.sq_ring_entries);
        if tail.wrapping_sub(head) >= entries {
            self.submit_pending()?;
            let head = load_u32(self.sq_head);
            let tail = load_u32(self.sq_tail);
            if tail.wrapping_sub(head) >= entries {
                return Err(io::Error::new(
                    io::ErrorKind::WouldBlock,
                    "io_uring submission queue is full",
                ));
            }
        }

        let tail = load_u32(self.sq_tail);
        let mask = load_u32(self.sq_ring_mask);
        let index = (tail & mask) as usize;
        let sqe = unsafe { &mut *self.sqes_ptr.add(index) };
        *sqe = IoUringSqe::default();
        fill(sqe);
        unsafe {
            ptr::write_volatile(self.sq_array.add(index), index as u32);
        }
        compiler_fence(Ordering::Release);
        store_u32(self.sq_tail, tail.wrapping_add(1));
        Ok(())
    }

    fn submit_pending(&self) -> io::Result<u32> {
        let head = load_u32(self.sq_head);
        let tail = load_u32(self.sq_tail);
        let to_submit = tail.wrapping_sub(head);
        if to_submit == 0 {
            return Ok(0);
        }
        self.enter(to_submit, 0, 0)
    }

    fn enter(&self, to_submit: u32, min_complete: u32, flags: u32) -> io::Result<u32> {
        cvt_long(unsafe {
            libc::syscall(
                libc::SYS_io_uring_enter,
                self.ring_fd,
                to_submit as libc::c_uint,
                min_complete as libc::c_uint,
                flags as libc::c_uint,
                ptr::null::<libc::c_void>(),
                0usize,
            )
        })
        .map(|value| value as u32)
    }
}

impl Drop for IoUring {
    fn drop(&mut self) {
        unsafe {
            libc::munmap(self.sqes_ptr.cast(), self.sqes_size);
            if self.single_mmap {
                libc::munmap(
                    self.sq_ring_ptr.cast(),
                    self.sq_ring_size.max(self.cq_ring_size),
                );
            } else {
                libc::munmap(self.sq_ring_ptr.cast(), self.sq_ring_size);
                libc::munmap(self.cq_ring_ptr.cast(), self.cq_ring_size);
            }
            libc::close(self.ring_fd);
        }
    }
}

unsafe impl Send for IoUring {}

fn offset_ptr<T>(base: *mut u8, offset: u32) -> *mut T {
    unsafe { base.add(offset as usize).cast::<T>() }
}

fn mmap_ring(length: usize, fd: RawFd, offset: libc::off_t) -> io::Result<*mut u8> {
    let ptr = unsafe {
        libc::mmap(
            ptr::null_mut(),
            length,
            libc::PROT_READ | libc::PROT_WRITE,
            libc::MAP_SHARED | libc::MAP_POPULATE,
            fd,
            offset,
        )
    };
    if ptr == libc::MAP_FAILED {
        Err(io::Error::last_os_error())
    } else {
        Ok(ptr.cast())
    }
}

fn load_u32(ptr: *const u32) -> u32 {
    let value = unsafe { ptr::read_volatile(ptr) };
    compiler_fence(Ordering::Acquire);
    value
}

fn store_u32(ptr: *mut u32, value: u32) {
    compiler_fence(Ordering::Release);
    unsafe {
        ptr::write_volatile(ptr, value);
    }
}

fn cvt_long(result: libc::c_long) -> io::Result<libc::c_long> {
    if result == -1 {
        Err(io::Error::last_os_error())
    } else {
        Ok(result)
    }
}

fn global_submitter() -> &'static Mutex<Option<IoUring>> {
    GLOBAL_SUBMITTER.get_or_init(|| Mutex::new(None))
}

fn duration_to_kernel_timespec(duration: Duration) -> KernelTimespec {
    KernelTimespec {
        tv_sec: duration.as_secs() as i64,
        tv_nsec: duration.subsec_nanos() as i64,
    }
}
