//! Daemon-level encrypted transport proofs.
//!
//! These tests deliberately use the public client/server runtime entrypoints:
//! an application UDP socket talks to a client runtime, which talks through an
//! authenticated encrypted listener to a per-session UDP backend socket.

use std::{
    fs,
    net::SocketAddr,
    os::unix::fs::PermissionsExt,
    path::Path,
    process::Command,
    sync::{
        atomic::{AtomicU16, Ordering},
        Arc,
    },
    time::Duration,
};

use tempfile::TempDir;
use tokio::{
    net::UdpSocket,
    sync::mpsc,
    time::{sleep, timeout},
};
use wgcloak::{
    client,
    config::{self, Config, TransportMode},
    control, server,
    transport::{
        quic::QuicClientConnector, tls_tcp::client_tls_config, TlsTcpConnector, TransportConnector,
        TransportFuture, TransportSession,
    },
};

const SECRET_ONE: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
const SECRET_TWO: &str = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB";
const WRONG_SECRET: &str = "CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC";

fn free_udp_port() -> u16 {
    std::net::UdpSocket::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
        .port()
}

fn free_tcp_port() -> u16 {
    std::net::TcpListener::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
        .port()
}

fn private_runtime_dir(root: &Path) -> std::path::PathBuf {
    let path = root.join("runtime");
    fs::create_dir(&path).unwrap();
    fs::set_permissions(&path, fs::Permissions::from_mode(0o700)).unwrap();
    path
}

fn make_identity(dir: &TempDir) -> (std::path::PathBuf, std::path::PathBuf) {
    let cert = dir.path().join("cert.der");
    let key = dir.path().join("key.der");
    let pem_cert = dir.path().join("cert.pem");
    let pem_key = dir.path().join("key.pem");
    let result = Command::new("openssl")
        .args(["req", "-x509", "-newkey", "rsa:2048", "-nodes", "-keyout"])
        .arg(&pem_key)
        .args(["-out"])
        .arg(&pem_cert)
        .args([
            "-subj",
            "/CN=localhost",
            "-addext",
            "subjectAltName=DNS:localhost,IP:127.0.0.1",
            "-addext",
            "basicConstraints=critical,CA:FALSE",
            "-addext",
            "extendedKeyUsage=serverAuth",
            "-days",
            "1",
        ])
        .output()
        .unwrap();
    assert!(result.status.success(), "openssl failed: {:?}", result);
    let result = Command::new("openssl")
        .args(["x509", "-in"])
        .arg(&pem_cert)
        .args(["-outform", "DER", "-out"])
        .arg(&cert)
        .output()
        .unwrap();
    assert!(result.status.success());
    let result = Command::new("openssl")
        .args(["pkcs8", "-topk8", "-nocrypt", "-in"])
        .arg(&pem_key)
        .args(["-outform", "DER", "-out"])
        .arg(&key)
        .output()
        .unwrap();
    assert!(result.status.success());
    (cert, key)
}

fn server_config(
    cert: &Path,
    key: &Path,
    backend: SocketAddr,
    quic_port: u16,
    tls_port: u16,
    quic: bool,
    tls_tcp: bool,
) -> config::ServerConfig {
    Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='{backend}'\n[transports.direct_udp]\nenabled=false\n[transports.quic]\nenabled={quic}\nport={quic_port}\n[transports.http3]\nenabled=false\n[transports.tls_tcp]\nenabled={tls_tcp}\nport={tls_port}\n[tls]\ncertificate='{}'\nprivate_key='{}'\n[[auth.clients]]\nid='one'\nsecret='{SECRET_ONE}'\n[[auth.clients]]\nid='two'\nsecret='{SECRET_TWO}'\n[limits]\nhandshake_timeout_ms=2000\n",
        cert.display(),
        key.display(),
    ))
    .unwrap()
    .into_server()
    .unwrap()
}

fn client_config(
    id: &str,
    secret: &str,
    mode: &str,
    port: u16,
    local_port: u16,
) -> config::ClientConfig {
    Config::from_str(&format!(
        "[client]\nserver='localhost'\nid='{id}'\nsecret='{secret}'\n[local]\nlisten='127.0.0.1:{local_port}'\n[transport]\nmode='{mode}'\n[transport.auto]\nprobe_timeout_ms=500\nretries=1\n[transport.quic]\nport={port}\n[transport.tls_tcp]\nport={port}\n[network]\nmax_datagram_size=1200\n"
    ))
    .unwrap()
    .into_client()
    .unwrap()
}

