#![cfg(feature = "quic")]

use std::future;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
use std::sync::Arc;

use axum::http::Request;
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use bytes::{Buf, Bytes};
use narsil_engine::{BurnBackend, ObservabilityState, SharedInferenceService, serve_quic};
use quinn::crypto::rustls::QuicClientConfig;
use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use serde_json::{Value, json};
use tokio::sync::oneshot;

/// TEST-ONLY certificate verifier that accepts any server certificate so the integration
/// test can talk to the dev self-signed QUIC endpoint. This MUST NEVER be moved into the
/// library or any client code: it disables TLS server authentication entirely.
#[derive(Debug)]
struct SkipServerVerification(Arc<rustls::crypto::CryptoProvider>);

impl SkipServerVerification {
    fn new() -> Arc<Self> {
        Arc::new(Self(
            Arc::new(rustls::crypto::aws_lc_rs::default_provider()),
        ))
    }
}

impl ServerCertVerifier for SkipServerVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp: &[u8],
        _now: UnixTime,
    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &rustls::DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(
            message,
            cert,
            dss,
            &self.0.signature_verification_algorithms,
        )
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &rustls::DigitallySignedStruct,
    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(
            message,
            cert,
            dss,
            &self.0.signature_verification_algorithms,
        )
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        self.0.signature_verification_algorithms.supported_schemes()
    }
}

fn f32_bytes(values: &[f32]) -> Vec<u8> {
    values
        .iter()
        .flat_map(|value| value.to_le_bytes())
        .collect()
}

fn f32_values(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(4)
        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
        .collect()
}

fn unused_udp_addr() -> SocketAddr {
    let socket = UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap();
    socket.local_addr().unwrap()
}

fn client_config() -> h3_quinn::quinn::ClientConfig {
    let mut tls_config = rustls::ClientConfig::builder_with_provider(Arc::new(
        rustls::crypto::aws_lc_rs::default_provider(),
    ))
    .with_protocol_versions(&[&rustls::version::TLS13])
    .unwrap()
    .dangerous()
    .with_custom_certificate_verifier(SkipServerVerification::new())
    .with_no_client_auth();
    tls_config.alpn_protocols = vec![b"h3".to_vec()];

    h3_quinn::quinn::ClientConfig::new(Arc::new(QuicClientConfig::try_from(tls_config).unwrap()))
}

async fn connect(
    addr: SocketAddr,
) -> std::result::Result<
    (h3_quinn::quinn::Endpoint, h3_quinn::quinn::Connection),
    Box<dyn std::error::Error>,
> {
    let mut endpoint =
        h3_quinn::quinn::Endpoint::client(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0).into())?;
    endpoint.set_default_client_config(client_config());

    let connection = endpoint.connect(addr, "localhost")?.await?;
    Ok((endpoint, connection))
}

