//! Simple, fast connectivity tests with explicit timeouts.
//! - Two-node loopback connect: Tests bidirectional data exchange between client and server
//! - Three-node ring connect: Tests ring topology where each node connects to the next (1->2->3->1)
//! - Connection error scenarios: Tests timeout, certificate validation, and connection failure handling
//! - Connection lifecycle test: Tests graceful connection establishment, data exchange, and cleanup (each uses the others' endpoints)

#![allow(clippy::unwrap_used, clippy::expect_used)]

use ant_quic::{
    config::{ClientConfig, ServerConfig, TransportConfig},
    high_level::Endpoint,
};
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::time::{Duration, timeout};

fn gen_self_signed_cert() -> (Vec<CertificateDer<'static>>, PrivateKeyDer<'static>) {
    let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])
        .expect("generate self-signed");
    let cert_der = CertificateDer::from(cert.cert);
    let key_der = PrivateKeyDer::Pkcs8(cert.signing_key.serialize_der().into());
    (vec![cert_der], key_der)
}

fn transport_config_no_pqc() -> Arc<TransportConfig> {
    let mut transport = TransportConfig::default();
    transport.enable_pqc(false);
    Arc::new(transport)
}

async fn mk_server() -> (Endpoint, SocketAddr, Vec<CertificateDer<'static>>) {
    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();

    let (chain, key) = gen_self_signed_cert();
    let mut server_cfg = ServerConfig::with_single_cert(chain.clone(), key).expect("server cfg");
    server_cfg.transport_config(transport_config_no_pqc());
    let ep = Endpoint::server(server_cfg, ([127, 0, 0, 1], 0).into()).expect("server ep");
    let addr = ep.local_addr().expect("server addr");
    (ep, addr, chain)
}

fn mk_client_config(chain: &[CertificateDer<'static>]) -> ClientConfig {
    let mut roots = rustls::RootCertStore::empty();
    for c in chain.iter().cloned() {
        roots.add(c).expect("add root");
    }
    let mut config = ClientConfig::with_root_certificates(Arc::new(roots)).expect("client cfg");
    config.transport_config(transport_config_no_pqc());
    config
}

#[tokio::test]
async fn two_node_loopback_connect() {
    let (server, server_addr, chain) = mk_server().await;

    let accept = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(10), server.accept())
            .await
            .expect("accept wait")
            .expect("incoming");
        let conn = timeout(Duration::from_secs(10), inc)
            .await
            .expect("hs wait")
            .expect("server hs ok");

        // Test bidirectional data exchange on server side
        // Validates that the server can receive data from client and send responses
        let (mut send, mut recv) = conn.accept_bi().await.expect("accept bi");

        // Receive message from client
        let mut buf = [0; 1024];
        let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
            .await
            .expect("server read timeout")
            .expect("server read")
            .expect("server read some data");
        let received = std::str::from_utf8(&buf[..len]).expect("valid utf8");
        assert_eq!(
            received, "Hello from client!",
            "server received correct message"
        );

        // Send response back to client
        let response = b"Hello from server!";
        timeout(Duration::from_secs(5), send.write_all(response))
            .await
            .expect("server write timeout")
            .expect("server write");

        // Finish the stream first
        send.finish().expect("server finish stream");

        // Wait a bit for client to finish reading
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Gracefully close the connection
        conn.close(0u32.into(), b"test complete");
    });

    let mut client = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("client ep");
    client.set_default_client_config(mk_client_config(&chain));
    let connecting = client
        .connect(server_addr, "localhost")
        .expect("start connect");
    let conn = timeout(Duration::from_secs(10), connecting)
        .await
        .expect("client wait")
        .expect("client ok");

    // Test bidirectional data exchange on client side
    let (mut send, mut recv) = conn.open_bi().await.expect("open bi");

    // Send message to server
    let message = b"Hello from client!";
    timeout(Duration::from_secs(5), send.write_all(message))
        .await
        .expect("client write timeout")
        .expect("client write");

    // Receive response from server
    let mut buf = [0; 1024];
    let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
        .await
        .expect("client read timeout")
        .expect("client read")
        .expect("client read some data");
    let received = std::str::from_utf8(&buf[..len]).expect("valid utf8");
    assert_eq!(
        received, "Hello from server!",
        "client received correct response"
    );

    // Wait for server task to complete
    accept.await.expect("join");

    // Verify connection statistics
    let stats = conn.stats();
    assert!(stats.frame_rx.stream > 0, "received stream frames");
    assert!(stats.frame_tx.stream > 0, "sent stream frames");
}

