use std::{collections::HashMap, fs, process::Command, sync::Arc, time::Duration};

use base64::Engine;
use bytes::Bytes;
use quinn::{ClientConfig, Endpoint, ServerConfig};
use rustls::pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer};
use tokio::time::timeout;
use wgcloak::config;
use wgcloak::transport::{
    http3::{Http3ClientConnection, Http3Error, Http3Server, Http3ServerConnection},
    AdmissionController, TransportSession,
};

const SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
const WRONG_SECRET: &str = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB";

fn test_endpoints() -> (
    tempfile::TempDir,
    Endpoint,
    ClientConfig,
    std::net::SocketAddr,
) {
    let directory = tempfile::tempdir().unwrap();
    let certificate = directory.path().join("localhost-cert.pem");
    let private_key = directory.path().join("localhost-key.pem");
    let result = Command::new("openssl")
        .args([
            "req",
            "-x509",
            "-newkey",
            "rsa:2048",
            "-nodes",
            "-keyout",
            private_key.to_str().unwrap(),
            "-out",
            certificate.to_str().unwrap(),
            "-subj",
            "/CN=localhost",
            "-addext",
            "subjectAltName=DNS:localhost",
            "-addext",
            "basicConstraints=critical,CA:FALSE",
            "-addext",
            "keyUsage=critical,digitalSignature,keyEncipherment",
            "-addext",
            "extendedKeyUsage=serverAuth",
            "-days",
            "1",
        ])
        .output()
        .unwrap();
    assert!(
        result.status.success(),
        "openssl failed: {}",
        String::from_utf8_lossy(&result.stderr)
    );

    let cert_pem = fs::read(&certificate).unwrap();
    let certs: Vec<CertificateDer<'static>> = CertificateDer::pem_slice_iter(&cert_pem)
        .collect::<Result<_, _>>()
        .unwrap();
    let key_pem = fs::read(&private_key).unwrap();
    let key = PrivateKeyDer::from_pem_slice(&key_pem).unwrap();
    let mut server_config = ServerConfig::with_single_cert(certs.clone(), key).unwrap();
    let mut server_transport = quinn::TransportConfig::default();
    server_transport.datagram_send_buffer_size(1024 * 1024);
    server_config.transport_config(Arc::new(server_transport));
    let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap();
    let server_address = server.local_addr().unwrap();

    let mut roots = rustls::RootCertStore::empty();
    roots.add(certs[0].clone()).unwrap();
    let mut client_config = ClientConfig::with_root_certificates(Arc::new(roots)).unwrap();
    let mut client_transport = quinn::TransportConfig::default();
    client_transport.datagram_send_buffer_size(1024 * 1024);
    client_config.transport_config(Arc::new(client_transport));
    (directory, server, client_config, server_address)
}

async fn connect_client(address: std::net::SocketAddr, config: ClientConfig) -> quinn::Connection {
    let mut endpoint = Endpoint::client("[::]:0".parse().unwrap()).unwrap();
    endpoint.set_default_client_config(config);
    endpoint
        .connect(address, "localhost")
        .unwrap()
        .await
        .unwrap()
}

#[tokio::test]
async fn connect_udp_authenticates_then_exchanges_contextual_http_datagrams() {
    let (_certificate_dir, server_endpoint, client_tls, address) = test_endpoints();
    let mut clients = HashMap::new();
    clients.insert(
        "laptop".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let server_task = tokio::spawn(async move {
        let incoming = server_endpoint.accept().await.unwrap();
        let quic_connection = incoming.await.unwrap();
        Http3ServerConnection::new(quic_connection)
            .await
            .unwrap()
            .accept_connect_udp("/tunnel", &clients, 1200, Duration::from_secs(3))
            .await
            .unwrap()
    });

    let quic_connection = connect_client(address, client_tls).await;
    let client = Http3ClientConnection::new(quic_connection).await.unwrap();
    let client_session = client
        .connect_udp(
            &format!("localhost:{}", address.port()),
            "/tunnel",
            "laptop",
            &base64::engine::general_purpose::STANDARD
                .decode(SECRET)
                .unwrap(),
            1200,
            Duration::from_secs(3),
        )
        .await
        .unwrap();
    let server_session = timeout(Duration::from_secs(3), server_task)
        .await
        .unwrap()
        .unwrap();

    assert_eq!(
        client_session.max_datagram_size(),
        server_session.max_datagram_size()
    );
    assert!(client_session.max_datagram_size() < 1200);
    client_session
        .send_datagram(b"opaque-upstream")
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), server_session.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"opaque-upstream")
    );
    server_session
        .send_datagram(b"opaque-downstream")
        .await
        .unwrap();
    assert_eq!(
        timeout(Duration::from_secs(2), client_session.recv_datagram())
            .await
            .unwrap()
            .unwrap(),
        Bytes::from_static(b"opaque-downstream")
    );
    let too_large = vec![0; client_session.max_datagram_size() + 1];
    assert!(matches!(
        client_session.send_datagram(&too_large).await,
        Err(Http3Error::Oversized { .. })
    ));
    client_session.close(Duration::from_secs(2)).await.unwrap();
    assert!(matches!(
        client_session.send_datagram(b"after-close").await,
        Err(Http3Error::Closed)
    ));
}

