//! End-to-end responder-identity proofs over the production QUIC/TLS path.
//!
//! Only the deliberately hostile certificate-A/key-B resolver is test-only.
//! Identities, peer IDs, TLS policy, client startup, SNI construction, and
//! connection establishment all come from the production implementation.

use std::{
    future::Future,
    net::SocketAddr,
    ops::{Deref, DerefMut},
    panic::AssertUnwindSafe,
    sync::Arc,
    time::Duration,
};

use bytes::Bytes;
use eyre::{WrapErr as _, ensure};
use futures::FutureExt as _;
use lanspread_proto::{PeerEndpoint, PeerId};
use s2n_quic::{
    Connection,
    Server as QuicServer,
    provider::tls::rustls::server::Server as S2nRustlsServer,
};
use tokio_util::sync::CancellationToken;

use crate::{
    identity::PeerIdentity,
    quic_runtime::{
        EndpointTask,
        QuicClientRuntime,
        QuicConnector,
        start_quic_client,
        tracked_quic_io,
    },
    tls::{self, TestVerificationCounter, hostile_mismatched_server_provider},
};

const TEST_TIMEOUT: Duration = Duration::from_secs(5);
const LOOPBACK_EPHEMERAL: SocketAddr =
    SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST), 0);

struct CloseOnDrop(Connection);

