//! Tokio UDP orchestration around [`ServerCore`](crate::ServerCore).

use std::{future::Future, io, net::SocketAddr, time::Duration};

use socket2::SockRef;
use thiserror::Error;
use tokio::{net::UdpSocket, time::MissedTickBehavior};

use crate::{socket_options::set_reply_traffic_class, ServerConfig, ServerCore, ServerError};

/// Fixed transport receive capacity.
///
/// Standard UDP payloads fit in this buffer. A datagram reported as filling it
/// exactly is dropped conservatively instead of passing potentially truncated
/// bytes to the protocol core. IPv6 jumbograms are not supported.
const RECEIVE_BUFFER_LEN: usize = 65_536;
const MAINTENANCE_INTERVAL: Duration = Duration::from_secs(1);

/// A Tokio UDP listener backed by one deterministic server core.
///
/// Processing is sequential: one task receives a datagram, lets the core
/// process it, and sends the optional reply from the same socket. No
/// per-datagram tasks or locks are involved.
///
/// **That sequential ownership is what makes the reply traffic class correct.**
/// The class is a socket-wide setting, and every reply explicitly applies its
/// own — zero included — immediately before its send, so the reply about to go
/// out is always the one that set it. Nothing else sends from this socket, not
/// even while a send is suspended, so no reply can inherit another session's
/// marking. Adding a concurrent sender would break that and require per-packet
/// control messages instead.
///
/// Bind an explicit local address when reply source-address identity matters.
/// Wildcard binds rely on the kernel's source-address choice; destination
/// packet metadata and per-packet source selection are not yet implemented for
/// multi-homed hosts.
#[derive(Debug)]
pub struct Server {
    socket: UdpSocket,
    core: ServerCore,
    recv_buffer: Vec<u8>,
    /// The listener's own address family, which selects the IPv4 TOS or IPv6
    /// Traffic Class option. `None` only if the bound address could not be
    /// queried at construction, in which case the send path asks again rather
    /// than guessing a family.
    listener_is_ipv4: Option<bool>,
    /// Whether this server has ever successfully applied a nonzero traffic
    /// class to the socket.
    ///
    /// This is not a cache and never elides a call — the class is applied
    /// before every send regardless. It is consulted only when applying one
    /// *fails*, to tell "the socket may be carrying a marking of ours" from
    /// "it cannot be carrying one".
    socket_is_marked: bool,
}

impl Server {
    /// Binds one UDP listener at `addr` and creates its independent session
    /// namespace.
    pub async fn bind(addr: SocketAddr, config: ServerConfig) -> Result<Self, ServerRuntimeError> {
        let socket = UdpSocket::bind(addr)
            .await
            .map_err(|source| ServerRuntimeError::Bind { addr, source })?;
        Ok(Self::from_socket(socket, config))
    }

    /// Wraps an already prepared Tokio UDP socket in a server.
    #[must_use]
    pub fn from_socket(socket: UdpSocket, config: ServerConfig) -> Self {
        let listener_is_ipv4 = socket.local_addr().ok().map(|addr| addr.is_ipv4());
        Self {
            socket,
            core: ServerCore::new(config),
            recv_buffer: vec![0; RECEIVE_BUFFER_LEN],
            listener_is_ipv4,
            socket_is_marked: false,
        }
    }

    /// Applies one reply's raw traffic class to the listener socket, and
    /// reports whether that reply may be sent.
    ///
    /// The class is applied before *every* send, and zero as deliberately as
    /// any other value: a listener serves many sessions from one socket, so
    /// skipping the call for an unmarked reply — an open reply, or a session
    /// that negotiated nothing — would send it under whichever marking the
    /// previous reply left behind.
    ///
    /// When the option cannot be applied, [`may_send_unappliable`] decides. A
    /// marking this server asked for and did not get must never be replaced by
    /// silently sending under the previous one, so those replies are dropped;
    /// but a host that refuses the option outright — some Windows builds do not
    /// support `IP_TOS`, and a few targets have no safe setter at all — must
    /// still be able to run a server, and it can, because a socket this server
    /// has never marked has no marking of ours to clear.
    fn prepare_reply_traffic_class(&mut self, traffic_class: u8) -> bool {
        match self.apply_reply_traffic_class(traffic_class) {
            Ok(()) => {
                self.socket_is_marked = traffic_class != 0;
                true
            }
            Err(_) => may_send_unappliable(traffic_class, self.socket_is_marked),
        }
    }

    fn apply_reply_traffic_class(&self, traffic_class: u8) -> io::Result<()> {
        let is_ipv4 = match self.listener_is_ipv4 {
            Some(is_ipv4) => is_ipv4,
            None => self.socket.local_addr()?.is_ipv4(),
        };
        set_reply_traffic_class(SockRef::from(&self.socket), is_ipv4, traffic_class)
    }

