wireshift-uring 0.1.1

use std::collections::VecDeque;
use std::os::fd::RawFd;
use std::sync::mpsc;

use crossbeam_channel::Receiver;
use io_uring::{opcode, types, IoUring};
use wireshift_core::backend::{BackendCompletion, BackendSubmission, CancellationHandle};
use wireshift_core::{Error, Result};

use crate::buffer_registry::RegisteredBuffers;
use crate::file_registry::FileRegistryTable;
use crate::helpers::finish_descriptor_batch;
use crate::sqe_builder::prepare_submission;
use crate::types::Inflight;
use crate::uring_sys;

#[derive(Debug)]
#[non_exhaustive]
pub enum Message {
    Submit(BackendSubmission),
    Cancel(CancellationHandle),
    RegisterFiles {
        fds: Vec<RawFd>,
        response: mpsc::Sender<Result<()>>,
    },
    UnregisterFiles {
        response: mpsc::Sender<Result<()>>,
    },
    Shutdown,
}

pub const WAKEUP_ID: u64 = 0xC000_0000_0000_0000_u64;
pub const CANCEL_BIT: u64 = 0x8000_0000_0000_0000_u64;

pub struct WakeupFdGuard(pub RawFd);
impl Drop for WakeupFdGuard {
    fn drop(&mut self) {
        if self.0 >= 0 {
            unsafe {
                libc::close(self.0);
            }
        }
    }
}

