unb-server 1.0.0

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use unb_core::RetirementReason;
use unb_server::{
    ConnectError, ConnectionStatus, DisconnectReason, Endpoint, EndpointSet, HostConfig, Node,
    PeerConnection, TcpTransport, TransportKind,
};

fn assert_clone_send_sync<T: Clone + Send + Sync>() {}

fn websocket_endpoint(address: String) -> Endpoint {
    Endpoint {
        kind: TransportKind::WebSocket,
        address,
        cert_hash: None,
    }
}

async fn hosted(name: &str) -> ((Arc<Node>, unb_server::Hosting), String) {
    let node = Node::builder(name)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    host_node(node).await
}

async fn host_node(node: Arc<Node>) -> ((Arc<Node>, unb_server::Hosting), String) {
    let hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
        .start(&node)
        .await
        .unwrap();
    let url = format!("ws://{}", hosting.websocket_addr().unwrap());
    ((node, hosting), url)
}

#[test]
fn peer_connection_is_clone_send_and_sync() {
    assert_clone_send_sync::<PeerConnection>();
}

#[tokio::test(flavor = "multi_thread")]
async fn connect_returns_a_connected_peer_handle() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();

    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();

    assert_eq!(connection.peer(), "peer");
    assert_eq!(connection.status(), ConnectionStatus::Connected);
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn failed_initial_connect_does_not_poison_a_later_connection() {
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();

    let error = match node.connect(EndpointSet::new()).await {
        Ok(_) => panic!("empty endpoint set unexpectedly connected"),
        Err(error) => error,
    };
    assert_eq!(error, ConnectError::NoSupportedEndpoint);

    let (server, url) = hosted("peer").await;
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    assert_eq!(connection.peer(), "peer");
    assert_eq!(connection.status(), ConnectionStatus::Connected);
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn clones_observe_the_same_status_transition() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    let observer = connection.clone();
    let changed = tokio::spawn(observer.changed());

    connection.disconnect();

    assert_eq!(
        tokio::time::timeout(Duration::from_secs(1), changed)
            .await
            .unwrap()
            .unwrap(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ExplicitDisconnect,
        }
    );
    assert_eq!(
        connection.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ExplicitDisconnect,
        }
    );
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn repeated_connects_converge_on_one_logical_connection() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let endpoints = EndpointSet::from(websocket_endpoint(url));

    let first = node.connect(endpoints.clone()).await.unwrap();
    let second = node.connect(endpoints).await.unwrap();
    first.disconnect();

    assert_eq!(
        second.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ExplicitDisconnect,
        }
    );
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn connecting_a_different_peer_keeps_connections_independent() {
    let (first_server, first_url) = hosted("peer-a").await;
    let (second_server, second_url) = hosted("peer-b").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();

    let first = node
        .connect(EndpointSet::from(websocket_endpoint(first_url)))
        .await
        .unwrap();
    let second = node
        .connect(EndpointSet::from(websocket_endpoint(second_url)))
        .await
        .unwrap();
    second.disconnect();

    assert_eq!(first.peer(), "peer-a");
    assert_eq!(first.status(), ConnectionStatus::Connected);
    assert_eq!(second.peer(), "peer-b");
    assert_eq!(
        second.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ExplicitDisconnect,
        }
    );
    drop(first_server);
    drop(second_server);
}

