//! UDP relay for a Wolfenstein: Enemy Territory server reachable only through a tunnel.
//!
//! - Proxies every client (players, server browsers, masters) to the ET server, one
//!   upstream socket per client address, like nginx `stream`.
//! - Sends master heartbeats from the public socket, so masters challenge (`getinfo`) the
//!   public address: the challenge goes through the relay like any client packet, the ET
//!   server answers it, and the master lists the public address.
use anyhow::{Context, Result};
use clap::Parser;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::net::{UdpSocket, lookup_host};
use tokio::sync::mpsc::{self, error::TrySendError};
use tokio::time::{Instant, interval, sleep_until, timeout};

/// Connectionless packet prefix (Quake 3 / ET protocol).
const OOB: &[u8] = b"\xff\xff\xff\xff";
/// Same heartbeat period as the ET server itself.
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(300);
const HEARTBEAT_ALIVE: &str = "EnemyTerritory-1";
const HEARTBEAT_DEAD: &str = "ETFlatline-1";
const PROBE_TIMEOUT: Duration = Duration::from_secs(3);
const PROBE_ATTEMPTS: usize = 3;
const SESSION_TIMEOUT: Duration = Duration::from_secs(120);
const SESSION_QUEUE: usize = 64;
const MAX_PACKET: usize = 65_535;
/// Longest connectionless command name shown in debug logs.
const MAX_COMMAND: usize = 32;

static DEBUG: AtomicBool = AtomicBool::new(false);

/// println! only with --debug.
macro_rules! debug {
    ($($arg:tt)*) => {
        if DEBUG.load(Ordering::Relaxed) {
            println!($($arg)*);
        }
    };
}

/// UDP relay + master heartbeat for a Wolfenstein: Enemy Territory server.
#[derive(Parser)]
#[command(version, name = "etrelay")]
struct Cli {
    /// ET server address, e.g. 10.0.0.2:27960
    upstream: SocketAddr,

    /// Public address to listen on
    #[arg(long, default_value = "0.0.0.0:27960")]
    listen: SocketAddr,

    /// Master server host:port (repeatable)
    #[arg(
        long = "master",
        default_values = ["etmaster.idsoftware.com:27950", "etmaster.net:27950"]
    )]
    masters: Vec<String>,

    /// Log sessions, connectionless packets (getinfo, getstatus, connect...), probes and heartbeats
    #[arg(long)]
    debug: bool,
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    DEBUG.store(cli.debug, Ordering::Relaxed);

    let public = UdpSocket::bind(cli.listen)
        .await
        .with_context(|| format!("bind {}", cli.listen))?;
    let public = Arc::new(public);

    println!(
        "etrelay: {} -> {}, masters: {}",
        cli.listen,
        cli.upstream,
        cli.masters.join(", ")
    );

    tokio::spawn(heartbeat(public.clone(), cli.upstream, cli.masters));
    relay(public, cli.upstream).await
}

/// Dispatches client packets to their session, opening one on first packet.
async fn relay(public: Arc<UdpSocket>, upstream: SocketAddr) -> Result<()> {
    let mut sessions: HashMap<SocketAddr, mpsc::Sender<Vec<u8>>> = HashMap::new();
    let mut buf = vec![0u8; MAX_PACKET];

    loop {
        let (n, client) = public.recv_from(&mut buf).await.context("recv")?;
        let mut packet = buf[..n].to_vec();

        if let Some(tx) = sessions.get(&client) {
            match tx.try_send(packet) {
                // Full: drop it, like any congested UDP hop.
                Ok(()) | Err(TrySendError::Full(_)) => continue,
                // Session expired: open a new one below.
                Err(TrySendError::Closed(p)) => packet = p,
            }
        }

        sessions.retain(|_, tx| !tx.is_closed());

        match open_session(public.clone(), upstream, client).await {
            Ok(tx) => {
                let _ = tx.try_send(packet);
                sessions.insert(client, tx);
                debug!("session + {client} ({} active)", sessions.len());
            }
            Err(e) => eprintln!("session {client}: {e:#}"),
        }
    }
}