pub fn manager_loop(
    ring: IoUring,
    queue_depth: u32,
    wakeup_fd: RawFd,
    receiver: Receiver<Message>,
    completion_tx: std::sync::mpsc::Sender<BackendCompletion>,
    registered: Option<RegisteredBuffers>,
) {
    struct KernelTeardownGuard {
        ring: IoUring,
        registered: Option<RegisteredBuffers>,
    }

    impl Drop for KernelTeardownGuard {
        fn drop(&mut self) {
            if self.registered.is_some() {
                let _ = self.ring.submitter().unregister_buffers();
            }
        }
    }

    let mut guard = KernelTeardownGuard { ring, registered };

    let mut pending = VecDeque::new();
    let mut inflight: Vec<Option<Inflight>> = (0..queue_depth as usize).map(|_| None).collect();
    let mut inflight_count = 0_usize;
    let mut shutting_down = false;
    let mut shutdown_cancels_issued = false;
    let mut cancel_queue = VecDeque::new();
    let mut fixed_files = FileRegistryTable::default();
    let _wakeup_guard = WakeupFdGuard(wakeup_fd);
    let mut wakeup_armed = false;

    loop {
        // Arm the eventfd poll trigger
        if !wakeup_armed && inflight_count < queue_depth.saturating_sub(1) as usize {
            let poll_sqe = opcode::PollAdd::new(types::Fd(wakeup_fd), libc::POLLIN as u32)
                .build()
                .user_data(WAKEUP_ID);
            if uring_sys::push_entry(&mut guard.ring, &poll_sqe).is_ok() {
                wakeup_armed = true;
            }
        }

        // Phase 1: Accept new submissions
        if !shutting_down {
            if pending.is_empty() && inflight_count == 0 {
                match receiver.recv() {
                    Ok(message) => pending.push_back(message),
                    Err(_) => {
                        shutting_down = true;
                    }
                }
            }
            while let Ok(message) = receiver.try_recv() {
                pending.push_back(message);
            }
        }

        // Phase 2: Submit pending operations
        if shutting_down {
            while let Some(message) = pending.pop_front() {
                match message {
                    Message::Submit(submission) => {
                        crate::helpers::disarm_descriptor(submission.descriptor);
                        let _ = completion_tx.send(BackendCompletion {
                            id: submission.id,
                            result: Err(Error::canceled(
                                "operation rejected because backend is shutting down",
                                "wait for shutdown to complete before submitting",
                            )),
                        });
                    }
                    Message::RegisterFiles { response, .. } => {
                        let _ = response.send(Err(Error::canceled(
                            "register files rejected because backend is shutting down",
                            "wait for shutdown to complete before registering files",
                        )));
                    }
                    Message::UnregisterFiles { response } => {
                        let _ = response.send(Err(Error::canceled(
                            "unregister files rejected because backend is shutting down",
                            "wait for shutdown to complete before unregistering files",
                        )));
                    }
                    Message::Cancel(_) | Message::Shutdown => {}
                }
            }
        } else {
            while inflight_count < queue_depth.saturating_sub(1) as usize {
                let Some(message) = pending.pop_front() else {
                    break;
                };
                match message {
                    Message::Submit(submission) => {
                        match prepare_submission(&mut guard.ring, submission, &fixed_files) {
                            Ok(state) => {
                                let slot = inflight_slot(state.id, queue_depth as usize);
                                if inflight[slot].is_some() {
                                    let _ = completion_tx.send(BackendCompletion {
                                        id: state.id,
                                        result: Err(Error::completion(
                                            "io_uring inflight slot collision",
                                            "report this as a backend bug",
                                        )),
                                    });
                                } else {
                                    inflight[slot] = Some(state);
                                    inflight_count += 1;
                                }
                            }
                            Err((submission, error)) => {
                                crate::helpers::disarm_descriptor(submission.descriptor);
                                let _ = completion_tx.send(BackendCompletion {
                                    id: submission.id,
                                    result: Err(error),
                                });
                            }
                        }
                    }
                    Message::Cancel(cancellation) => {
                        let cancel_id = cancellation.target;
                        let entry = opcode::AsyncCancel::new(cancel_id)
                            .build()
                            .user_data(cancel_id | CANCEL_BIT);
                        if let Err(error) = uring_sys::push_entry(&mut guard.ring, &entry) {
                            tracing::warn!(
                                target_id = cancel_id,
                                "failed to submit native cancel SQE: {error}"
                            );
                        }
                    }
                    Message::RegisterFiles { fds, response } => {
                        let submitter = guard.ring.submitter();
                        let _ = response.send(fixed_files.register(&submitter, &fds));
                    }
                    Message::UnregisterFiles { response } => {
                        let submitter = guard.ring.submitter();
                        let _ = response.send(fixed_files.unregister(&submitter));
                    }
                    Message::Shutdown => {
                        shutting_down = true;
                        if !shutdown_cancels_issued {
                            for entry in inflight.iter().flatten() {
                                cancel_queue.push_back(entry.id);
                            }
                            shutdown_cancels_issued = true;
                        }
                        break;
                    }
                }
            }
        }

        // Phase 2.5: Force kernel aborts during teardown
        if shutting_down {
            while let Some(cancel_id) = cancel_queue.pop_front() {
                let entry = opcode::AsyncCancel::new(cancel_id)
                    .build()
                    .user_data(cancel_id | CANCEL_BIT);
                if uring_sys::push_entry(&mut guard.ring, &entry).is_err() {
                    cancel_queue.push_front(cancel_id);
                    break;
                }
            }
        }

        if shutting_down && inflight_count == 0 {
            // Phase 3.5: Drain remaining CQEs (cancel completions) before shutdown
            drain_remaining_cqes(
                &mut guard.ring,
                &mut inflight,
                &mut inflight_count,
                &mut wakeup_armed,
                wakeup_fd,
                &completion_tx,
                queue_depth as usize,
            );
            fixed_files.clear();
            return;
        }

        // Phase 4: Wait for and process completions
        if inflight_count != 0 || wakeup_armed {
            if shutting_down {
                const MAX_SHUTDOWN_WAIT_MS: u64 = 5000;
                // Event-driven bounded wait instead of a fixed 10ms poll-sleep
                // spin (Law 7). `submit_with_args` blocks in the kernel until the
                // next completion arrives OR the timeout elapses: a CQE wakes it
                // immediately (sub-millisecond) rather than after the remainder
                // of a 10ms sleep quantum, and no CPU is burned spinning while
                // waiting. `ETIME` means the bounded budget elapsed with the
                // ring still stuck, which we surface loudly.
                guard.ring.completion().sync();
                if guard.ring.completion().is_empty() {
                    let timeout = types::Timespec::new()
                        .sec(MAX_SHUTDOWN_WAIT_MS / 1000)
                        .nsec(((MAX_SHUTDOWN_WAIT_MS % 1000) * 1_000_000) as u32);
                    let args = types::SubmitArgs::new().timespec(&timeout);
                    match guard.ring.submitter().submit_with_args(1, &args) {
                        Ok(_) => {}
                        Err(error) if error.raw_os_error() == Some(libc::ETIME) => {
                            tracing::warn!("shutdown wait timeout: operations may be stuck");
                        }
                        Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {}
                        Err(error) => {
                            tracing::warn!(%error, "shutdown completion wait failed");
                        }
                    }
                }
            } else {
                let _ = guard.ring.submitter().submit_and_wait(1);
            }
            process_completions(
                &mut guard.ring,
                &mut inflight,
                &mut inflight_count,
                &mut wakeup_armed,
                wakeup_fd,
                &completion_tx,
                queue_depth as usize,
            );
        }
    }
}

