//! Multi-client integration tests for the Aura UDP transport (the v2 master-loop demuxer).
//!
//! These prove that a single bound [`UdpServer`] can simultaneously serve **many** peers, that bad
//! peers do not poison the server, and that established connections survive other peers coming and
//! going. The single-client and lossy-channel tests live in `udp_loopback.rs`; here we focus on
//! demuxer correctness.
//!
//! * [`udp_multi_client_two_concurrent`] — bind one server, drive two clients (different client CNs)
//!   to it concurrently, accept twice, and verify both connections are independent (no cross-talk;
//!   each side learns the correct peer id).
//! * [`udp_bad_ca_does_not_block_other_clients`] — a third client with a foreign CA fails the
//!   handshake; the server must keep accepting subsequent legitimate clients on the same port.
//! * [`udp_dropped_connection_does_not_block_other_clients`] — drop one client's connection mid-flight
//!   and prove the server keeps serving the other plus accepts a fresh one.

use std::sync::Arc;
use std::time::Duration;

use aura_pki::AuraCa;
use aura_proto::{ClientConfig, PacketConnection, ServerConfig};
use aura_transport::{UdpClient, UdpConnection, UdpOpts, UdpServer};

const SERVER_NAME: &str = "localhost";

/// Mint a CA, a server cert, and a set of client certs whose CNs are taken from `client_ids`.
fn make_configs(client_ids: &[&str]) -> (ServerConfig, Vec<ClientConfig>) {
    let ca = AuraCa::generate("Aura UDP Multi-Client Test CA").expect("generate CA");
    let server_cert = ca
        .issue_server_cert(SERVER_NAME)
        .expect("issue server cert");
    let ca_pem = ca.ca_cert_pem();
    let server_cfg = ServerConfig {
        ca_cert_pem: ca_pem.clone(),
        server_cert_pem: server_cert.cert_pem,
        server_key_pem: server_cert.key_pem,
    };
    let client_cfgs: Vec<ClientConfig> = client_ids
        .iter()
        .map(|id| {
            let c = ca.issue_client_cert(id).expect("issue client cert");
            ClientConfig {
                ca_cert_pem: ca_pem.clone(),
                client_cert_pem: c.cert_pem,
                client_key_pem: c.key_pem,
                server_name: SERVER_NAME.to_string(),
            }
        })
        .collect();
    (server_cfg, client_cfgs)
}

/// Mint a **separate** CA + matching client cert; the resulting `ClientConfig` will trust this CA
/// for the server (so it will reject the real server) and present a cert the real server will not
/// verify either. Used to drive a handshake failure that must NOT take down the server.
fn make_foreign_ca_client(server_name: &str, client_cn: &str) -> ClientConfig {
    let foreign = AuraCa::generate("Foreign CA").expect("generate foreign CA");
    let client_cert = foreign
        .issue_client_cert(client_cn)
        .expect("issue client cert under foreign CA");
    ClientConfig {
        ca_cert_pem: foreign.ca_cert_pem(),
        client_cert_pem: client_cert.cert_pem,
        client_key_pem: client_cert.key_pem,
        server_name: server_name.to_string(),
    }
}

/// Round-trip a payload `pkt` from `tx` to `rx` and assert byte equality.
async fn round_trip(tx: &Arc<dyn PacketConnection>, rx: &Arc<dyn PacketConnection>, pkt: &[u8]) {
    tx.send_packet(pkt).await.expect("send");
    let got = tokio::time::timeout(Duration::from_secs(5), rx.recv_packet())
        .await
        .expect("recv did not arrive within 5s")
        .expect("recv");
    assert_eq!(got, pkt, "payload mismatch over round trip");
}

