//! Linux ancillary receive and transmit metadata.
//!
//! Receive kernel timing uses `SO_TIMESTAMPING` with `RX_SOFTWARE` +
//! `SOFTWARE` (reported via `SCM_TIMESTAMPING`), not the older
//! `SO_TIMESTAMPNS`/`SCM_TIMESTAMPNS`. Linux documents fallback quirks when
//! `SO_TIMESTAMP[NS]` and `SO_TIMESTAMPING` are both enabled on one socket,
//! so this establishes one unambiguous software-timestamp facility. `ReceiveMeta`'s
//! semantics — an optional observed kernel receive wall time, used for
//! downstream one-way delay only when plausible relative to the userspace
//! receive sample — are unchanged by any of this module's TX-side additions.
//!
//! After a successful Open, the client adapter best-effort upgrades the
//! socket from [`RX_TIMESTAMPING_FLAGS`] to [`error_queue::TX_TIMESTAMPING_FLAGS`]
//! via [`try_enable_tx_timestamping`]. When that succeeds, the automatic
//! `SOF_TIMESTAMPING_OPT_ID` counter the kernel assigns to each successfully
//! submitted datagram is used directly as the probe's wire sequence number
//! (see the client crate's `AGENTS.md` for the correlation invariant this
//! relies on). [`drain_tx_timestamps`] performs a small bounded, nonblocking
//! read of `MSG_ERRQUEUE` so the adapter can opportunistically collect those
//! timestamps without ever waiting for one. The observed kernel TX wall time
//! is retained as metadata only; it is not consumed by any measurement in
//! this change.