fn quic_connector(config: config::ClientConfig, cert: &Path, key: &Path) -> QuicClientConnector {
    let certs = wgcloak::transport::quic::load_tls_identity(cert, key)
        .map(|(certs, _)| certs)
        .unwrap();
    let mut roots = rustls::RootCertStore::empty();
    roots.add(certs[0].clone()).unwrap();
    let tls = quinn::ClientConfig::with_root_certificates(Arc::new(roots)).unwrap();
    QuicClientConnector::with_client_config(config, tls)
}

fn tls_connector(config: config::ClientConfig, cert: &Path) -> TlsTcpConnector {
    let certs = wgcloak::transport::tls_tcp::load_certificates(cert).unwrap();
    let mut roots = rustls::RootCertStore::empty();
    roots.add(certs[0].clone()).unwrap();
    TlsTcpConnector::with_client_config(config, (*client_tls_config(roots)).clone())
}

struct RotatingQuicConnector {
    config: config::ClientConfig,
    cert: std::path::PathBuf,
    key: std::path::PathBuf,
    port: Arc<AtomicU16>,
}

impl TransportConnector for RotatingQuicConnector {
    fn connect<'a>(&'a self) -> TransportFuture<'a, Box<dyn TransportSession>> {
        Box::pin(async move {
            let mut config = self.config.clone();
            config.transport.quic.port = self.port.load(Ordering::Relaxed);
            quic_connector(config, &self.cert, &self.key)
                .connect()
                .await
        })
    }
}

async fn wait_for(path: &Path, command: &str, predicate: impl Fn(&str) -> bool) -> String {
    let mut last = None;
    timeout(Duration::from_secs(10), async {
        loop {
            if let Ok(response) = control::send_command(path, command).await {
                if predicate(&response) {
                    return response;
                }
                last = Some(response);
            }
            sleep(Duration::from_millis(25)).await;
        }
    })
    .await
    .unwrap_or_else(|_| {
        panic!(
            "control {} {command} did not reach expected state: {last:?}",
            path.display()
        )
    })
}

async fn echo_backend(socket: UdpSocket) {
    echo_backend_observed(socket, None).await;
}

async fn echo_backend_observed(
    socket: UdpSocket,
    observed: Option<mpsc::UnboundedSender<SocketAddr>>,
) {
    let mut buf = [0u8; 2048];
    while let Ok((len, peer)) = socket.recv_from(&mut buf).await {
        if let Some(observed) = &observed {
            let _ = observed.send(peer);
        }
        let _ = socket.send_to(&buf[..len], peer).await;
    }
}

async fn stop(path: &Path, task: tokio::task::JoinHandle<anyhow::Result<()>>) {
    let response = control::send_command(path, "shutdown").await.unwrap();
    assert!(response.contains("\"ok\":true"));
    timeout(Duration::from_secs(5), task)
        .await
        .expect("runtime did not stop")
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn explicit_quic_runtime_forwards_two_clients_and_reports_identity() {
    let dir = tempfile::tempdir().unwrap();
    let (cert, key) = make_identity(&dir);
    let runtime = private_runtime_dir(dir.path());
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let (peer_tx, mut peer_rx) = mpsc::unbounded_channel();
    let backend_task = tokio::spawn(echo_backend_observed(backend, Some(peer_tx)));
    let quic_port = free_udp_port();
    let server_path = runtime.join("server.sock");
    let server_task = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            quic_port,
            free_udp_port(),
            true,
            false,
        ),
        server_path.clone(),
    ));
    wait_for(&server_path, "status", |s| s.contains("transports=quic")).await;

    let mut clients = Vec::new();
    for (id, secret) in [("one", SECRET_ONE), ("two", SECRET_TWO)] {
        let local = free_udp_port();
        let app = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let cfg = client_config(id, secret, "quic", quic_port, local);
        let control_path = runtime.join(format!("client-{id}.sock"));
        let connector =
            Arc::new(quic_connector(cfg.clone(), &cert, &key)) as Arc<dyn TransportConnector>;
        let task = tokio::spawn(client::run_with_transport_connector_and_control_path(
            cfg,
            control_path.clone(),
            TransportMode::Quic,
            connector,
        ));
        wait_for(&control_path, "stats", |s| {
            s.contains("\"state\":\"connected\"")
        })
        .await;
        clients.push((local, app, control_path, task));
    }
    wait_for(&server_path, "stats", |s| {
        s.contains("\"active_sessions_by_transport\":{\"quic\":2")
            && s.contains("\"one\":1")
            && s.contains("\"two\":1")
    })
    .await;

    for (local, app, _, _) in &clients {
        app.send_to(b"encrypted-echo", ("127.0.0.1", *local))
            .await
            .unwrap();
        let mut reply = [0u8; 64];
        let (len, _) = timeout(Duration::from_secs(5), app.recv_from(&mut reply))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(&reply[..len], b"encrypted-echo");
    }
    let peer_one = timeout(Duration::from_secs(2), peer_rx.recv())
        .await
        .unwrap()
        .unwrap();
    let peer_two = timeout(Duration::from_secs(2), peer_rx.recv())
        .await
        .unwrap()
        .unwrap();
    assert_ne!(
        peer_one, peer_two,
        "each encrypted session needs its own backend socket"
    );
    for (_, _, control_path, task) in clients {
        stop(&control_path, task).await;
    }
    stop(&server_path, server_task).await;
    backend_task.abort();
}