    /// Returns the local endpoint selected for the listener.
    pub fn local_addr(&self) -> Result<SocketAddr, ServerRuntimeError> {
        self.socket
            .local_addr()
            .map_err(|source| ServerRuntimeError::LocalAddr { source })
    }

    /// Serves datagrams until `shutdown` completes.
    ///
    /// Graceful shutdown sends no session-close packets and leaves no hidden
    /// receive task behind, even while a reply send is waiting for socket
    /// writability. Per-packet send failures and short sends drop only that
    /// reply, and so does a failure to apply a marking that reply needed —
    /// preparing one packet is not a reason to stop serving every other
    /// session, and the next reply makes its own independent attempt. Receive
    /// failures and internal core failures terminate the loop with an error.
    ///
    /// A reply whose marking could not be applied is not sent instead: sending
    /// it would put it on the wire under the previous reply's marking. An
    /// *unmarked* reply on a socket this server has never marked is the one
    /// exception, so a host that does not support the option at all still runs
    /// a working server. See [`Server::prepare_reply_traffic_class`].
    pub async fn run<F>(&mut self, shutdown: F) -> Result<(), ServerRuntimeError>
    where
        F: Future<Output = ()>,
    {
        tokio::pin!(shutdown);
        let mut maintenance = tokio::time::interval(MAINTENANCE_INTERVAL);
        maintenance.set_missed_tick_behavior(MissedTickBehavior::Skip);
        maintenance.tick().await;

        loop {
            tokio::select! {
                _ = &mut shutdown => return Ok(()),
                _ = maintenance.tick() => self.core.maintain(),
                received = self.socket.recv_from(&mut self.recv_buffer) => {
                    let (len, peer) = match received {
                        Ok(received) => received,
                        Err(source) if source.kind() == io::ErrorKind::Interrupted => continue,
                        Err(source) => return Err(ServerRuntimeError::Receive { source }),
                    };
                    if len == RECEIVE_BUFFER_LEN {
                        continue;
                    }

                    if let Some(reply) = self.core.handle_datagram(peer, &self.recv_buffer[..len])? {
                        if !self.prepare_reply_traffic_class(reply.traffic_class()) {
                            continue;
                        }
                        let send = self.socket.send_to(reply.bytes(), peer);
                        tokio::pin!(send);
                        let sent = loop {
                            tokio::select! {
                                _ = &mut shutdown => return Ok(()),
                                _ = maintenance.tick() => self.core.maintain(),
                                sent = &mut send => break sent,
                            }
                        };
                        if !matches!(sent, Ok(len) if len == reply.bytes().len()) {
                            continue;
                        }
                    }
                }
            }
        }
    }
}

/// Whether a reply may still be sent after its traffic class could not be
/// applied.
///
/// Only an unmarked reply on a socket this server has never successfully marked
/// qualifies. Both halves matter: a reply that wanted a marking must not go out
/// without one, and once this server has put a nonzero class on the socket, a
/// failure to clear it means the socket may still be carrying it.
///
/// A socket handed to [`Server::from_socket`] pre-marked by its creator is
/// outside what this can know; explicit marking is the server's own to manage
/// from a fresh listener.
fn may_send_unappliable(traffic_class: u8, socket_is_marked: bool) -> bool {
    traffic_class == 0 && !socket_is_marked
}

/// Failure to create or run a Tokio UDP server.
#[derive(Debug, Error)]
pub enum ServerRuntimeError {
    /// Binding the requested listener failed.
    #[error("could not bind UDP listener at {addr}: {source}")]
    Bind {
        addr: SocketAddr,
        #[source]
        source: io::Error,
    },
    /// Querying the listener's selected local endpoint failed.
    #[error("could not query UDP listener address: {source}")]
    LocalAddr {
        #[source]
        source: io::Error,
    },
    /// Receiving from the listener failed irrecoverably.
    #[error("UDP receive failed: {source}")]
    Receive {
        #[source]
        source: io::Error,
    },
    /// The deterministic protocol/session core failed internally.
    #[error(transparent)]
    Core(#[from] ServerError),
}

#[cfg(test)]
mod tests {
    use std::{net::Ipv4Addr, time::Duration};

    use irtt_proto::{
        decode_echo_reply, decode_open_reply, encode_request, Clock, Params, ReceivedStats,
        RequestToEncode, StampAt,
    };
    use tokio::{net::UdpSocket, sync::oneshot, time::timeout};

    use super::*;

    #[tokio::test(flavor = "current_thread")]
    async fn bind_reports_local_address_and_shutdown_stops_run() {
        let requested = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
        let mut server = Server::bind(requested, ServerConfig::default())
            .await
            .unwrap();
        let bound = server.local_addr().unwrap();
        assert_eq!(bound.ip(), requested.ip());
        assert_ne!(bound.port(), 0);

        let (shutdown_tx, shutdown_rx) = oneshot::channel();
        shutdown_tx.send(()).unwrap();
        server
            .run(async {
                let _ = shutdown_rx.await;
            })
            .await
            .unwrap();
    }