#[tokio::test]
async fn connect_udp_rejects_bad_secret_before_returning_session() {
    let (_certificate_dir, server_endpoint, client_tls, address) = test_endpoints();
    let mut clients = HashMap::new();
    clients.insert(
        "laptop".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let server_task = tokio::spawn(async move {
        let incoming = server_endpoint.accept().await.unwrap();
        let quic_connection = incoming.await.unwrap();
        let observer = quic_connection.clone();
        let result = Http3ServerConnection::new(quic_connection)
            .await
            .unwrap()
            .accept_connect_udp("/tunnel", &clients, 1200, Duration::from_secs(3))
            .await;
        (result, observer)
    });

    let quic_connection = connect_client(address, client_tls).await;
    let client = Http3ClientConnection::new(quic_connection).await.unwrap();
    let client_result = client
        .connect_udp(
            &format!("localhost:{}", address.port()),
            "/tunnel",
            "laptop",
            &base64::engine::general_purpose::STANDARD
                .decode(WRONG_SECRET)
                .unwrap(),
            1200,
            Duration::from_secs(3),
        )
        .await;
    assert!(
        matches!(client_result, Err(Http3Error::AuthenticationRejected)),
        "unexpected client error: {:?}",
        client_result.as_ref().err()
    );
    let (server_result, server_connection) = timeout(Duration::from_secs(3), server_task)
        .await
        .unwrap()
        .unwrap();
    assert!(
        matches!(server_result, Err(Http3Error::AuthenticationRejected)),
        "unexpected server error: {:?}",
        server_result.as_ref().err()
    );
    timeout(Duration::from_secs(2), async {
        loop {
            if server_connection.close_reason().is_some() {
                break;
            }
            tokio::task::yield_now().await;
        }
    })
    .await
    .expect("authentication rejection must close the QUIC connection");
}

#[tokio::test]
async fn http3_server_admission_permit_tracks_authenticated_session_lifetime() {
    let (_certificate_dir, endpoint, client_tls, address) = test_endpoints();
    let mut clients = HashMap::new();
    clients.insert(
        "laptop".to_owned(),
        base64::engine::general_purpose::STANDARD
            .decode(SECRET)
            .unwrap(),
    );
    let limits = config::LimitsConfig {
        max_sessions: 1,
        max_sessions_per_ip: 1,
        ..config::LimitsConfig::default()
    };
    let admission = AdmissionController::new(&limits);
    let server = Arc::new(Http3Server::from_endpoint_with_admission(
        endpoint,
        "/tunnel",
        clients,
        1200,
        Duration::from_secs(3),
        admission.clone(),
    ));
    let accept_server = server.clone();
    let server_task = tokio::spawn(async move { accept_server.accept_session().await.unwrap() });

    let quic_connection = connect_client(address, client_tls).await;
    let client = Http3ClientConnection::new(quic_connection).await.unwrap();
    let client_session = client
        .connect_udp(
            &format!("localhost:{}", address.port()),
            "/tunnel",
            "laptop",
            &base64::engine::general_purpose::STANDARD
                .decode(SECRET)
                .unwrap(),
            1200,
            Duration::from_secs(3),
        )
        .await
        .unwrap();
    let server_session = timeout(Duration::from_secs(3), server_task)
        .await
        .unwrap()
        .unwrap();

    assert_eq!(admission.active_sessions(), 1);
    assert_eq!(admission.pending_handshakes(), 0);
    server_session.close().await.unwrap();
    assert_eq!(admission.active_sessions(), 0);
    timeout(
        Duration::from_secs(2),
        client_session.close(Duration::from_millis(250)),
    )
    .await
    .expect("client cleanup after peer close must be bounded")
    .unwrap_or_else(|error| panic!("client cleanup returned an unexpected error: {error}"));
}