#[tokio::test]
async fn explicit_tls_runtime_forwards_and_rejects_hostname_mismatch() {
    let dir = tempfile::tempdir().unwrap();
    let (cert, key) = make_identity(&dir);
    let runtime = private_runtime_dir(dir.path());
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let backend_task = tokio::spawn(echo_backend(backend));
    let tls_port = free_tcp_port();
    let server_path = runtime.join("server.sock");
    let server_task = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            free_udp_port(),
            tls_port,
            false,
            true,
        ),
        server_path.clone(),
    ));
    wait_for(&server_path, "status", |s| s.contains("transports=tls_tcp")).await;
    let local = free_udp_port();
    let app = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let cfg = client_config("one", SECRET_ONE, "tls_tcp", tls_port, local);
    let client_path = runtime.join("client.sock");
    let connector = Arc::new(tls_connector(cfg.clone(), &cert)) as Arc<dyn TransportConnector>;
    let client_task = tokio::spawn(client::run_with_transport_connector_and_control_path(
        cfg.clone(),
        client_path.clone(),
        TransportMode::TlsTcp,
        connector,
    ));
    wait_for(&client_path, "stats", |s| {
        s.contains("\"state\":\"connected\"")
    })
    .await;
    app.send_to(b"tls-echo", ("127.0.0.1", local))
        .await
        .unwrap();
    let mut reply = [0u8; 64];
    let (len, _) = timeout(Duration::from_secs(5), app.recv_from(&mut reply))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(&reply[..len], b"tls-echo");
    wait_for(&server_path, "stats", |s| {
        s.contains("\"active_sessions_by_transport\":{\"tls_tcp\":1")
            && s.contains("\"active_sessions_by_client_id\":{\"one\":1")
    })
    .await;
    // A trusted certificate with an unrelated hostname must still fail normal
    // hostname verification; trusting the issuing certificate is not enough.
    let mut mismatch = cfg.clone();
    mismatch.client.server = "not-localhost".to_owned();
    let result = tls_connector(mismatch, &cert).connect().await;
    assert!(matches!(
        result,
        Err(wgcloak::transport::TransportError::Io(_))
    ));
    stop(&client_path, client_task).await;
    stop(&server_path, server_task).await;
    backend_task.abort();
}

#[tokio::test]
async fn encrypted_runtime_cancellation_releases_server_sessions() {
    let dir = tempfile::tempdir().unwrap();
    let (cert, key) = make_identity(&dir);
    let runtime = private_runtime_dir(dir.path());
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let backend_task = tokio::spawn(echo_backend(backend));
    let quic_port = free_udp_port();
    let server_path = runtime.join("server.sock");
    let server_task = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            quic_port,
            free_udp_port(),
            true,
            false,
        ),
        server_path.clone(),
    ));
    wait_for(&server_path, "status", |s| s.contains("transports=quic")).await;
    let local = free_udp_port();
    let cfg = client_config("one", SECRET_ONE, "quic", quic_port, local);
    let client_path = runtime.join("client.sock");
    let connector =
        Arc::new(quic_connector(cfg.clone(), &cert, &key)) as Arc<dyn TransportConnector>;
    let client_task = tokio::spawn(client::run_with_transport_connector_and_control_path(
        cfg,
        client_path.clone(),
        TransportMode::Quic,
        connector,
    ));
    wait_for(&client_path, "stats", |s| {
        s.contains("\"state\":\"connected\"")
    })
    .await;
    wait_for(&server_path, "stats", |s| {
        s.contains("\"active_sessions\":1")
    })
    .await;
    stop(&client_path, client_task).await;
    wait_for(&server_path, "stats", |s| {
        s.contains("\"active_sessions\":0")
    })
    .await;
    stop(&server_path, server_task).await;
    backend_task.abort();
}