/// Connect with a bounded retry loop instead of a fixed sleep so the test isn't flaky when
/// the server is still binding the UDP socket (common under CI load).
async fn connect_with_retry(
    addr: SocketAddr,
) -> (h3_quinn::quinn::Endpoint, h3_quinn::quinn::Connection) {
    let mut attempts = 0;
    loop {
        match connect(addr).await {
            Ok(res) => break res,
            Err(err) => {
                attempts += 1;
                if attempts >= 50 {
                    panic!("failed to connect after {attempts} attempts: {err}");
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        }
    }
}

/// Drive an HTTP/3 request to completion and return the response status plus body bytes.
async fn send_request(
    connection: h3_quinn::quinn::Connection,
    method: &str,
    path: &str,
    body: Option<Vec<u8>>,
) -> (axum::http::StatusCode, Vec<u8>) {
    let (mut driver, mut client) = h3::client::new(h3_quinn::Connection::new(connection))
        .await
        .unwrap();
    let driver = tokio::spawn(async move {
        let _ = future::poll_fn(|cx| driver.poll_close(cx)).await;
    });

    let uri = format!("https://localhost{path}");
    let request = Request::builder().method(method).uri(uri).body(()).unwrap();
    let mut stream = client.send_request(request).await.unwrap();
    if let Some(body) = body {
        stream.send_data(Bytes::from(body)).await.unwrap();
    }
    stream.finish().await.unwrap();

    let response = stream.recv_response().await.unwrap();
    let status = response.status();

    let mut bytes = Vec::new();
    while let Some(mut chunk) = stream.recv_data().await.unwrap() {
        while chunk.has_remaining() {
            let data = chunk.copy_to_bytes(chunk.remaining());
            bytes.extend_from_slice(&data);
        }
    }

    driver.abort();
    (status, bytes)
}

fn affine_input(values: &[f32]) -> Value {
    json!({
        "model_name": "burn-affine",
        "inputs": [{
            "name": "input",
            "dtype": "F32",
            "shape": [values.len()],
            "data": BASE64.encode(f32_bytes(values)),
        }],
        "parameters": {},
    })
}

fn spawn_server() -> (
    SocketAddr,
    oneshot::Sender<()>,
    tokio::task::JoinHandle<narsil_engine::Result<()>>,
) {
    let addr = unused_udp_addr();
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    let service = SharedInferenceService::new(BurnBackend::with_affine(2.0, 1.0));
    let observability = ObservabilityState::new(&service).unwrap();
    let server = tokio::spawn(serve_quic(addr, service, observability, async move {
        let _ = shutdown_rx.await;
    }));
    (addr, shutdown_tx, server)
}

#[tokio::test]
async fn quic_transport_serves_http3_infer() {
    let (addr, shutdown_tx, server) = spawn_server();

    let (_endpoint, connection) = connect_with_retry(addr).await;
    let request = serde_json::to_vec(&affine_input(&[2.0, 4.0])).unwrap();
    let (status, body) = send_request(connection, "POST", "/v1/infer", Some(request)).await;
    assert_eq!(status, axum::http::StatusCode::OK);

    let response: Value = serde_json::from_slice(&body).unwrap();
    let output = BASE64
        .decode(response["outputs"][0]["data"].as_str().unwrap())
        .unwrap();
    assert_eq!(f32_values(&output), vec![5.0, 9.0]);

    let _ = shutdown_tx.send(());
    server.await.unwrap().unwrap();
}

#[tokio::test]
async fn quic_transport_serves_http3_infer_batch() {
    let (addr, shutdown_tx, server) = spawn_server();

    let (_endpoint, connection) = connect_with_retry(addr).await;
    let request = json!({
        "requests": [
            affine_input(&[1.0, 2.0]),
            affine_input(&[3.0]),
        ],
    });
    let body = serde_json::to_vec(&request).unwrap();
    let (status, body) = send_request(connection, "POST", "/v1/infer_batch", Some(body)).await;
    assert_eq!(status, axum::http::StatusCode::OK);

    let response: Value = serde_json::from_slice(&body).unwrap();
    let responses = response["responses"].as_array().unwrap();
    assert_eq!(responses.len(), 2);

    let first = BASE64
        .decode(
            responses[0]["response"]["outputs"][0]["data"]
                .as_str()
                .unwrap(),
        )
        .unwrap();
    // affine: y = 2x + 1
    assert_eq!(f32_values(&first), vec![3.0, 5.0]);
    let second = BASE64
        .decode(
            responses[1]["response"]["outputs"][0]["data"]
                .as_str()
                .unwrap(),
        )
        .unwrap();
    assert_eq!(f32_values(&second), vec![7.0]);

    let _ = shutdown_tx.send(());
    server.await.unwrap().unwrap();
}

#[tokio::test]
async fn quic_transport_unknown_path_returns_404() {
    let (addr, shutdown_tx, server) = spawn_server();

    let (_endpoint, connection) = connect_with_retry(addr).await;
    let (status, body) = send_request(connection, "GET", "/does/not/exist", None).await;
    assert_eq!(status, axum::http::StatusCode::NOT_FOUND);

    let response: Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(response["code"], "not_found");

    let _ = shutdown_tx.send(());
    server.await.unwrap().unwrap();
}