#[tokio::test]
async fn three_node_ring_connect() {
    // Three servers
    let (s1, a1, c1) = mk_server().await;
    let (s2, a2, c2) = mk_server().await;
    let (s3, a3, c3) = mk_server().await;

    // Accept connections and test data exchange on each server
    let t1 = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(10), s1.accept())
            .await
            .expect("acc1 wait")
            .expect("incoming1");
        let conn = timeout(Duration::from_secs(10), inc)
            .await
            .expect("hs1 wait")
            .expect("hs1 ok");

        // Server 1 receives from client 3, sends to client 3
        let (mut send, mut recv) = conn.accept_bi().await.expect("s1 accept bi");

        let mut buf = [0; 1024];
        let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
            .await
            .expect("s1 read timeout")
            .expect("s1 read")
            .expect("s1 read data");
        let received = std::str::from_utf8(&buf[..len]).expect("s1 valid utf8");
        assert_eq!(received, "Hello from client 3!", "s1 received from c3");

        let response = b"Hello back from server 1!";
        timeout(Duration::from_secs(5), send.write_all(response))
            .await
            .expect("s1 write timeout")
            .expect("s1 write");

        send.finish().expect("s1 finish stream");

        // Wait a bit for client to finish reading
        tokio::time::sleep(Duration::from_millis(100)).await;

        conn.close(0u32.into(), b"ring test complete");
    });

    let t2 = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(10), s2.accept())
            .await
            .expect("acc2 wait")
            .expect("incoming2");
        let conn = timeout(Duration::from_secs(10), inc)
            .await
            .expect("hs2 wait")
            .expect("hs2 ok");

        // Server 2 receives from client 1, sends to client 1
        let (mut send, mut recv) = conn.accept_bi().await.expect("s2 accept bi");

        let mut buf = [0; 1024];
        let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
            .await
            .expect("s2 read timeout")
            .expect("s2 read")
            .expect("s2 read data");
        let received = std::str::from_utf8(&buf[..len]).expect("s2 valid utf8");
        assert_eq!(received, "Hello from client 1!", "s2 received from c1");

        let response = b"Hello back from server 2!";
        timeout(Duration::from_secs(5), send.write_all(response))
            .await
            .expect("s2 write timeout")
            .expect("s2 write");

        send.finish().expect("s2 finish stream");

        // Wait a bit for client to finish reading
        tokio::time::sleep(Duration::from_millis(100)).await;

        conn.close(0u32.into(), b"ring test complete");
    });

    let t3 = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(10), s3.accept())
            .await
            .expect("acc3 wait")
            .expect("incoming3");
        let conn = timeout(Duration::from_secs(10), inc)
            .await
            .expect("hs3 wait")
            .expect("hs3 ok");

        // Server 3 receives from client 2, sends to client 2
        let (mut send, mut recv) = conn.accept_bi().await.expect("s3 accept bi");

        let mut buf = [0; 1024];
        let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
            .await
            .expect("s3 read timeout")
            .expect("s3 read")
            .expect("s3 read data");
        let received = std::str::from_utf8(&buf[..len]).expect("s3 valid utf8");
        assert_eq!(received, "Hello from client 2!", "s3 received from c2");

        let response = b"Hello back from server 3!";
        timeout(Duration::from_secs(5), send.write_all(response))
            .await
            .expect("s3 write timeout")
            .expect("s3 write");

        send.finish().expect("s3 finish stream");

        // Wait a bit for client to finish reading
        tokio::time::sleep(Duration::from_millis(100)).await;

        conn.close(0u32.into(), b"ring test complete");
    });

    // Three clients each connecting to the next server (ring): 1->2, 2->3, 3->1
    let mut c_ep1 = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("c1 ep");
    c_ep1.set_default_client_config(mk_client_config(&c2));
    let mut c_ep2 = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("c2 ep");
    c_ep2.set_default_client_config(mk_client_config(&c3));
    let mut c_ep3 = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("c3 ep");
    c_ep3.set_default_client_config(mk_client_config(&c1));

    let c1_conn = timeout(
        Duration::from_secs(10),
        c_ep1.connect(a2, "localhost").expect("c1 start"),
    )
    .await
    .expect("c1 hs wait")
    .expect("c1 hs ok");
    let c2_conn = timeout(
        Duration::from_secs(10),
        c_ep2.connect(a3, "localhost").expect("c2 start"),
    )
    .await
    .expect("c2 hs wait")
    .expect("c2 hs ok");
    let c3_conn = timeout(
        Duration::from_secs(10),
        c_ep3.connect(a1, "localhost").expect("c3 start"),
    )
    .await
    .expect("c3 hs wait")
    .expect("c3 hs ok");

    // Test data exchange in the ring
    // Client 1 sends to server 2, receives from server 2
    let (mut c1_send, mut c1_recv) = c1_conn.open_bi().await.expect("c1 open bi");
    timeout(
        Duration::from_secs(5),
        c1_send.write_all(b"Hello from client 1!"),
    )
    .await
    .expect("c1 write timeout")
    .expect("c1 write");

    c1_send.finish().expect("c1 finish send");

    let mut buf = [0; 1024];
    let len = timeout(Duration::from_secs(5), c1_recv.read(&mut buf))
        .await
        .expect("c1 read timeout")
        .expect("c1 read")
        .expect("c1 read data");
    let received = std::str::from_utf8(&buf[..len]).expect("c1 response valid utf8");
    assert_eq!(received, "Hello back from server 2!", "c1 received from s2");

    // Client 2 sends to server 3, receives from server 3
    let (mut c2_send, mut c2_recv) = c2_conn.open_bi().await.expect("c2 open bi");
    timeout(
        Duration::from_secs(5),
        c2_send.write_all(b"Hello from client 2!"),
    )
    .await
    .expect("c2 write timeout")
    .expect("c2 write");

    c2_send.finish().expect("c2 finish send");

    let len = timeout(Duration::from_secs(5), c2_recv.read(&mut buf))
        .await
        .expect("c2 read timeout")
        .expect("c2 read")
        .expect("c2 read data");
    let received = std::str::from_utf8(&buf[..len]).expect("c2 response valid utf8");
    assert_eq!(received, "Hello back from server 3!", "c2 received from s3");

    // Client 3 sends to server 1, receives from server 1
    let (mut c3_send, mut c3_recv) = c3_conn.open_bi().await.expect("c3 open bi");
    timeout(
        Duration::from_secs(5),
        c3_send.write_all(b"Hello from client 3!"),
    )
    .await
    .expect("c3 write timeout")
    .expect("c3 write");

    c3_send.finish().expect("c3 finish send");

    let len = timeout(Duration::from_secs(5), c3_recv.read(&mut buf))
        .await
        .expect("c3 read timeout")
        .expect("c3 read")
        .expect("c3 read data");
    let received = std::str::from_utf8(&buf[..len]).expect("c3 response valid utf8");
    assert_eq!(received, "Hello back from server 1!", "c3 received from s1");

    // All servers accepted and completed data exchange
    t1.await.expect("t1 join");
    t2.await.expect("t2 join");
    t3.await.expect("t3 join");

    // Verify connection statistics for all clients
    let c1_stats = c1_conn.stats();
    let c2_stats = c2_conn.stats();
    let c3_stats = c3_conn.stats();

    assert!(
        c1_stats.frame_rx.stream > 0 && c1_stats.frame_tx.stream > 0,
        "c1 had data exchange"
    );
    assert!(
        c2_stats.frame_rx.stream > 0 && c2_stats.frame_tx.stream > 0,
        "c2 had data exchange"
    );
    assert!(
        c3_stats.frame_rx.stream > 0 && c3_stats.frame_tx.stream > 0,
        "c3 had data exchange"
    );
}

