use std::sync::{Arc, Mutex};

use boring::ssl::{select_next_proto, AlpnError, SslAcceptor};
use specter::{Client, CookieJar, Message};
use tokio::sync::RwLock;

#[path = "helpers/mock_ws_server.rs"]
mod mock_ws_server;
#[path = "helpers/tls.rs"]
mod tls;

use mock_ws_server::{server_text_frame, AcceptMode, MockWsServer, WsResponse};
use tls::generate_cert_bundle;

#[tokio::test]
async fn valid_ws_handshake_sends_required_headers_and_16_byte_key() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/chat?room=blue");
    let expected_host = url
        .strip_prefix("ws://")
        .unwrap()
        .strip_suffix("/chat?room=blue")
        .unwrap()
        .to_string();
    let handle = server.start_once(WsResponse::default());

    let _ws = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect("websocket handshake should succeed");

    let exchange = handle.await.unwrap();
    let request = exchange.request;

    assert_eq!(request.request_line, "GET /chat?room=blue HTTP/1.1");
    assert_eq!(request.header("Host"), Some(expected_host.as_str()));
    assert_eq!(request.header("Upgrade"), Some("websocket"));
    assert!(request
        .header("Connection")
        .unwrap_or_default()
        .split(',')
        .any(|token| token.trim().eq_ignore_ascii_case("upgrade")));
    assert_eq!(request.header("Sec-WebSocket-Version"), Some("13"));
    assert_eq!(request.sec_websocket_key_len(), 16);
    assert_eq!(request.header("Sec-WebSocket-Extensions"), None);
}

#[tokio::test]
async fn explicit_ws_port_is_preserved_in_host_header() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/port");
    let expected_host = url
        .strip_prefix("ws://")
        .unwrap()
        .strip_suffix("/port")
        .unwrap()
        .to_string();
    let handle = server.start_once(WsResponse::default());

    let _ws = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect("websocket handshake should succeed");

    let exchange = handle.await.unwrap();
    assert_eq!(
        exchange.request.header("Host"),
        Some(expected_host.as_str())
    );
}

#[tokio::test]
async fn wrong_accept_header_fails_the_handshake() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/bad-accept");
    let handle = server.start_once(WsResponse {
        accept: AcceptMode::Wrong,
        ..WsResponse::default()
    });

    let err = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect_err("invalid Sec-WebSocket-Accept must fail");

    let debug = format!("{err:?}");
    assert!(
        debug.contains("InvalidAccept") || debug.contains("Sec-WebSocket-Accept"),
        "unexpected error: {debug}"
    );
    assert!(handle.await.unwrap().client_frame.is_none());
}

#[tokio::test]
async fn unexpected_extension_header_fails_the_handshake() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/extension");
    let handle = server.start_once(WsResponse {
        headers: vec![(
            "Sec-WebSocket-Extensions".to_string(),
            "permessage-deflate".to_string(),
        )],
        ..WsResponse::default()
    });

    let err = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect_err("unexpected extension must fail");

    let debug = format!("{err:?}");
    assert!(
        debug.contains("UnexpectedExtension") || debug.contains("Sec-WebSocket-Extensions"),
        "unexpected error: {debug}"
    );
    assert!(handle.await.unwrap().client_frame.is_none());
}

#[tokio::test]
async fn unexpected_subprotocol_header_fails_when_none_was_offered() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/subprotocol");
    let handle = server.start_once(WsResponse {
        headers: vec![("Sec-WebSocket-Protocol".to_string(), "chat".to_string())],
        ..WsResponse::default()
    });

    let err = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect_err("unoffered subprotocol must fail");

    let debug = format!("{err:?}");
    assert!(
        debug.contains("UnexpectedSubprotocol") || debug.contains("Sec-WebSocket-Protocol"),
        "unexpected error: {debug}"
    );
    assert!(handle.await.unwrap().client_frame.is_none());
}

#[tokio::test]
async fn offered_subprotocol_is_accepted_when_echoed_exactly() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/subprotocol-ok");
    let handle = server.start_once(WsResponse {
        headers: vec![("Sec-WebSocket-Protocol".to_string(), "chat.v2".to_string())],
        ..WsResponse::default()
    });

    let ws = Client::new()
        .unwrap()
        .websocket(url)
        .subprotocol("chat.v2")
        .connect()
        .await
        .expect("offered subprotocol should succeed");

    assert_eq!(ws.protocol(), Some("chat.v2"));
    let exchange = handle.await.unwrap();
    assert_eq!(
        exchange.request.header("Sec-WebSocket-Protocol"),
        Some("chat.v2")
    );
}

#[tokio::test]
async fn first_frame_sent_with_101_is_available_to_next() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.ws_url("/early-frame");
    let handle = server.start_once(WsResponse {
        first_frame: Some(server_text_frame("ready")),
        ..WsResponse::default()
    });

    let mut ws = Client::new()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect("websocket handshake should succeed");

    let message = ws
        .next()
        .await
        .expect("read first frame")
        .expect("message available");
    assert!(matches!(message, Message::Text(text) if text == "ready"));

    let _ = handle.await.unwrap();
}