#[tokio::test]
async fn wrong_secret_is_authentication_rejection_without_fallback() {
    let dir = tempfile::tempdir().unwrap();
    let (cert, key) = make_identity(&dir);
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let backend_task = tokio::spawn(echo_backend(backend));
    let quic_port = free_udp_port();
    let server_path = private_runtime_dir(dir.path()).join("server.sock");
    let server_task = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            quic_port,
            free_udp_port(),
            true,
            false,
        ),
        server_path.clone(),
    ));
    wait_for(&server_path, "status", |s| s.contains("transports=quic")).await;
    let cfg = client_config("one", WRONG_SECRET, "quic", quic_port, free_udp_port());
    let result = timeout(
        Duration::from_secs(5),
        QuicClientConnector::with_client_config(cfg.clone(), {
            let certs = wgcloak::transport::quic::load_tls_identity(&cert, &key)
                .unwrap()
                .0;
            let mut roots = rustls::RootCertStore::empty();
            roots.add(certs[0].clone()).unwrap();
            quinn::ClientConfig::with_root_certificates(Arc::new(roots)).unwrap()
        })
        .connect(),
    )
    .await
    .unwrap();
    assert!(matches!(
        result,
        Err(wgcloak::transport::TransportError::AuthenticationRejected(
            _
        ))
    ));
    stop(&server_path, server_task).await;
    backend_task.abort();
}

#[tokio::test]
async fn quic_runtime_reconnects_after_server_restart_and_cleans_on_cancel() {
    let dir = tempfile::tempdir().unwrap();
    let (cert, key) = make_identity(&dir);
    let runtime = private_runtime_dir(dir.path());
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let backend_task = tokio::spawn(echo_backend(backend));
    let quic_port = free_udp_port();
    let server_path = runtime.join("server.sock");
    let first_server = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            quic_port,
            free_udp_port(),
            true,
            false,
        ),
        server_path.clone(),
    ));
    wait_for(&server_path, "status", |s| s.contains("transports=quic")).await;
    let local = free_udp_port();
    let app = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let cfg = client_config("one", SECRET_ONE, "quic", quic_port, local);
    let client_path = runtime.join("client.sock");
    let rotating_port = Arc::new(AtomicU16::new(quic_port));
    let connector = Arc::new(RotatingQuicConnector {
        config: cfg.clone(),
        cert: cert.clone(),
        key: key.clone(),
        port: rotating_port.clone(),
    }) as Arc<dyn TransportConnector>;
    let client_task = tokio::spawn(client::run_with_transport_connector_and_control_path(
        cfg,
        client_path.clone(),
        TransportMode::Quic,
        connector,
    ));
    wait_for(&client_path, "stats", |s| {
        s.contains("\"state\":\"connected\"")
    })
    .await;
    app.send_to(b"before-restart", ("127.0.0.1", local))
        .await
        .unwrap();
    let mut reply = [0u8; 64];
    let (len, _) = timeout(Duration::from_secs(5), app.recv_from(&mut reply))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(&reply[..len], b"before-restart");

    control::send_command(&server_path, "shutdown")
        .await
        .unwrap();
    timeout(Duration::from_secs(5), first_server)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
    let replacement_port = free_udp_port();
    rotating_port.store(replacement_port, Ordering::Relaxed);
    let replacement_path = runtime.join("server-replacement.sock");
    let replacement = tokio::spawn(server::run_with_control_path(
        server_config(
            &cert,
            &key,
            backend_addr,
            replacement_port,
            free_udp_port(),
            true,
            false,
        ),
        replacement_path.clone(),
    ));
    wait_for(&replacement_path, "status", |s| {
        s.contains("transports=quic")
    })
    .await;
    wait_for(&client_path, "stats", |s| {
        s.contains("\"state\":\"connected\"") && s.contains("\"reconnections\":1")
    })
    .await;
    app.send_to(b"after-restart", ("127.0.0.1", local))
        .await
        .unwrap();
    let (len, _) = timeout(Duration::from_secs(5), app.recv_from(&mut reply))
        .await
        .unwrap()
        .unwrap();
    assert_eq!(&reply[..len], b"after-restart");
    stop(&client_path, client_task).await;
    stop(&replacement_path, replacement).await;
    backend_task.abort();
}