    #[tokio::test(flavor = "current_thread")]
    async fn ipv4_open_echo_close_crosses_a_real_udp_socket() {
        exercise_udp_path(SocketAddr::from((Ipv4Addr::LOCALHOST, 0))).await;
    }

    #[tokio::test(flavor = "current_thread")]
    async fn ipv6_open_echo_close_when_loopback_is_available() {
        let addr = "[::1]:0".parse().unwrap();
        let server = match Server::bind(addr, unthrottled()).await {
            Ok(server) => server,
            Err(ServerRuntimeError::Bind { source, .. })
                if matches!(
                    source.kind(),
                    io::ErrorKind::AddrNotAvailable | io::ErrorKind::Unsupported
                ) =>
            {
                return;
            }
            Err(error) => panic!("unexpected IPv6 loopback bind failure: {error}"),
        };
        exercise_bound_server(server).await;
    }

    async fn exercise_udp_path(addr: SocketAddr) {
        let server = Server::bind(addr, unthrottled()).await.unwrap();
        exercise_bound_server(server).await;
    }

    async fn exercise_bound_server(mut server: Server) {
        let server_addr = server.local_addr().unwrap();
        let client_bind = match server_addr {
            SocketAddr::V4(_) => "127.0.0.1:0",
            SocketAddr::V6(_) => "[::1]:0",
        };
        let client = UdpSocket::bind(client_bind).await.unwrap();
        let (shutdown_tx, shutdown_rx) = oneshot::channel();
        let server_task = tokio::spawn(async move {
            server
                .run(async {
                    let _ = shutdown_rx.await;
                })
                .await
        });

        let requested = Params {
            protocol_version: 1,
            duration_ns: 1_000_000_000,
            interval_ns: 100_000_000,
            received_stats: ReceivedStats::Both,
            stamp_at: StampAt::Both,
            clock: Clock::Both,
            ..Params::default()
        };
        let open = encode_request(
            RequestToEncode::Open {
                params: &requested,
                no_test: false,
            },
            None,
        )
        .unwrap();
        client.send_to(&open, server_addr).await.unwrap();

        let mut buffer = vec![0; RECEIVE_BUFFER_LEN];
        let (len, reply_source) = timeout(Duration::from_secs(1), client.recv_from(&mut buffer))
            .await
            .expect("OPEN reply timeout")
            .unwrap();
        assert_eq!(reply_source, server_addr);
        let open_reply = decode_open_reply(&buffer[..len], None).unwrap();
        assert_ne!(open_reply.token, 0);

        let echo = encode_request(
            RequestToEncode::Echo {
                token: open_reply.token,
                sequence: 0,
                params: &open_reply.params,
                payload: &[],
            },
            None,
        )
        .unwrap();
        client.send_to(&echo, server_addr).await.unwrap();
        let (len, reply_source) = timeout(Duration::from_secs(1), client.recv_from(&mut buffer))
            .await
            .expect("ECHO reply timeout")
            .unwrap();
        assert_eq!(reply_source, server_addr);
        let echo_reply = decode_echo_reply(&buffer[..len], &open_reply.params, None).unwrap();
        assert_eq!(echo_reply.token, open_reply.token);
        assert_eq!(echo_reply.sequence, 0);
        assert_eq!(echo_reply.recv_count, Some(1));
        assert_eq!(echo_reply.recv_window, Some(1));

        let close = encode_request(
            RequestToEncode::Close {
                token: open_reply.token,
            },
            None,
        )
        .unwrap();
        client.send_to(&close, server_addr).await.unwrap();
        client.send_to(&echo, server_addr).await.unwrap();
        assert!(
            timeout(Duration::from_millis(100), client.recv_from(&mut buffer))
                .await
                .is_err()
        );

        shutdown_tx.send(()).unwrap();
        server_task.await.unwrap().unwrap();
    }

    fn unthrottled() -> ServerConfig {
        ServerConfig::default().with_min_send_interval(Duration::ZERO)
    }

    /// The rule that decides what an unappliable traffic class means, which no
    /// normal interface can reach: it needs a host that refuses the socket
    /// option, and the hosts these tests run on do not.
    #[test]
    fn only_an_unmarked_reply_on_a_never_marked_socket_survives_an_apply_failure() {
        for (traffic_class, socket_is_marked, sendable, why) in [
            (
                0,
                false,
                true,
                "no marking wanted, and none of ours to clear",
            ),
            (0, true, false, "our own marking may still be on the socket"),
            (
                0xb8,
                false,
                false,
                "a marking that was wanted but not applied",
            ),
            (0xb8, true, false, "and likewise over a marking already set"),
        ] {
            assert_eq!(
                may_send_unappliable(traffic_class, socket_is_marked),
                sendable,
                "{why}"
            );
        }
    }
}