#[tokio::test(flavor = "multi_thread")]
async fn selected_session_retirement_publishes_its_typed_reason() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    let changed = connection.changed();

    server.0.shutdown();

    let status = tokio::time::timeout(Duration::from_secs(1), changed)
        .await
        .unwrap();
    assert!(matches!(
        status,
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::SessionRetired {
                reason: RetirementReason::SessionClosed
                    | RetirementReason::SendClosed
                    | RetirementReason::SendTransport
                    | RetirementReason::TransportFailed
            }
        }
    ));
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn node_shutdown_is_visible_to_every_connection_clone() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    let observer = connection.clone();
    let changed = observer.changed();

    node.shutdown();

    assert_eq!(
        tokio::time::timeout(Duration::from_secs(1), changed)
            .await
            .unwrap(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::NodeShutdown,
        }
    );
    assert_eq!(
        connection.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::NodeShutdown,
        }
    );
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn explicit_reconnect_runs_one_sweep_and_restores_connected_status() {
    let (server, url) = hosted("peer").await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    connection.disconnect();

    connection.reconnect().await.unwrap();

    assert_eq!(connection.status(), ConnectionStatus::Connected);
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn disconnect_rejects_a_late_reconnect_candidate() {
    let admissions = Arc::new(AtomicUsize::new(0));
    let release = Arc::new(tokio::sync::Notify::new());
    let peer = Node::builder("peer")
        .peer_layer_fn({
            let admissions = admissions.clone();
            let release = release.clone();
            move |mut request, next| {
                let admissions = admissions.clone();
                let release = release.clone();
                async move {
                    if admissions.fetch_add(1, Ordering::SeqCst) > 0 {
                        release.notified().await;
                    }
                    request.accept_declared();
                    next.admit(request).await
                }
            }
        })
        .build()
        .unwrap();
    let (server, url) = host_node(peer).await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    connection.disconnect();
    let reconnecting = connection.changed();
    let reconnect = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().await }
    });
    assert_eq!(reconnecting.await, ConnectionStatus::Connecting);
    tokio::time::timeout(Duration::from_secs(1), async {
        while admissions.load(Ordering::SeqCst) < 2 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    assert_eq!(connection.status(), ConnectionStatus::Connecting);

    connection.disconnect();
    release.notify_waiters();

    assert_eq!(reconnect.await.unwrap(), Err(ConnectError::Cancelled));
    tokio::time::sleep(Duration::from_millis(20)).await;
    assert_eq!(
        connection.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ExplicitDisconnect,
        }
    );
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn concurrent_reconnect_callers_share_one_failed_attempt() {
    let admissions = Arc::new(AtomicUsize::new(0));
    let reject = Arc::new(tokio::sync::Notify::new());
    let peer = Node::builder("peer")
        .peer_layer_fn({
            let admissions = admissions.clone();
            let reject = reject.clone();
            move |mut request, next| {
                let admissions = admissions.clone();
                let reject = reject.clone();
                async move {
                    if admissions.fetch_add(1, Ordering::SeqCst) > 0 {
                        reject.notified().await;
                        return Err(unb_server::HandlerError::new(
                            unb_core::ErrorCode::Unauthorized,
                            "reconnect rejected",
                        ));
                    }
                    request.accept_declared();
                    next.admit(request).await
                }
            }
        })
        .build()
        .unwrap();
    let (server, url) = host_node(peer).await;
    let node = Node::builder("caller")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from(websocket_endpoint(url)))
        .await
        .unwrap();
    connection.disconnect();

    let first = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().await }
    });
    tokio::time::timeout(Duration::from_secs(1), async {
        while admissions.load(Ordering::SeqCst) < 2 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    let second = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().await }
    });
    let third = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().await }
    });
    tokio::task::yield_now().await;
    reject.notify_waiters();

    let first = first.await.unwrap();
    let second = second.await.unwrap();
    let third = third.await.unwrap();
    assert_eq!(first, second);
    assert_eq!(second, third);
    assert!(first.is_err());
    assert_eq!(admissions.load(Ordering::SeqCst), 2);
    assert!(matches!(
        connection.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ReconnectFailed { .. }
        }
    ));

    let later = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().await }
    });
    tokio::time::timeout(Duration::from_secs(1), async {
        while admissions.load(Ordering::SeqCst) < 3 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    reject.notify_waiters();
    assert!(later.await.unwrap().is_err());
    assert_eq!(admissions.load(Ordering::SeqCst), 3);
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn reconnect_accepts_a_new_instance_of_the_same_node_id() {
    let (first, first_url) = hosted("peer").await;
    let (second, second_url) = hosted("peer").await;
    let node = Node::builder("caller")
        .connect_timeout(Duration::from_secs(1))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from([
            websocket_endpoint(first_url),
            websocket_endpoint(second_url),
        ]))
        .await
        .unwrap();
    let disconnected = connection.changed();

    first.0.shutdown();
    drop(first);
    assert!(matches!(
        tokio::time::timeout(Duration::from_secs(1), disconnected)
            .await
            .unwrap(),
        ConnectionStatus::Disconnected { .. }
    ));

    connection.reconnect().await.unwrap();
    assert_eq!(connection.peer(), "peer");
    assert_eq!(connection.status(), ConnectionStatus::Connected);
    drop(second);
}

#[tokio::test(flavor = "multi_thread")]
async fn reconnect_rejects_an_endpoint_for_a_different_node_id() {
    let (first, first_url) = hosted("peer").await;
    let (other, other_url) = hosted("other").await;
    let node = Node::builder("caller")
        .connect_timeout(Duration::from_secs(1))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let connection = node
        .connect(EndpointSet::from([
            websocket_endpoint(first_url),
            websocket_endpoint(other_url),
        ]))
        .await
        .unwrap();
    let disconnected = connection.changed();
    first.0.shutdown();
    drop(first);
    let _ = tokio::time::timeout(Duration::from_secs(1), disconnected)
        .await
        .unwrap();

    let error = connection.reconnect().await.unwrap_err();

    assert_eq!(
        error,
        ConnectError::IdentityMismatch {
            expected: "peer".into(),
            actual: Some("other".into()),
        }
    );
    assert_eq!(
        connection.status(),
        ConnectionStatus::Disconnected {
            reason: DisconnectReason::ReconnectFailed { error },
        }
    );
    drop(other);
}