wireshift-uring 0.1.1

use std::net::SocketAddr;
use std::os::fd::{AsRawFd, FromRawFd};

use io_uring::squeue;
use wireshift_core::op::{CompletionPayload, OpDescriptor};
use wireshift_core::{Error, Result};

use crate::types::InflightData;

/// Returns true if the flags indicate a sequential I/O hint.
pub fn has_sequential_hint(flags: u32) -> bool {
    flags & sequential_hint_bit() != 0
}

/// The bit used to represent a sequential I/O hint in `OpenFlags`.
pub const fn sequential_hint_bit() -> u32 {
    1 << 31
}

/// Applies sequential I/O hints to a file using `posix_fadvise`.
pub fn apply_open_hints(file: &std::fs::File, flags: u32) -> Result<()> {
    if !has_sequential_hint(flags) {
        return Ok(());
    }

    rustix::fs::fadvise(file, 0, None, rustix::fs::Advice::Sequential).map_err(|error| {
        Error::io(
            "posix_fadvise(SEQUENTIAL) failed after native openat",
            std::io::Error::from(error),
            "remove OpenFlags::SEQUENTIAL or ensure the target filesystem supports fadvise",
        )
    })?;

    Ok(())
}

/// Converts a `SocketAddr` into a `sockaddr_storage` and its length.
pub fn sockaddr_to_storage(addr: &SocketAddr) -> (libc::sockaddr_storage, libc::socklen_t) {
    let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
    let len = match addr {
        SocketAddr::V4(v4) => {
            let sockaddr_in = libc::sockaddr_in {
                sin_family: libc::AF_INET as libc::sa_family_t,
                sin_port: v4.port().to_be(),
                sin_addr: libc::in_addr {
                    s_addr: u32::from_ne_bytes(v4.ip().octets()),
                },
                sin_zero: [0; 8],
            };
            unsafe {
                std::ptr::copy_nonoverlapping(
                    std::ptr::from_ref::<libc::sockaddr_in>(&sockaddr_in).cast::<u8>(),
                    std::ptr::from_mut::<libc::sockaddr_storage>(&mut storage).cast::<u8>(),
                    std::mem::size_of::<libc::sockaddr_in>(),
                );
            }
            std::mem::size_of::<libc::sockaddr_in>() as libc::socklen_t
        }
        SocketAddr::V6(v6) => {
            let sockaddr_in6 = libc::sockaddr_in6 {
                sin6_family: libc::AF_INET6 as libc::sa_family_t,
                sin6_port: v6.port().to_be(),
                sin6_flowinfo: v6.flowinfo(),
                sin6_addr: libc::in6_addr {
                    s6_addr: v6.ip().octets(),
                },
                sin6_scope_id: v6.scope_id(),
            };
            unsafe {
                std::ptr::copy_nonoverlapping(
                    std::ptr::from_ref::<libc::sockaddr_in6>(&sockaddr_in6).cast::<u8>(),
                    std::ptr::from_mut::<libc::sockaddr_storage>(&mut storage).cast::<u8>(),
                    std::mem::size_of::<libc::sockaddr_in6>(),
                );
            }
            std::mem::size_of::<libc::sockaddr_in6>() as libc::socklen_t
        }
    };
    (storage, len)
}

/// Builds Linux `open(2)` flags from high-level `wireshift` open options.
#[allow(clippy::cast_possible_wrap)]
pub fn build_open_flags(
    extra_flags: u32,
    read: bool,
    write: bool,
    create: bool,
    truncate: bool,
) -> Result<i32> {
    let mut flags = match (read, write) {
        (true, true) => libc::O_RDWR,
        (true, false) => libc::O_RDONLY,
        (false, true) => libc::O_WRONLY,
        (false, false) => {
            return Err(Error::validation(
                "openat must request read and/or write access",
                "enable read_only or create a read-write open request",
            ));
        }
    };
    if create {
        flags |= libc::O_CREAT;
    }
    if truncate {
        flags |= libc::O_TRUNC;
    }
    flags |= (extra_flags & !sequential_hint_bit()) as i32;
    flags |= libc::O_CLOEXEC;
    Ok(flags)
}

/// Determines the appropriate `io_uring` link flags between two operations.
pub fn link_flag_between(current: &OpDescriptor, next: &OpDescriptor) -> squeue::Flags {
    if matches!(current, OpDescriptor::ReadFixed { .. })
        && matches!(next, OpDescriptor::CloseFixed { .. })
    {
        squeue::Flags::IO_HARDLINK
    } else {
        squeue::Flags::IO_LINK
    }
}

/// Returns the number of CQEs we expect for a given descriptor.
pub fn descriptor_expected_cqes(desc: &OpDescriptor) -> usize {
    match desc {
        OpDescriptor::Linked { descriptors } => {
            descriptors.iter().map(descriptor_expected_cqes).sum()
        }
        _ => 1,
    }
}