impl Deref for CloseOnDrop {
    type Target = Connection;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for CloseOnDrop {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl Drop for CloseOnDrop {
    fn drop(&mut self) {
        self.0.close(0_u32.into());
    }
}

struct TestServer {
    server: Option<QuicServer>,
    endpoint: Option<EndpointTask>,
    addr: SocketAddr,
}

impl TestServer {
    async fn start(bind_addr: SocketAddr, tls: S2nRustlsServer) -> eyre::Result<Self> {
        let (io, control) = tracked_quic_io(bind_addr)?;
        let server = QuicServer::builder().with_tls(tls)?.with_io(io)?.start()?;
        let endpoint = control.take_started()?;
        let addr_result = server
            .local_addr()
            .map_err(eyre::Report::from)
            .and_then(|addr| {
                ensure!(
                    addr.port() != 0,
                    "test server did not resolve an ephemeral port"
                );
                Ok(addr)
            });
        match addr_result {
            Ok(addr) => Ok(Self {
                server: Some(server),
                endpoint: Some(endpoint),
                addr,
            }),
            Err(error) => {
                drop(server);
                let cleanup = endpoint.shutdown_and_join().await;
                merge_results(Err(error), cleanup)
            }
        }
    }

    fn server_mut(&mut self) -> eyre::Result<&mut QuicServer> {
        self.server
            .as_mut()
            .ok_or_else(|| eyre::eyre!("test server was already stopped"))
    }

    async fn shutdown(mut self) -> eyre::Result<()> {
        drop(self.server.take());
        let endpoint = self
            .endpoint
            .take()
            .ok_or_else(|| eyre::eyre!("test server endpoint was already joined"))?;
        endpoint.shutdown_and_join().await
    }
}

struct TestClient {
    runtime: Option<QuicClientRuntime>,
    connector: Option<QuicConnector>,
}

impl TestClient {
    fn start() -> eyre::Result<Self> {
        let (runtime, connector) = start_quic_client()?;
        Ok(Self {
            runtime: Some(runtime),
            connector: Some(connector),
        })
    }

    fn connector(&self) -> eyre::Result<&QuicConnector> {
        self.connector
            .as_ref()
            .ok_or_else(|| eyre::eyre!("test client was already stopped"))
    }

    async fn shutdown(mut self) -> eyre::Result<()> {
        drop(self.connector.take());
        self.runtime
            .take()
            .ok_or_else(|| eyre::eyre!("test client runtime was already joined"))?
            .shutdown()
            .await
    }

    async fn shutdown_rejected_handshake_fixture(mut self) -> eyre::Result<()> {
        drop(self.connector.take());
        self.runtime
            .take()
            .ok_or_else(|| eyre::eyre!("test client runtime was already joined"))?
            .shutdown_rejected_handshake_fixture()
            .await
    }
}

struct TestPair {
    server: TestServer,
    client: TestClient,
}

impl TestPair {
    async fn start(server_tls: S2nRustlsServer) -> eyre::Result<Self> {
        let server = TestServer::start(LOOPBACK_EPHEMERAL, server_tls).await?;
        match TestClient::start() {
            Ok(client) => Ok(Self { server, client }),
            Err(error) => {
                let cleanup = server.shutdown().await;
                merge_results(Err(error), cleanup)
            }
        }
    }

    async fn shutdown(self) -> eyre::Result<()> {
        let (server_result, client_result) =
            tokio::join!(self.server.shutdown(), self.client.shutdown());
        merge_results(server_result, client_result)
    }

    async fn shutdown_rejected_handshake_fixture(self) -> eyre::Result<()> {
        let server_result = self.server.shutdown().await;
        let client_result = self.client.shutdown_rejected_handshake_fixture().await;
        merge_results(server_result, client_result)
    }
}

fn merge_results<T>(primary: eyre::Result<T>, cleanup: eyre::Result<()>) -> eyre::Result<T> {
    match (primary, cleanup) {
        (Ok(value), Ok(())) => Ok(value),
        (Err(error), Ok(())) | (Ok(_), Err(error)) => Err(error),
        (Err(error), Err(cleanup_error)) => Err(eyre::eyre!(
            "operation failed: {error:#}; cleanup also failed: {cleanup_error:#}"
        )),
    }
}

async fn bounded<T>(
    cancellation: &CancellationToken,
    label: &'static str,
    operation: impl Future<Output = eyre::Result<T>>,
) -> eyre::Result<T> {
    tokio::select! {
        biased;
        () = cancellation.cancelled() => Err(eyre::eyre!("{label} cancelled")),
        result = tokio::time::timeout(TEST_TIMEOUT, operation) => {
            result.wrap_err_with(|| format!("{label} timed out"))?
        }
    }
}

async fn supervise<T>(
    future: impl Future<Output = eyre::Result<T>>,
    sibling_stop: CancellationToken,
) -> eyre::Result<T> {
    match AssertUnwindSafe(future).catch_unwind().await {
        Ok(Ok(value)) => Ok(value),
        Ok(Err(error)) => {
            sibling_stop.cancel();
            Err(error)
        }
        Err(_) => {
            sibling_stop.cancel();
            Err(eyre::eyre!("QUIC responder-identity proof branch panicked"))
        }
    }
}

async fn serve_one(
    server: &mut QuicServer,
    cancellation: CancellationToken,
    exchange_finished: Arc<tokio::sync::Barrier>,
) -> eyre::Result<()> {
    let connection = bounded(&cancellation, "server accept", async {
        server
            .accept()
            .await
            .ok_or_else(|| eyre::eyre!("test server closed before accepting a connection"))
    })
    .await?;
    let mut connection = CloseOnDrop(connection);
    let expected_alpn = tls::protocol_alpn();
    ensure!(
        connection.application_protocol()?.as_ref() == expected_alpn.as_slice(),
        "server negotiated the wrong production ALPN"
    );

    let stream = bounded(&cancellation, "server stream accept", async {
        connection
            .accept_bidirectional_stream()
            .await?
            .ok_or_else(|| eyre::eyre!("client closed before opening its request stream"))
    })
    .await?;
    let mut stream = stream;
    let mut request = Vec::new();
    while let Some(chunk) = bounded(&cancellation, "server request receive", async {
        Ok(stream.receive().await?)
    })
    .await?
    {
        request.extend_from_slice(&chunk);
        ensure!(request.len() <= 1, "client sent an oversized proof request");
    }
    ensure!(request == b"?", "client sent the wrong proof request");

    bounded(&cancellation, "server response send", async {
        stream.send(Bytes::from_static(&[0xa5])).await?;
        Ok(())
    })
    .await?;
    stream.finish()?;
    bounded(&cancellation, "server exchange barrier", async {
        exchange_finished.wait().await;
        Ok(())
    })
    .await
}

async fn request_one(
    client: &TestClient,
    endpoint: &PeerEndpoint,
    cancellation: CancellationToken,
    exchange_finished: Arc<tokio::sync::Barrier>,
) -> eyre::Result<()> {
    let connection = bounded(&cancellation, "client connect", async {
        client.connector()?.connect(endpoint).await
    })
    .await?;
    let mut connection = connection;
    let expected_alpn = tls::protocol_alpn();
    ensure!(
        connection.application_protocol()?.as_ref() == expected_alpn.as_slice(),
        "client negotiated the wrong production ALPN"
    );
    let stream = bounded(&cancellation, "client stream open", async {
        Ok(connection.open_bidirectional_stream().await?)
    })
    .await?;
    let mut stream = stream;
    bounded(&cancellation, "client request send", async {
        stream.send(Bytes::from_static(b"?")).await?;
        stream.finish()?;
        Ok(())
    })
    .await?;

    let mut response = Vec::new();
    while let Some(chunk) = bounded(&cancellation, "client response receive", async {
        Ok(stream.receive().await?)
    })
    .await?
    {
        response.extend_from_slice(&chunk);
        ensure!(
            response.len() <= 1,
            "server sent an oversized proof response"
        );
    }
    ensure!(response == [0xa5], "server sent the wrong proof response");
    bounded(&cancellation, "client exchange barrier", async {
        exchange_finished.wait().await;
        Ok(())
    })
    .await
}

async fn run_successful_exchange(
    server: &mut TestServer,
    client: &TestClient,
    endpoint: PeerEndpoint,
) -> eyre::Result<()> {
    let cancellation = CancellationToken::new();
    let exchange_finished = Arc::new(tokio::sync::Barrier::new(2));
    let server_future = supervise(
        serve_one(
            server.server_mut()?,
            cancellation.clone(),
            exchange_finished.clone(),
        ),
        cancellation.clone(),
    );
    let client_future = supervise(
        request_one(client, &endpoint, cancellation.clone(), exchange_finished),
        cancellation.clone(),
    );
    let (server_result, client_result) = tokio::join!(server_future, client_future);
    cancellation.cancel();
    merge_results(server_result, client_result)
}

async fn expect_rejected_connection(
    server: &mut TestServer,
    client: &TestClient,
    endpoint: PeerEndpoint,
) -> eyre::Result<()> {
    let cancellation = CancellationToken::new();
    let server_stop = cancellation.clone();
    let server_future = supervise(
        async {
            tokio::select! {
                biased;
                () = server_stop.cancelled() => Ok(()),
                accepted = server.server_mut()?.accept() => {
                    if let Some(connection) = accepted {
                        connection.close(0_u32.into());
                        Err(eyre::eyre!("server accepted a connection that should fail authentication"))
                    } else {
                        Err(eyre::eyre!("server endpoint closed during rejected handshake"))
                    }
                }
            }
        },
        cancellation.clone(),
    );
    let client_stop = cancellation.clone();
    let client_future = supervise(
        async {
            let result = tokio::select! {
                biased;
                () = client_stop.cancelled() => {
                    return Err(eyre::eyre!("client connection attempt was cancelled"));
                }
                result = tokio::time::timeout(
                    TEST_TIMEOUT,
                    client.connector()?.connect(&endpoint),
                ) => result,
            };
            let outcome = match result {
                Err(_) => Err(eyre::eyre!("client authentication rejection timed out")),
                Ok(Err(_)) => Ok(()),
                Ok(Ok(connection)) => {
                    drop(connection);
                    Err(eyre::eyre!("client accepted the wrong responder identity"))
                }
            };
            client_stop.cancel();
            outcome
        },
        cancellation.clone(),
    );

    let (server_result, client_result) = tokio::join!(server_future, client_future);
    cancellation.cancel();
    merge_results(client_result, server_result)
}

async fn execute_successful_exchanges(identity: &PeerIdentity, count: usize) -> eyre::Result<()> {
    let server_tls = tls::server_provider(identity)?;
    let mut pair = TestPair::start(server_tls).await?;
    let endpoint = PeerEndpoint {
        peer_id: identity.peer_id(),
        addr: pair.server.addr,
    };
    let mut operation = Ok(());
    for _ in 0..count {
        if operation.is_ok() {
            operation = run_successful_exchange(&mut pair.server, &pair.client, endpoint).await;
        }
    }
    let cleanup = pair.shutdown().await;
    merge_results(operation, cleanup)
}

async fn execute_rejected_exchange(
    server_tls: S2nRustlsServer,
    expected_peer: PeerId,
    verify_rejection: impl FnOnce() -> eyre::Result<()>,
) -> eyre::Result<()> {
    let mut pair = TestPair::start(server_tls).await?;
    let endpoint = PeerEndpoint {
        peer_id: expected_peer,
        addr: pair.server.addr,
    };
    let operation = expect_rejected_connection(&mut pair.server, &pair.client, endpoint)
        .await
        .and_then(|()| verify_rejection());
    let cleanup = pair.shutdown_rejected_handshake_fixture().await;
    merge_results(operation, cleanup)
}

#[test]
fn peer_id_sni_and_alpn_golden_vectors_are_production_values() -> eyre::Result<()> {
    let zero_vector = PeerId::from_bytes([0_u8; 32]);
    let zero_encoded = "a".repeat(52);
    ensure!(zero_vector.to_string() == zero_encoded);
    ensure!(zero_encoded.parse::<PeerId>()? == zero_vector);

    let ones_vector = PeerId::from_bytes([u8::MAX; 32]);
    let ones_encoded = format!("{}q", "7".repeat(51));
    ensure!(ones_vector.to_string() == ones_encoded);
    ensure!(ones_encoded.parse::<PeerId>()? == ones_vector);

    let identity = PeerIdentity::generate()?;
    let encoded = identity.peer_id().to_string();
    ensure!(encoded.len() == 52);
    ensure!(encoded.parse::<PeerId>()? == identity.peer_id());
    ensure!(encoded.to_uppercase().parse::<PeerId>().is_err());
    ensure!(tls::sni_for_peer(identity.peer_id())?.ends_with(".peer.lanspread.invalid"));
    ensure!(tls::protocol_alpn() == lanspread_proto::ALPN_PROTOCOL);
    ensure!(tls::protocol_alpn() != b"lanspread/7");
    Ok(())
}

#[tokio::test]
async fn correct_certificate_and_key_complete_a_full_production_quic_exchange() -> eyre::Result<()>
{
    let identity = PeerIdentity::generate()?;
    execute_successful_exchanges(&identity, 1).await
}

#[tokio::test]
async fn certificate_a_signed_by_private_key_b_fails_production_certificate_verify()
-> eyre::Result<()> {
    let identity_a = PeerIdentity::generate()?;
    let identity_b = PeerIdentity::generate()?;
    let counter = TestVerificationCounter::install()?;
    let server_tls = hostile_mismatched_server_provider(&identity_a, &identity_b)?;
    execute_rejected_exchange(server_tls, identity_a.peer_id(), || {
        let snapshot = counter.snapshot();
        ensure!(snapshot.certificate_calls == 1);
        ensure!(snapshot.certificate_accepts == 1);
        ensure!(snapshot.tls12_calls == 0);
        ensure!(snapshot.tls13_calls == 1);
        ensure!(snapshot.tls13_accepts == 0);
        Ok(())
    })
    .await
}

#[tokio::test]
async fn different_valid_peer_at_reused_address_fails_production_expected_id_pin()
-> eyre::Result<()> {
    let identity_a = PeerIdentity::generate()?;
    let identity_b = PeerIdentity::generate()?;
    let counter = TestVerificationCounter::install()?;
    let client = TestClient::start()?;
    let mut server_a =
        match TestServer::start(LOOPBACK_EPHEMERAL, tls::server_provider(&identity_a)?).await {
            Ok(server) => server,
            Err(error) => {
                let cleanup = client.shutdown().await;
                return merge_results(Err(error), cleanup);
            }
        };
    let reused_addr = server_a.addr;
    let endpoint_a = PeerEndpoint {
        peer_id: identity_a.peer_id(),
        addr: reused_addr,
    };
    let first_result = run_successful_exchange(&mut server_a, &client, endpoint_a).await;
    let server_a_cleanup = server_a.shutdown().await;
    if let Err(error) = merge_results(first_result, server_a_cleanup) {
        let cleanup = client.shutdown().await;
        return merge_results(Err(error), cleanup);
    }

    let mut server_b =
        match TestServer::start(reused_addr, tls::server_provider(&identity_b)?).await {
            Ok(server) => server,
            Err(error) => {
                let cleanup = client.shutdown().await;
                return merge_results(Err(error), cleanup);
            }
        };
    let rejection = if server_b.addr == reused_addr {
        expect_rejected_connection(&mut server_b, &client, endpoint_a).await
    } else {
        Err(eyre::eyre!(
            "test server did not reuse the exact socket address"
        ))
    };
    let (server_cleanup, client_cleanup) = tokio::join!(server_b.shutdown(), client.shutdown());
    let cleanup = merge_results(server_cleanup, client_cleanup);
    merge_results(rejection, cleanup)?;

    let snapshot = counter.snapshot();
    ensure!(snapshot.certificate_calls == 2);
    ensure!(snapshot.certificate_accepts == 1);
    ensure!(snapshot.tls12_calls == 0);
    ensure!(snapshot.tls13_calls == 1);
    ensure!(snapshot.tls13_accepts == 1);
    Ok(())
}

#[tokio::test]
async fn two_reconnects_repeat_production_certificate_and_signature_checks() -> eyre::Result<()> {
    let identity = PeerIdentity::generate()?;
    let counter = TestVerificationCounter::install()?;
    execute_successful_exchanges(&identity, 2).await?;
    let snapshot = counter.snapshot();
    ensure!(snapshot.certificate_calls == 2);
    ensure!(snapshot.certificate_accepts == 2);
    ensure!(snapshot.tls12_calls == 0);
    ensure!(snapshot.tls13_calls == 2);
    ensure!(snapshot.tls13_accepts == 2);
    Ok(())
}