#[tokio::test]
async fn wss_handshake_offers_only_http1_alpn() {
    let server = MockWsServer::new().await.unwrap();
    let url = server.wss_url("/secure-alpn");
    let (acceptor, ca_cert, client_alpn) = h1_tls_acceptor();
    let handle = server.start_tls_once(acceptor, WsResponse::default());

    let ws = Client::builder()
        .add_root_certificate(ca_cert)
        .localhost_allows_invalid_certs(false)
        .build()
        .unwrap()
        .websocket(url)
        .connect()
        .await
        .expect("wss websocket handshake should succeed");
    drop(ws);

    let exchange = handle.await.unwrap();
    assert_eq!(&*client_alpn.lock().unwrap(), b"\x08http/1.1");
    assert_eq!(
        exchange.selected_alpn.as_deref(),
        Some(b"http/1.1".as_ref())
    );
}

#[tokio::test]
async fn wss_set_cookie_on_101_is_stored_and_replayed() {
    let cookie_jar = Arc::new(RwLock::new(CookieJar::new()));
    let (acceptor1, ca_cert1, _) = h1_tls_acceptor();
    let (acceptor2, ca_cert2, _) = h1_tls_acceptor();

    let server1 = MockWsServer::new().await.unwrap();
    let url1 = server1.wss_url("/set-cookie");
    let handle1 = server1.start_tls_once(
        acceptor1,
        WsResponse {
            headers: vec![(
                "Set-Cookie".to_string(),
                "sid=abc; Secure; Path=/".to_string(),
            )],
            ..WsResponse::default()
        },
    );

    let client = Client::builder()
        .add_root_certificate(ca_cert1)
        .add_root_certificate(ca_cert2)
        .cookie_jar(cookie_jar)
        .build()
        .unwrap();

    let ws = client
        .websocket(url1)
        .connect()
        .await
        .expect("first wss handshake should store cookie");
    drop(ws);
    let _ = handle1.await.unwrap();

    let server2 = MockWsServer::new().await.unwrap();
    let url2 = server2.wss_url("/uses-cookie");
    let handle2 = server2.start_tls_once(acceptor2, WsResponse::default());

    let ws = client
        .websocket(url2)
        .connect()
        .await
        .expect("second wss handshake should send cookie");
    drop(ws);

    let exchange = handle2.await.unwrap();
    assert_eq!(exchange.request.header("Cookie"), Some("sid=abc"));
}

#[tokio::test]
async fn explicit_wss_cookie_header_is_not_overwritten_by_cookie_jar() {
    let cookie_jar = Arc::new(RwLock::new(CookieJar::new()));
    let (acceptor1, ca_cert1, _) = h1_tls_acceptor();
    let (acceptor2, ca_cert2, _) = h1_tls_acceptor();

    let server1 = MockWsServer::new().await.unwrap();
    let url1 = server1.wss_url("/set-cookie");
    let handle1 = server1.start_tls_once(
        acceptor1,
        WsResponse {
            headers: vec![(
                "Set-Cookie".to_string(),
                "sid=abc; Secure; Path=/".to_string(),
            )],
            ..WsResponse::default()
        },
    );

    let client = Client::builder()
        .add_root_certificate(ca_cert1)
        .add_root_certificate(ca_cert2)
        .cookie_jar(cookie_jar)
        .build()
        .unwrap();

    let ws = client
        .websocket(url1)
        .connect()
        .await
        .expect("first wss handshake should store cookie");
    drop(ws);
    let _ = handle1.await.unwrap();

    let server2 = MockWsServer::new().await.unwrap();
    let url2 = server2.wss_url("/manual-cookie");
    let handle2 = server2.start_tls_once(acceptor2, WsResponse::default());

    let ws = client
        .websocket(url2)
        .header("Cookie", "manual=1")
        .connect()
        .await
        .expect("second wss handshake should respect explicit cookie header");
    drop(ws);

    let exchange = handle2.await.unwrap();
    assert_eq!(exchange.request.header("Cookie"), Some("manual=1"));
}

fn h1_tls_acceptor() -> (SslAcceptor, Vec<u8>, Arc<Mutex<Vec<u8>>>) {
    let (mut builder, ca_cert) = generate_cert_bundle();
    let client_alpn = Arc::new(Mutex::new(Vec::new()));
    let captured_alpn = Arc::clone(&client_alpn);

    builder.set_alpn_select_callback(move |_, client_protos| {
        *captured_alpn.lock().unwrap() = client_protos.to_vec();
        select_next_proto(b"\x08http/1.1", client_protos).ok_or(AlpnError::NOACK)
    });

    (builder.build(), ca_cert, client_alpn)
}