#[tokio::test]
async fn udp_multi_client_two_concurrent() {
    let (server_cfg, client_cfgs) = make_configs(&["client-a", "client-b"]);
    let opts = UdpOpts::default();

    let server =
        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
    let server_addr = server.local_addr().expect("server addr");
    let server = Arc::new(server);

    // Spawn two server-side accepts in parallel; they must each pull their own connection from the
    // master-loop's accept queue.
    let s_a = server.clone();
    let accept_a = tokio::spawn(async move { s_a.accept().await });
    let s_b = server.clone();
    let accept_b = tokio::spawn(async move { s_b.accept().await });

    // Spawn the two clients concurrently. They share the server's bound port.
    let cfg_a = client_cfgs[0].clone();
    let cfg_b = client_cfgs[1].clone();
    let connect_a = tokio::spawn(async move { UdpClient::connect(server_addr, cfg_a, opts).await });
    let connect_b = tokio::spawn(async move { UdpClient::connect(server_addr, cfg_b, opts).await });

    // Wait for everything to settle (generous timeout — handshake should be sub-second on loopback).
    let timeout = Duration::from_secs(15);
    let server_a: UdpConnection = tokio::time::timeout(timeout, accept_a)
        .await
        .expect("accept_a within timeout")
        .expect("accept_a join")
        .expect("accept_a result");
    let server_b: UdpConnection = tokio::time::timeout(timeout, accept_b)
        .await
        .expect("accept_b within timeout")
        .expect("accept_b join")
        .expect("accept_b result");
    let client_a: UdpConnection = tokio::time::timeout(timeout, connect_a)
        .await
        .expect("connect_a within timeout")
        .expect("connect_a join")
        .expect("connect_a result");
    let client_b: UdpConnection = tokio::time::timeout(timeout, connect_b)
        .await
        .expect("connect_b within timeout")
        .expect("connect_b join")
        .expect("connect_b result");

    // Each server-side connection has a `peer_id` of either `client-a` or `client-b`; the accept
    // order is *not* guaranteed (whichever handshake finishes first), so detect which is which
    // and pair them with the matching client connection.
    let id_a = server_a.peer_id().map(str::to_owned);
    let id_b = server_b.peer_id().map(str::to_owned);
    let mut ids = vec![id_a.clone(), id_b.clone()];
    ids.sort();
    assert_eq!(
        ids,
        vec![Some("client-a".to_string()), Some("client-b".to_string())],
        "the two server-side connections must carry client-a and client-b CNs (no duplicates)"
    );

    let (srv_for_a, srv_for_b) = if id_a.as_deref() == Some("client-a") {
        (server_a, server_b)
    } else {
        (server_b, server_a)
    };

    // Each side sees its own peer id (client side sees the server name).
    assert_eq!(client_a.peer_id(), Some(SERVER_NAME));
    assert_eq!(client_b.peer_id(), Some(SERVER_NAME));

    let client_a: Arc<dyn PacketConnection> = Arc::new(client_a);
    let client_b: Arc<dyn PacketConnection> = Arc::new(client_b);
    let server_for_a: Arc<dyn PacketConnection> = Arc::new(srv_for_a);
    let server_for_b: Arc<dyn PacketConnection> = Arc::new(srv_for_b);

    // No cross-talk: A's payload reaches A's server-side conn (not B's), and vice versa.
    round_trip(&client_a, &server_for_a, b"hi from a").await;
    round_trip(&client_b, &server_for_b, b"hi from b").await;
    round_trip(&server_for_a, &client_a, b"reply to a").await;
    round_trip(&server_for_b, &client_b, b"reply to b").await;

    // And both directions still work concurrently (no head-of-line blocking via the master loop).
    let a_send = {
        let c = client_a.clone();
        let s = server_for_a.clone();
        tokio::spawn(async move {
            c.send_packet(b"a-concurrent").await.unwrap();
            s.recv_packet().await.unwrap()
        })
    };
    let b_send = {
        let c = client_b.clone();
        let s = server_for_b.clone();
        tokio::spawn(async move {
            c.send_packet(b"b-concurrent").await.unwrap();
            s.recv_packet().await.unwrap()
        })
    };
    assert_eq!(a_send.await.unwrap(), b"a-concurrent");
    assert_eq!(b_send.await.unwrap(), b"b-concurrent");
}