#[tokio::test]
async fn connection_error_scenarios() {
    // Test various connection failure scenarios to ensure robust error handling:
    // 1. Connection refused when no server is listening
    // 2. Certificate validation failures
    // 3. Connection timeouts during handshake

    // Test 1: Connection refused (no server listening)
    let mut client = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("client ep");
    let (chain, _) = gen_self_signed_cert();
    client.set_default_client_config(mk_client_config(&chain));

    let connecting = client
        .connect(([127, 0, 0, 1], 12345).into(), "localhost")
        .expect("connect call should succeed");
    let result = timeout(Duration::from_secs(5), connecting).await;

    // Should timeout or fail to connect (connection refused)
    match result {
        Ok(Ok(_)) => panic!("Expected connection to fail, but it succeeded"),
        Ok(Err(e)) => {
            // Connection failed as expected
            println!("Connection correctly failed: {:?}", e);
        }
        Err(_) => {
            // Timeout occurred as expected
            println!("Connection correctly timed out");
        }
    }

    // Test 2: Invalid certificate (client rejects server cert)
    let (server, server_addr, _) = mk_server().await;

    let server_task = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(5), server.accept()).await;
        match inc {
            Ok(Some(incoming)) => {
                let conn_result = timeout(Duration::from_secs(5), incoming).await;
                match conn_result {
                    Ok(Ok(_)) => println!("Server accepted connection (unexpected)"),
                    Ok(Err(e)) => println!("Server handshake failed as expected: {:?}", e),
                    Err(_) => println!("Server handshake timed out"),
                }
            }
            _ => println!("Server accept timed out or failed"),
        }
    });

    let mut client = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("client ep");
    // Use empty root store - should reject any certificate
    let roots = rustls::RootCertStore::empty();
    match ClientConfig::with_root_certificates(Arc::new(roots)) {
        Ok(client_config) => {
            client.set_default_client_config(client_config);
        }
        Err(e) => {
            // Config creation failed with empty roots - this is expected and acceptable
            // The test validates that certificate validation prevents insecure connections
            println!(
                "Certificate validation correctly prevented config creation: {:?}",
                e
            );
            return; // Test passes - certificate validation worked
        }
    }

    let connecting = client
        .connect(server_addr, "localhost")
        .expect("connect call should succeed");
    let connect_result = timeout(Duration::from_secs(5), connecting).await;

    match connect_result {
        Ok(Ok(_)) => panic!("Expected certificate validation to fail"),
        Ok(Err(e)) => println!("Certificate validation correctly failed: {:?}", e),
        Err(_) => println!("Certificate validation timed out (also acceptable)"),
    }

    server_task.await.expect("server task join");

    // Test 3: Connection timeout during handshake
    let (server, server_addr, chain) = mk_server().await;

    // Start server but don't accept connections immediately
    let server_task = tokio::spawn(async move {
        // Delay accepting to simulate slow server
        tokio::time::sleep(Duration::from_secs(10)).await;
        let _ = server.accept().await;
    });

    let mut client = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("client ep");
    client.set_default_client_config(mk_client_config(&chain));

    let connecting = client
        .connect(server_addr, "localhost")
        .expect("connect call should succeed");
    let connect_result = timeout(Duration::from_millis(100), connecting).await;

    // Should timeout before handshake completes
    match connect_result {
        Ok(_) => panic!("Expected handshake to timeout"),
        Err(_) => println!("Handshake correctly timed out"),
    }

    // Clean up server task
    server_task.abort();
}