/// Forwards packets between one client and the ET server until SESSION_TIMEOUT of silence.
async fn open_session(
    public: Arc<UdpSocket>,
    upstream: SocketAddr,
    client: SocketAddr,
) -> Result<mpsc::Sender<Vec<u8>>> {
    let sock = connect(upstream).await?;
    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(SESSION_QUEUE);

    tokio::spawn(async move {
        let mut buf = vec![0u8; MAX_PACKET];
        let mut deadline = Instant::now() + SESSION_TIMEOUT;
        let (mut sent, mut received) = (0u64, 0u64);

        loop {
            tokio::select! {
                Some(packet) = rx.recv() => {
                    if let Some(command) = command(&packet) {
                        debug!("{client} -> ET  {command} ({} B)", packet.len());
                    }
                    if let Err(e) = sock.send(&packet).await {
                        debug!("{client} -> ET  send failed: {e}");
                    }
                    sent += 1;
                    deadline = Instant::now() + SESSION_TIMEOUT;
                }
                res = sock.recv(&mut buf) => match res {
                    Ok(n) => {
                        if let Some(command) = command(&buf[..n]) {
                            debug!("{client} <- ET  {command} ({n} B)");
                        }
                        let _ = public.send_to(&buf[..n], client).await;
                        received += 1;
                        deadline = Instant::now() + SESSION_TIMEOUT;
                    }
                    // ICMP unreachable while the ET server is down: keep waiting.
                    Err(e) => debug!("{client} <- ET  recv failed: {e}"),
                },
                _ = sleep_until(deadline) => break,
            }
        }

        debug!("session - {client}: {sent} packets to ET, {received} from ET");
    });

    Ok(tx)
}

/// Heartbeats masters while the ET server answers, sends one flatline when it stops.
async fn heartbeat(public: Arc<UdpSocket>, upstream: SocketAddr, masters: Vec<String>) {
    let mut ticker = interval(HEARTBEAT_INTERVAL);
    let mut alive = false;

    loop {
        ticker.tick().await;

        let up = probe(upstream).await;
        debug!("probe {upstream}: {}", if up { "up" } else { "no answer" });
        if up != alive {
            println!("ET server {upstream} is {}", if up { "up" } else { "down" });
        }

        let message = match (up, alive) {
            (true, _) => HEARTBEAT_ALIVE,
            (false, true) => HEARTBEAT_DEAD,
            (false, false) => continue,
        };
        alive = up;

        let packet = [OOB, format!("heartbeat {message}\n").as_bytes()].concat();
        for master in &masters {
            match send_to_host(&public, &packet, master).await {
                Ok(addr) => debug!("heartbeat {message} -> {master} ({addr})"),
                Err(e) => eprintln!("heartbeat {master}: {e:#}"),
            }
        }
    }
}

/// True if the ET server answers `getinfo`.
async fn probe(upstream: SocketAddr) -> bool {
    for _ in 0..PROBE_ATTEMPTS {
        let attempt = async {
            let sock = connect(upstream).await?;
            sock.send(&[OOB, b"getinfo etrelay"].concat()).await?;

            let mut buf = vec![0u8; MAX_PACKET];
            let n = sock.recv(&mut buf).await?;
            Ok::<_, anyhow::Error>(buf[..n].starts_with(&[OOB, b"infoResponse"].concat()))
        };

        if let Ok(Ok(true)) = timeout(PROBE_TIMEOUT, attempt).await {
            return true;
        }
    }
    false
}

/// Resolves `host` (DNS may change between heartbeats) and sends from `sock`.
async fn send_to_host(sock: &UdpSocket, packet: &[u8], host: &str) -> Result<SocketAddr> {
    let ipv4 = sock.local_addr()?.is_ipv4();
    let addr = lookup_host(host)
        .await?
        .find(|a| a.is_ipv4() == ipv4)
        .context("no address in the listen socket family")?;

    sock.send_to(packet, addr).await?;
    Ok(addr)
}

/// Command name of a connectionless packet (`getinfo`, `infoResponse`, `connect`...),
/// None for in-game packets. Arguments are left out: `connect` carries the userinfo.
fn command(packet: &[u8]) -> Option<String> {
    let body = packet.strip_prefix(OOB)?;
    let end = body
        .iter()
        .position(|b| b.is_ascii_whitespace() || *b == b'\\')
        .unwrap_or(body.len())
        .min(MAX_COMMAND);

    Some(String::from_utf8_lossy(&body[..end]).into_owned())
}

/// Ephemeral socket connected to `upstream`.
async fn connect(upstream: SocketAddr) -> Result<UdpSocket> {
    let any: SocketAddr = if upstream.is_ipv4() {
        "0.0.0.0:0".parse()?
    } else {
        "[::]:0".parse()?
    };

    let sock = UdpSocket::bind(any).await?;
    sock.connect(upstream).await?;
    Ok(sock)
}