#[tokio::test]
async fn udp_bad_ca_does_not_block_other_clients() {
    let (server_cfg, client_cfgs) = make_configs(&["client-good"]);
    // Use a tighter handshake timeout so the failing peer fails quickly and the test finishes
    // even if the rogue client retransmits its ClientHello for a while.
    let opts = UdpOpts {
        hs_timeout: Duration::from_secs(3),
        ..UdpOpts::default()
    };

    let server =
        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
    let server_addr = server.local_addr().expect("server addr");
    let server = Arc::new(server);

    // A rogue client with a foreign CA: its server-side handshake task will fail. The server must
    // log + drop and keep accepting OTHER peers.
    let foreign_cfg = make_foreign_ca_client(SERVER_NAME, "rogue");
    let rogue =
        tokio::spawn(async move { UdpClient::connect(server_addr, foreign_cfg, opts).await });

    // Give the rogue task a head start so the server's master loop registers it first.
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Now the legitimate client connects. The server must still accept it.
    let cfg = client_cfgs[0].clone();
    let s = server.clone();
    let accept_good = tokio::spawn(async move { s.accept().await });
    let connect_good =
        tokio::spawn(async move { UdpClient::connect(server_addr, cfg, opts).await });

    let timeout = Duration::from_secs(15);
    let server_good: UdpConnection = tokio::time::timeout(timeout, accept_good)
        .await
        .expect("accept_good within timeout")
        .expect("accept_good join")
        .expect("accept_good result");
    let client_good: UdpConnection = tokio::time::timeout(timeout, connect_good)
        .await
        .expect("connect_good within timeout")
        .expect("connect_good join")
        .expect("connect_good result");

    assert_eq!(
        server_good.peer_id(),
        Some("client-good"),
        "server must learn the good client's CN despite the rogue peer"
    );

    let server_good: Arc<dyn PacketConnection> = Arc::new(server_good);
    let client_good: Arc<dyn PacketConnection> = Arc::new(client_good);
    round_trip(&client_good, &server_good, b"still serving").await;
    round_trip(&server_good, &client_good, b"yes we are").await;

    // The rogue connect should eventually fail (foreign CA → server's handshake rejects, the
    // client's handshake adapter then errors out on the deadline / chain mismatch). We do not
    // care about the exact error; we only require that it *errors*, not that it succeeds.
    let rogue_result = tokio::time::timeout(Duration::from_secs(10), rogue)
        .await
        .expect("rogue task should terminate")
        .expect("rogue task join");
    assert!(
        rogue_result.is_err(),
        "rogue client (foreign CA) must NOT succeed in establishing a connection"
    );
}

/// Establish ONE client against a running multi-client server, then verify the server-side conn
/// has the expected CN. The accept happens in its own spawned task to avoid blocking the connect.
async fn establish_one(
    server: &Arc<UdpServer>,
    server_addr: std::net::SocketAddr,
    cfg: ClientConfig,
    opts: UdpOpts,
    expect_cn: &str,
) -> (UdpConnection, UdpConnection) {
    let s = server.clone();
    let acc = tokio::spawn(async move { s.accept().await });
    let con = tokio::spawn(async move { UdpClient::connect(server_addr, cfg, opts).await });
    let timeout = Duration::from_secs(15);
    let srv = tokio::time::timeout(timeout, acc)
        .await
        .expect("accept timely")
        .expect("accept join")
        .expect("accept result");
    let cli = tokio::time::timeout(timeout, con)
        .await
        .expect("connect timely")
        .expect("connect join")
        .expect("connect result");
    assert_eq!(
        srv.peer_id(),
        Some(expect_cn),
        "server learned wrong CN for this client"
    );
    (srv, cli)
}

#[tokio::test]
async fn udp_dropped_connection_does_not_block_other_clients() {
    let (server_cfg, client_cfgs) = make_configs(&["client-1", "client-2", "client-3"]);
    let opts = UdpOpts::default();

    let server =
        UdpServer::bind("127.0.0.1:0".parse().unwrap(), server_cfg, opts).expect("bind server");
    let server_addr = server.local_addr().expect("server addr");
    let server = Arc::new(server);

    // Connect clients sequentially so the (server-side, client-side) pairing is unambiguous.
    let (srv1, cli1) = establish_one(
        &server,
        server_addr,
        client_cfgs[0].clone(),
        opts,
        "client-1",
    )
    .await;
    let (srv2, cli2) = establish_one(
        &server,
        server_addr,
        client_cfgs[1].clone(),
        opts,
        "client-2",
    )
    .await;

    let srv2: Arc<dyn PacketConnection> = Arc::new(srv2);
    let cli2: Arc<dyn PacketConnection> = Arc::new(cli2);
    // Sanity: client-2 works.
    round_trip(&cli2, &srv2, b"keep-1").await;

    // Drop both ends of client-1's pair: the server-side `UdpConnection` is dropped, its
    // `PeerSocket` (with the master's per-peer inbox receiver) is dropped, and the master loop's
    // next datagram from client-1's address — if any — will `Closed` and evict the entry.
    drop(srv1);
    drop(cli1);
    tokio::time::sleep(Duration::from_millis(50)).await;

    // The other client must keep working.
    round_trip(&cli2, &srv2, b"keep-2").await;
    round_trip(&srv2, &cli2, b"keep-3").await;

    // A fresh client-3 must also still be accepted.
    let (srv3, cli3) = establish_one(
        &server,
        server_addr,
        client_cfgs[2].clone(),
        opts,
        "client-3",
    )
    .await;
    let srv3: Arc<dyn PacketConnection> = Arc::new(srv3);
    let cli3: Arc<dyn PacketConnection> = Arc::new(cli3);
    round_trip(&cli3, &srv3, b"hi-3").await;
}