fn process_completions(
    ring: &mut IoUring,
    inflight: &mut [Option<Inflight>],
    inflight_count: &mut usize,
    wakeup_armed: &mut bool,
    wakeup_fd: RawFd,
    completion_tx: &mpsc::Sender<BackendCompletion>,
    capacity: usize,
) {
    let mut completion_queue = ring.completion();
    completion_queue.sync();
    for cqe in &mut completion_queue {
        let id = cqe.user_data();

        if id == WAKEUP_ID {
            *wakeup_armed = false;
            let mut val_bytes = [0u8; 8];
            // SAFETY: wakeup_fd is a valid eventfd during the manager loop.
            let fd = unsafe { std::os::fd::BorrowedFd::borrow_raw(wakeup_fd) };
            if let Err(error) = rustix::io::read(fd, &mut val_bytes) {
                tracing::error!(%error, "wakeup eventfd read failed");
            }
            continue;
        }

        if (id & CANCEL_BIT) != 0 && (id & !CANCEL_BIT) != 0 {
            let target_id = id & !CANCEL_BIT;
            if cqe.result() < 0 {
                tracing::debug!(
                    target_id,
                    errno = -cqe.result(),
                    "native cancel completed with error"
                );
            }
            continue;
        }

        let slot = inflight_slot(id, capacity);
        let mut finished = false;
        if let Some(Some(entry)) = inflight.get_mut(slot) {
            if entry.id != id {
                let _ = completion_tx.send(BackendCompletion {
                    id,
                    result: Err(Error::completion(
                        "io_uring completion arrived for unknown request id (slot mismatch)",
                        "keep submission ids unique until completions are drained",
                    )),
                });
                continue;
            }
            entry.results.push(cqe.result());
            if entry.results.len() == entry.expected_cqes {
                finished = true;
            }
        } else {
            let _ = completion_tx.send(BackendCompletion {
                id,
                result: Err(Error::completion(
                    "io_uring completion arrived for unknown request id (no inflight entry)",
                    "keep submission ids unique until completions are drained",
                )),
            });
            continue;
        }

        if finished {
            if let Some(entry) = inflight[slot].take() {
                *inflight_count = inflight_count.saturating_sub(1);
                let result = finish_descriptor_batch(entry.descriptor, entry.pinned, entry.results);
                let _ = completion_tx.send(BackendCompletion { id, result });
            }
        } else if inflight[slot].is_none() && (id & CANCEL_BIT) == 0 {
            // Finding 6.2.2: Improve error message for completions arriving after cancel.
            // If the slot is empty, but this isn't a cancel CQE itself, it might be the
            // original operation completing AFTER it was already cancelled and removed.
            tracing::debug!(
                id,
                "io_uring completion arrived for already removed/cancelled request"
            );
        }
    }
}

/// Drain remaining CQEs during shutdown without blocking
fn drain_remaining_cqes(
    ring: &mut IoUring,
    inflight: &mut [Option<Inflight>],
    inflight_count: &mut usize,
    wakeup_armed: &mut bool,
    wakeup_fd: RawFd,
    completion_tx: &mpsc::Sender<BackendCompletion>,
    capacity: usize,
) {
    // Process any remaining completions (cancel CQEs) without blocking
    loop {
        process_completions(
            ring,
            inflight,
            inflight_count,
            wakeup_armed,
            wakeup_fd,
            completion_tx,
            capacity,
        );

        // Check if there are more CQEs without waiting
        ring.completion().sync();
        if ring.completion().is_empty() {
            break;
        }
    }
}

pub fn inflight_slot(id: u64, capacity: usize) -> usize {
    ((id.saturating_sub(1)) % capacity as u64) as usize
}