/// Finishes a batch of completions for a descriptor.

/// Manually closes a File descriptor using libc::close (ignoring EBADF) and forgets it to prevent OwnedFd drop aborts
pub fn disarm_and_close_file(file: std::fs::File) {
    let fd = file.as_raw_fd();
    eprintln!("--- DEBUG: DISARMING FILE FD: {} ---", fd);
    if fd >= 0 {
        unsafe {
            libc::close(fd);
        }
    }
    std::mem::forget(file);
}

/// Manually closes a TcpStream descriptor using libc::close (ignoring EBADF) and forgets it to prevent OwnedFd drop aborts
pub fn disarm_and_close_stream(stream: std::net::TcpStream) {
    let fd = stream.as_raw_fd();
    eprintln!("--- DEBUG: DISARMING STREAM FD: {} ---", fd);
    if fd >= 0 {
        unsafe {
            libc::close(fd);
        }
    }
    std::mem::forget(stream);
}

/// Intentionally leaks any File or TcpStream with negative fds to bypass OwnedFd safety aborts on drop
pub fn disarm_descriptor(descriptor: OpDescriptor) {
    match descriptor {
        OpDescriptor::Read { file, .. }
        | OpDescriptor::Write { file, .. }
        | OpDescriptor::ReadVectored { file, .. }
        | OpDescriptor::ReadGpu { file, .. }
        | OpDescriptor::WriteVectored { file, .. }
        | OpDescriptor::Fsync { file, .. }
        | OpDescriptor::Madvise { file, .. } => {
            disarm_and_close_file(file);
        }
        OpDescriptor::OpenAt { dir, .. }
        | OpDescriptor::OpenAtDirect { dir, .. } => {
            if let Some(dir_file) = dir {
                disarm_and_close_file(dir_file);
            }
        }
        OpDescriptor::Send { stream, .. }
        | OpDescriptor::Recv { stream, .. } => {
            disarm_and_close_stream(stream);
        }
        OpDescriptor::Splice { fd_in, fd_out, .. } => {
            disarm_and_close_file(fd_in);
            disarm_and_close_file(fd_out);
        }
        OpDescriptor::Linked { descriptors } => {
            for desc in descriptors {
                disarm_descriptor(desc);
            }
        }
        _ => {}
    }
}

pub fn finish_descriptor_batch(
    descriptor: OpDescriptor,
    pinned: InflightData,
    mut results: Vec<i32>,
) -> Result<CompletionPayload> {
    if results.iter().any(|&r| r < 0 && r != -libc::ECANCELED) {
        let first_err = results
            .iter()
            .find(|&&r| r < 0 && r != -libc::ECANCELED)
            .copied()
            .unwrap_or(-1);
        disarm_descriptor(descriptor);
        return Err(Error::io(
            "io_uring operation failed natively in the kernel",
            std::io::Error::from_raw_os_error(-first_err),
            "inspect the descriptor state and submitted offsets before retrying",
        ));
    }
    if results.contains(&-libc::ECANCELED) {
        disarm_descriptor(descriptor);
        return Err(Error::canceled(
            "request was canceled",
            "avoid canceling the request if it still needs to complete",
        ));
    }

    match descriptor {
        OpDescriptor::Linked { descriptors } => {
            let mut pinned = pinned;
            let pinned_list = match &mut pinned {
                InflightData::Linked(pinned_list) => std::mem::take(pinned_list),
                _ => {
                    return Err(Error::completion(
                        "linked operations must have linked pinned data",
                        "internal error",
                    ))
                }
            };
            let mut pinned_list = pinned_list;
            let mut payloads = Vec::with_capacity(descriptors.len());
            for desc in descriptors {
                let count = descriptor_expected_cqes(&desc);
                let inner_results: Vec<i32> = results.drain(0..count).collect();
                let inner_pinned = if count == 1 {
                    pinned_list.remove(0)
                } else {
                    InflightData::Linked(pinned_list.drain(0..count).collect())
                };
                payloads.push(finish_descriptor_batch(desc, inner_pinned, inner_results)?);
            }
            Ok(CompletionPayload::LinkedChain(payloads))
        }
        other => finish_descriptor(other, pinned, results[0]),
    }
}