#[tokio::test]
async fn connection_lifecycle_test() {
    let (server, server_addr, chain) = mk_server().await;

    let server_task = tokio::spawn(async move {
        let inc = timeout(Duration::from_secs(10), server.accept())
            .await
            .expect("server accept")
            .expect("incoming connection");

        let conn = timeout(Duration::from_secs(10), inc)
            .await
            .expect("server handshake")
            .expect("server handshake ok");

        // Test bidirectional streams
        let (mut send, mut recv) = conn.accept_bi().await.expect("server accept bi");

        // Receive data
        let mut buf = [0; 1024];
        let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
            .await
            .expect("server read")
            .expect("server read data")
            .expect("server received data");

        let message = std::str::from_utf8(&buf[..len]).expect("valid message");
        assert_eq!(message, "lifecycle test message");

        // Send response
        let response = b"acknowledged";
        timeout(Duration::from_secs(5), send.write_all(response))
            .await
            .expect("server write")
            .expect("server write ok");

        // Finish the stream
        send.finish().expect("server finish stream");

        // Wait for client to close connection
        // The connection should close gracefully
        let start_time = std::time::Instant::now();
        loop {
            tokio::time::sleep(Duration::from_millis(100)).await;
            if conn.stats().frame_rx.stream == 0 {
                break; // Connection appears idle
            }
            // Timeout after 5 seconds to prevent hanging
            if start_time.elapsed() > Duration::from_secs(5) {
                println!("Warning: Connection did not close gracefully within timeout");
                break;
            }
        }

        // Connection should be closed by client
        println!("Server: connection lifecycle test completed");
    });

    let mut client = Endpoint::client(([127, 0, 0, 1], 0).into()).expect("client ep");
    client.set_default_client_config(mk_client_config(&chain));

    let connecting = client
        .connect(server_addr, "localhost")
        .expect("connect call should succeed");
    let conn = timeout(Duration::from_secs(10), connecting)
        .await
        .expect("client connect")
        .expect("client connect ok");

    // Test bidirectional streams
    let (mut send, mut recv) = conn.open_bi().await.expect("client open bi");

    // Send data
    let message = b"lifecycle test message";
    timeout(Duration::from_secs(5), send.write_all(message))
        .await
        .expect("client write")
        .expect("client write ok");

    // Finish sending
    send.finish().expect("client finish stream");

    // Receive response
    let mut buf = [0; 1024];
    let len = timeout(Duration::from_secs(5), recv.read(&mut buf))
        .await
        .expect("client read")
        .expect("client read data")
        .expect("client received response");

    let response = std::str::from_utf8(&buf[..len]).expect("valid response");
    assert_eq!(response, "acknowledged");

    // Gracefully close the connection
    conn.close(0u32.into(), b"lifecycle test complete");

    // Wait a bit for the close to propagate
    tokio::time::sleep(Duration::from_millis(500)).await;

    // Verify connection is closed
    let stats = conn.stats();
    println!(
        "Client connection stats: frames_rx={}, frames_tx={}",
        stats.frame_rx.stream, stats.frame_tx.stream
    );
    assert!(stats.frame_tx.stream > 0, "client sent stream frames");

    server_task.await.expect("server task completed");

    // Test endpoint cleanup
    drop(client); // Should clean up client endpoint
    // Server endpoint is already moved into the task and will be cleaned up when the task completes

    println!("Connection lifecycle test passed - graceful shutdown and cleanup verified");
}