use std::{
    io::{self, IoSliceMut},
    net::{SocketAddr, UdpSocket},
    os::fd::{AsFd, AsRawFd, RawFd},
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use nix::sys::{
    socket::{
        cmsg_space, recvmsg, setsockopt, sockopt, ControlMessageOwned, MsgFlags, RecvMsg,
        TimestampingFlag, Timestamps,
    },
    time::TimeSpec,
};

use crate::{metadata::ReceiveMeta, receive::ReceivedDatagram, timing::ClientTimestamp};

pub(crate) mod error_queue;

use error_queue::ErrorQueueRecord;

/// Starvation guard on one opportunistic `MSG_ERRQUEUE` drain.
///
/// Each successfully submitted, TX-timestamped datagram generates at most
/// one record, and drains happen at several natural choke points (after a
/// probe send, after a normal receive, before timeout expiry), so this is
/// not a queue-capacity promise: it just bounds the work one drain call can
/// do. A handful of in-flight probes' worth of backlog is the realistic
/// case between choke points, so a small constant comfortably covers it
/// without any unbounded loop.
pub(crate) const MAX_TX_TIMESTAMP_RECORDS_PER_DRAIN: usize = 32;

/// Best-effort upgrade from RX-only to RX+TX `SO_TIMESTAMPING` flags.
///
/// Returns `true` when the upgrade succeeded and the caller may now expect
/// `MSG_ERRQUEUE` TX timestamp records for subsequent sends. On failure, the
/// original [`RX_TIMESTAMPING_FLAGS`] are best-effort restored so a failed
/// upgrade cannot leave the socket without the RX timestamping this client
/// already depends on; that restore's own failure is not observable here
/// (both leave the caller without TX capability, which this return value
/// already reports).
pub(crate) fn try_enable_tx_timestamping<S: AsFd>(socket: &S) -> bool {
    if setsockopt(
        socket,
        sockopt::Timestamping,
        &error_queue::TX_TIMESTAMPING_FLAGS,
    )
    .is_ok()
    {
        return true;
    }
    let _ = setsockopt(socket, sockopt::Timestamping, &RX_TIMESTAMPING_FLAGS);
    false
}

/// Bounded, nonblocking drain of `fd`'s `MSG_ERRQUEUE`.
///
/// Reads and classifies up to [`MAX_TX_TIMESTAMP_RECORDS_PER_DRAIN`] records,
/// one at a time with no intermediate allocation. Each usable TX timestamp
/// completion is reported to `on_timestamp` immediately. Malformed or
/// unsupported timestamp notifications and unrelated records are silently
/// dropped: timestamp metadata is optional and never fails a probe. A
/// genuine non-timestamp socket/network error stops the drain and is
/// returned to the caller, which decides how to surface it.
pub(crate) fn drain_tx_timestamps<S: AsRawFd>(
    socket: &S,
    mut on_timestamp: impl FnMut(u32, SystemTime),
) -> io::Result<()> {
    let fd = socket.as_raw_fd();
    for _ in 0..MAX_TX_TIMESTAMP_RECORDS_PER_DRAIN {
        match error_queue::try_recv_error_queue_record(fd)? {
            None => return Ok(()),
            Some(ErrorQueueRecord::TxTimestamp { id, timestamp }) => on_timestamp(id, timestamp),
            Some(ErrorQueueRecord::MalformedOrUnsupportedTimestamp | ErrorQueueRecord::Ignored) => {
            }
            Some(ErrorQueueRecord::SocketError { errno, .. }) => {
                return Err(io::Error::from_raw_os_error(errno));
            }
        }
    }
    Ok(())
}

/// `SO_TIMESTAMPING` flags requested on every production socket at
/// connect time, before any TX upgrade is attempted.
///
/// `SOF_TIMESTAMPING_RX_SOFTWARE` asks the kernel to generate a software
/// receive timestamp for each datagram; `SOF_TIMESTAMPING_SOFTWARE` asks it
/// to report that timestamp back via `SCM_TIMESTAMPING`. Mixing this API with
/// the older `SO_TIMESTAMP`/`SO_TIMESTAMPNS` on the same socket has
/// documented fallback quirks, so the client uses `SO_TIMESTAMPING`
/// exclusively for receive timing. Hardware and TX-side flags are
/// deliberately excluded here: a fresh socket is receive-only until (and
/// unless) [`try_enable_tx_timestamping`] later upgrades it.
const RX_TIMESTAMPING_FLAGS: TimestampingFlag = TimestampingFlag::SOF_TIMESTAMPING_RX_SOFTWARE
    .union(TimestampingFlag::SOF_TIMESTAMPING_SOFTWARE);

/// Control buffer capacity for a normal receive: one `SCM_TIMESTAMPING`
/// record (three `TimeSpec`s) plus one traffic-class record. Only one of
/// `IP_TOS` (`u8`) or `IPV6_TCLASS` (`i32`) is ever requested on a given
/// socket, but `cmsg_space` already rounds both up to the same aligned
/// space, so sizing on the larger of the two costs nothing and covers
/// either family.
const CONTROL_LEN: usize = cmsg_space::<Timestamps>() + cmsg_space::<i32>();

pub(crate) fn configure_receive_metadata(socket: &UdpSocket, remote: SocketAddr) -> io::Result<()> {
    setsockopt(socket, sockopt::Timestamping, &RX_TIMESTAMPING_FLAGS)?;
    let socket = socket2::SockRef::from(socket);
    if remote.is_ipv4() {
        socket.set_recv_tos_v4(true)
    } else {
        socket.set_recv_tclass_v6(true)
    }
}

pub(crate) fn recv_datagram(
    socket: &UdpSocket,
    buf: &mut [u8],
) -> Result<ReceivedDatagram, io::Error> {
    recv_datagram_fd(socket.as_raw_fd(), buf, MsgFlags::empty())
}

#[cfg(feature = "tokio")]
pub(crate) fn try_recv_tokio_datagram(
    socket: &tokio::net::UdpSocket,
    buf: &mut [u8],
) -> Result<ReceivedDatagram, io::Error> {
    socket.try_io(tokio::io::Interest::READABLE, || {
        recv_datagram_fd(socket.as_raw_fd(), buf, MsgFlags::MSG_DONTWAIT)
    })
}

fn recv_datagram_fd(
    socket_fd: RawFd,
    buf: &mut [u8],
    flags: MsgFlags,
) -> Result<ReceivedDatagram, io::Error> {
    let mut control = ControlBuffer::new();
    let mut iov = [IoSliceMut::new(buf)];

    // The socket is connected, so the source address is not needed; `()` skips
    // copying it out.
    let msg = recvmsg::<()>(socket_fd, &mut iov, Some(control.as_mut_slice()), flags)?;
    let received_at = ClientTimestamp::now();

    Ok(ReceivedDatagram {
        len: msg.bytes,
        received_at,
        meta: receive_meta(&msg),
    })
}

fn receive_meta<S>(msg: &RecvMsg<'_, '_, S>) -> ReceiveMeta {
    let mut meta = ReceiveMeta::default();
    // A control buffer the kernel had to truncate (`MSG_CTRUNC`) cannot be
    // walked reliably, so `cmsgs` refuses to parse it. The datagram itself is
    // still valid; report it without ancillary metadata.
    let Ok(cmsgs) = msg.cmsgs() else {
        return meta;
    };

    for cmsg in cmsgs {
        match cmsg {
            ControlMessageOwned::ScmTimestampsns(timestamps) => {
                meta.kernel_rx_timestamp = system_time_from_timespec(timestamps.system);
            }
            ControlMessageOwned::Ipv4Tos(tos) => {
                meta.traffic_class = Some(tos);
            }
            ControlMessageOwned::Ipv6TClass(traffic_class) => {
                meta.traffic_class = u8::try_from(traffic_class).ok();
            }
            _ => {}
        }
    }
    meta
}

fn system_time_from_timespec(timespec: TimeSpec) -> Option<SystemTime> {
    let seconds = u64::try_from(timespec.tv_sec()).ok()?;
    let nanos = u32::try_from(timespec.tv_nsec()).ok()?;
    if nanos >= 1_000_000_000 {
        return None;
    }
    UNIX_EPOCH.checked_add(Duration::new(seconds, nanos))
}

/// Reusable per-receive control buffer. `recvmsg` control data must satisfy
/// `cmsghdr` alignment, which 8-byte alignment covers on supported targets.
#[repr(align(8))]
struct ControlBuffer([u8; CONTROL_LEN]);

impl ControlBuffer {
    fn new() -> Self {
        Self([0; CONTROL_LEN])
    }

    fn as_mut_slice(&mut self) -> &mut [u8] {
        &mut self.0
    }
}

#[cfg(test)]
mod tests {
    use std::{
        io,
        net::{SocketAddr, UdpSocket},
        time::{Duration, UNIX_EPOCH},
    };

    use nix::sys::time::TimeSpec;

    use crate::{
        event::PacketMeta,
        receive::{configure_receive_metadata, recv_datagram},
        socket_options::apply_traffic_class_to_socket,
        timing::ClientTimestamp,
    };

    fn connected_ipv4_loopback_pair() -> (UdpSocket, UdpSocket) {
        let a = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap();
        let b = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap();
        a.connect(b.local_addr().unwrap()).unwrap();
        b.connect(a.local_addr().unwrap()).unwrap();
        a.set_read_timeout(Some(Duration::from_millis(200)))
            .unwrap();
        b.set_read_timeout(Some(Duration::from_millis(200)))
            .unwrap();
        (a, b)
    }

    fn connected_ipv6_loopback_pair() -> Option<(UdpSocket, UdpSocket)> {
        let bind_addr = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 0));
        let a = match UdpSocket::bind(bind_addr) {
            Ok(socket) => socket,
            Err(error) if is_unavailable_ipv6_loopback(&error) => {
                eprintln!(
                    "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}"
                );
                return None;
            }
            Err(error) => panic!("{error}"),
        };
        let b = match UdpSocket::bind(bind_addr) {
            Ok(socket) => socket,
            Err(error) if is_unavailable_ipv6_loopback(&error) => {
                eprintln!(
                    "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}"
                );
                return None;
            }
            Err(error) => panic!("{error}"),
        };
        match a.connect(b.local_addr().unwrap()) {
            Ok(()) => {}
            Err(error) if is_unavailable_ipv6_loopback(&error) => {
                eprintln!(
                    "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}"
                );
                return None;
            }
            Err(error) => panic!("{error}"),
        }
        match b.connect(a.local_addr().unwrap()) {
            Ok(()) => {}
            Err(error) if is_unavailable_ipv6_loopback(&error) => {
                eprintln!(
                    "skipping IPv6 ancillary receive test: IPv6 loopback unavailable: {error}"
                );
                return None;
            }
            Err(error) => panic!("{error}"),
        }
        a.set_read_timeout(Some(Duration::from_millis(200)))
            .unwrap();
        b.set_read_timeout(Some(Duration::from_millis(200)))
            .unwrap();
        Some((a, b))
    }

    #[test]
    fn recvmsg_receive_returns_length_payload_and_timestamp() {
        let (sender, receiver) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        sender.send(b"hello").unwrap();

        let before = ClientTimestamp::now();
        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        let after = ClientTimestamp::now();

        assert_eq!(datagram.len, 5);
        assert_eq!(&buf[..datagram.len], b"hello");
        assert!(datagram.received_at.mono >= before.mono);
        assert!(datagram.received_at.mono <= after.mono);
    }

    #[cfg(feature = "tokio")]
    #[test]
    fn tokio_recvmsg_would_block_clears_readiness() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_io()
            .enable_time()
            .build()
            .unwrap();

        runtime.block_on(async {
            let (sender, receiver) = connected_ipv4_loopback_pair();
            configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
            receiver.set_nonblocking(true).unwrap();
            let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap();
            let mut buf = [0_u8; 16];

            sender.send(b"first").unwrap();
            tokio::time::timeout(Duration::from_secs(1), receiver.readable())
                .await
                .expect("initial readability timed out")
                .unwrap();

            let first = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap();
            assert_eq!(&buf[..first.len], b"first");

            let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err();
            assert_eq!(error.kind(), io::ErrorKind::WouldBlock);

            assert!(
                tokio::time::timeout(Duration::from_millis(20), receiver.readable())
                    .await
                    .is_err(),
                "readable completed from stale readiness after recvmsg WouldBlock"
            );

            sender.send(b"second").unwrap();
            tokio::time::timeout(Duration::from_secs(1), receiver.readable())
                .await
                .expect("second readability timed out")
                .unwrap();

            let second = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap();
            assert_eq!(&buf[..second.len], b"second");
        });
    }

    #[test]
    fn kernel_rx_timestamp_metadata_is_observed_when_kernel_provides_it() {
        let (sender, receiver) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        sender.send(b"stamp").unwrap();

        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        assert_eq!(datagram.len, 5);
        assert_eq!(&buf[..datagram.len], b"stamp");
        let Some(timestamp) = datagram.meta.kernel_rx_timestamp else {
            eprintln!(
                "skipping kernel timestamp assertion: kernel did not provide SCM_TIMESTAMPING"
            );
            return;
        };

        let duration = timestamp.duration_since(UNIX_EPOCH).unwrap();
        assert!(duration.as_nanos() > 0);
    }

    #[test]
    fn kernel_rx_timestamp_and_traffic_class_coexist() {
        let (sender, receiver) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        apply_traffic_class_to_socket(&sender, receiver.local_addr().unwrap(), 184).unwrap();
        sender.send(b"both").unwrap();

        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        assert_eq!(&buf[..datagram.len], b"both");

        // Both cmsgs are requested on this socket; the parser must not
        // depend on the order the kernel emits them in to observe either.
        if datagram.meta.kernel_rx_timestamp.is_none() {
            eprintln!("skipping combined assertion: kernel did not provide SCM_TIMESTAMPING");
            return;
        }
        if datagram.meta.traffic_class.is_none() {
            eprintln!("skipping combined assertion: kernel did not provide IP_TOS");
            return;
        }
        assert!(datagram.meta.kernel_rx_timestamp.is_some());
        assert_eq!(datagram.meta.traffic_class.unwrap() & 0xfc, 184);
    }

    #[test]
    fn ipv4_traffic_class_metadata_is_observed_when_kernel_provides_it() {
        let (sender, receiver) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        apply_traffic_class_to_socket(&sender, receiver.local_addr().unwrap(), 184).unwrap();
        sender.send(b"dscp").unwrap();

        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        let Some(traffic_class) = datagram.meta.traffic_class else {
            eprintln!("skipping IPv4 ancillary metadata assertion: kernel did not provide IP_TOS");
            return;
        };

        let packet_meta = PacketMeta::from(datagram.meta);
        assert_eq!(traffic_class & 0xfc, 184);
        assert_eq!(packet_meta.dscp, Some(46));
        assert_eq!(packet_meta.ecn, Some(0));
    }

    #[test]
    fn ipv6_traffic_class_metadata_is_observed_when_kernel_provides_it() {
        let Some((sender, receiver)) = connected_ipv6_loopback_pair() else {
            return;
        };
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        apply_traffic_class_to_socket(&sender, receiver.local_addr().unwrap(), 184).unwrap();
        sender.send(b"dscp").unwrap();

        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        let Some(traffic_class) = datagram.meta.traffic_class else {
            eprintln!(
                "skipping IPv6 ancillary metadata assertion: kernel did not provide IPV6_TCLASS"
            );
            return;
        };

        let packet_meta = PacketMeta::from(datagram.meta);
        assert_eq!(traffic_class & 0xfc, 184);
        assert_eq!(packet_meta.dscp, Some(46));
        assert_eq!(packet_meta.ecn, Some(0));
    }

    #[test]
    fn observed_zero_traffic_class_preserves_some_zero() {
        let (sender, receiver) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
        sender.send(b"zero").unwrap();

        let mut buf = [0_u8; 16];
        let datagram = recv_datagram(&receiver, &mut buf).unwrap();
        let Some(traffic_class) = datagram.meta.traffic_class else {
            eprintln!("skipping observed-zero assertion: kernel did not provide IP_TOS");
            return;
        };

        let packet_meta = PacketMeta::from(datagram.meta);
        assert_eq!(traffic_class, 0);
        assert_eq!(packet_meta.traffic_class, Some(0));
        assert_eq!(packet_meta.dscp, Some(0));
        assert_eq!(packet_meta.ecn, Some(0));
    }

    #[test]
    fn timespec_conversion_accepts_valid_unix_timestamp() {
        let timestamp = super::system_time_from_timespec(TimeSpec::new(1, 2)).unwrap();

        assert_eq!(
            timestamp.duration_since(UNIX_EPOCH).unwrap(),
            Duration::new(1, 2)
        );
    }

    #[test]
    fn timespec_conversion_rejects_negative_or_invalid_values() {
        assert_eq!(super::system_time_from_timespec(TimeSpec::new(-1, 0)), None);
        assert_eq!(super::system_time_from_timespec(TimeSpec::new(0, -1)), None);
        assert_eq!(
            super::system_time_from_timespec(TimeSpec::new(0, 1_000_000_000)),
            None
        );
    }

    fn is_unavailable_ipv6_loopback(error: &io::Error) -> bool {
        matches!(
            error.kind(),
            io::ErrorKind::AddrNotAvailable | io::ErrorKind::Unsupported
        )
    }

    #[test]
    fn try_enable_tx_timestamping_then_drain_reports_ids_through_the_adapter_api() {
        let (socket, _peer) = connected_ipv4_loopback_pair();
        // Mirrors production: RX-only first, then a best-effort upgrade.
        configure_receive_metadata(&socket, socket.peer_addr().unwrap()).unwrap();

        if !super::try_enable_tx_timestamping(&socket) {
            eprintln!(
                "skipping adapter-level TX timestamping smoke test: kernel/container denied \
                 SO_TIMESTAMPING TX flags"
            );
            return;
        }

        if let Err(error) = socket.send(b"one") {
            eprintln!("skipping adapter-level TX timestamping smoke test: send failed: {error}");
            return;
        }

        let deadline = std::time::Instant::now() + Duration::from_secs(2);
        let mut observed = None;
        while observed.is_none() && std::time::Instant::now() < deadline {
            super::drain_tx_timestamps(&socket, |id, timestamp| {
                observed = Some((id, timestamp));
            })
            .unwrap();
            std::thread::yield_now();
        }

        let Some((id, timestamp)) = observed else {
            eprintln!(
                "skipping adapter-level TX timestamping smoke test: no TX timestamp record \
                 arrived within the poll deadline"
            );
            return;
        };
        assert_eq!(id, 0);
        assert!(timestamp.duration_since(UNIX_EPOCH).unwrap() > Duration::ZERO);
    }

    #[test]
    fn drain_tx_timestamps_is_a_bounded_nonblocking_no_op_on_an_untouched_socket() {
        let (socket, _peer) = connected_ipv4_loopback_pair();
        configure_receive_metadata(&socket, socket.peer_addr().unwrap()).unwrap();

        let mut calls = 0_usize;
        super::drain_tx_timestamps(&socket, |_, _| calls += 1).unwrap();

        assert_eq!(calls, 0);
    }
}