/// Finishes a single descriptor completion.
pub fn finish_descriptor(
    descriptor: OpDescriptor,
    mut pinned: InflightData,
    result: i32,
) -> Result<CompletionPayload> {
    match descriptor {
        OpDescriptor::Read { file, buffer, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring read returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            let completed = buffer.into_completed(bytes)?;
            Ok(CompletionPayload::Read {
                buffer: completed,
                bytes,
            })
        }

        OpDescriptor::ReadGpu { file, buffer, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring gpu staged read returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            Ok(CompletionPayload::GpuRead { buffer, bytes })
        }

        OpDescriptor::Write { file, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring write returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            Ok(CompletionPayload::Bytes(bytes))
        }

        OpDescriptor::ReadVectored { file, buffers, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            let total = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring readv returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            let mut remaining = total;
            let mut completed = Vec::with_capacity(buffers.len());
            for buf in buffers {
                let chunk = remaining.min(buf.capacity());
                completed.push(buf.into_completed(chunk)?);
                remaining = remaining.saturating_sub(chunk);
            }
            Ok(CompletionPayload::ReadVectored {
                buffers: completed,
                bytes: total,
            })
        }

        OpDescriptor::WriteVectored { file, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring writev returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            Ok(CompletionPayload::Bytes(bytes))
        }

        OpDescriptor::OpenAt { dir, flags, .. } => {
            if let Some(dir_file) = dir {
                if dir_file.as_raw_fd() < 0 {
                    std::mem::forget(dir_file);
                }
            }
            let fd = result;
            // SAFETY: The kernel returned a new, exclusively owned file descriptor. We assume ownership via `File::from_raw_fd`.
            let file = unsafe { std::fs::File::from_raw_fd(fd) };
            apply_open_hints(&file, flags)?;
            Ok(CompletionPayload::File(file))
        }

        OpDescriptor::Statx { .. } => {
            let InflightData::Statx { statx_buf, .. } = &mut pinned else {
                return Err(Error::completion(
                    "statx completion is missing its pinned buffer",
                    "this is an internal consistency error  -  file a bug report",
                ));
            };
            let mode = u32::from(statx_buf.stx_mode);
            let is_file = (mode & libc::S_IFMT) == libc::S_IFREG;
            let is_dir = (mode & libc::S_IFMT) == libc::S_IFDIR;
            Ok(CompletionPayload::Metadata {
                size: statx_buf.stx_size,
                is_file,
                is_dir,
            })
        }

        OpDescriptor::Connect { .. } => {
            let socket_fd = match &mut pinned {
                InflightData::Sockaddr { socket_fd, .. } => {
                    let fd = *socket_fd;
                    *socket_fd = -1;
                    fd
                }
                _ => {
                    return Err(Error::completion(
                        "connect completion is missing its pinned socket fd",
                        "this is an internal consistency error  -  file a bug report",
                    ));
                }
            };
            let stream = unsafe { std::net::TcpStream::from_raw_fd(socket_fd) };
            Ok(CompletionPayload::Stream(stream))
        }

        OpDescriptor::Accept { .. } => {
            let fd = result;
            let stream = unsafe { std::net::TcpStream::from_raw_fd(fd) };
            Ok(CompletionPayload::Stream(stream))
        }

        OpDescriptor::Send { .. } => {
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring send returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            Ok(CompletionPayload::Bytes(bytes))
        }

        OpDescriptor::Recv { stream, buffer } => {
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring recv returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            let completed = buffer.into_completed(bytes)?;
            Ok(CompletionPayload::Recv {
                stream,
                buffer: completed,
                bytes,
            })
        }

        OpDescriptor::ReadFixed { buffer, .. } => {
            let bytes = usize::try_from(result).map_err(|_| {
                Error::completion(
                    "io_uring read returned an invalid negative byte count",
                    "inspect the kernel error path",
                )
            })?;
            let completed = buffer.into_completed(bytes)?;
            Ok(CompletionPayload::Read {
                buffer: completed,
                bytes,
            })
        }

        OpDescriptor::OpenAtDirect { dir, .. } => {
            if let Some(dir_file) = dir {
                if dir_file.as_raw_fd() < 0 {
                    std::mem::forget(dir_file);
                }
            }
            Ok(CompletionPayload::Unit)
        }
        OpDescriptor::Fsync { file } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            Ok(CompletionPayload::Unit)
        }
        OpDescriptor::Splice { fd_in, fd_out, .. } => {
            if fd_in.as_raw_fd() < 0 {
                std::mem::forget(fd_in);
            }
            if fd_out.as_raw_fd() < 0 {
                std::mem::forget(fd_out);
            }
            Ok(CompletionPayload::Unit)
        }
        OpDescriptor::Madvise { file, .. } => {
            if file.as_raw_fd() < 0 {
                std::mem::forget(file);
            }
            Ok(CompletionPayload::Unit)
        }
        OpDescriptor::CloseFixed { .. }
        | OpDescriptor::Nop
        | OpDescriptor::Cancel { .. } => Ok(CompletionPayload::Unit),

        OpDescriptor::Linked { .. } => Err(Error::Unsupported {
            message: "unsupported linked chain reached the native io_uring completion path".into(),
            fix: "submit only supported operations when the native backend is active".into(),
        }),
        _ => Ok(CompletionPayload::Unit),
    }
}