use super::*;
use futures_util::{SinkExt, StreamExt};
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::oneshot;

#[test]
fn websocket_client_scope_ends_all_aliases_and_timeout_stream_fails_once() {
    runtime()
        .run(async |scope| {
            let (listener, server) = server(scope).await;
            let port = listener.port().unwrap();
            let owner = scope
                .spawn(move |child| async move {
                    ws::connect("127.0.0.1", port, "/", &limits(), &child)
                        .await
                        .unwrap()
                })
                .unwrap();
            let escaped = owner.wait().await.unwrap();
            assert!(escaped.send(Message::Text("closed".into())).await.is_err());
            let socket = ws::connect(
                "127.0.0.1",
                port,
                "/",
                &ws::Limits {
                    idle_ms: 25,
                    ..limits()
                },
                scope,
            )
            .await
            .unwrap();
            let stream = socket.messages(|message| message);
            let alias = stream.clone();
            let failure = stream.next().await.unwrap().unwrap().unwrap_err();
            assert!(failure.contains("timed out"));
            assert!(alias.next().await.unwrap().is_none());
            assert!(socket.send(Message::Text("closed".into())).await.is_err());
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_listener_drain_keeps_session_and_parent_stop_cleans_handler() {
    runtime()
        .run(async |scope| {
            let (listener, mut server) = server(scope).await;
            let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
                .await
                .unwrap();
            listener.close().unwrap();
            tokio::select! {
                _ = &mut server => panic!("server must retain its upgraded session"),
                () = tokio::time::sleep(Duration::from_millis(25)) => (),
            }
            socket
                .send(Message::Text("after listener close".into()))
                .await
                .unwrap();
            assert_eq!(
                socket.receive().await.unwrap(),
                Some(Message::Text("after listener close".into()))
            );
            socket.close(1000, "").await.unwrap();
            server.wait().await.unwrap().unwrap();

            let listener = net::listen("127.0.0.1", 0, scope).await.unwrap();
            let incoming = listener.clone();
            let active = Arc::new(AtomicUsize::new(0));
            let running = active.clone();
            let (sender, receiver) = oneshot::channel::<WebSocket>();
            let sender = Arc::new(std::sync::Mutex::new(Some(sender)));
            let server = scope
                .spawn(move |child| async move {
                    http::serve_live(
                        &incoming,
                        &Limits::default(),
                        &http::LiveLimits::default(),
                        &child,
                        move |_, reply, scope| {
                            let sender = sender.clone();
                            let running = running.clone();
                            async move {
                                let socket = ws::accept(&reply, &limits(), &scope).await.unwrap();
                                running.fetch_add(1, Ordering::SeqCst);
                                let _running = Running(running);
                                let _ = sender.lock().unwrap().take().unwrap().send(socket.clone());
                                let _ = socket.receive().await;
                                Ok(())
                            }
                        },
                    )
                    .await
                })
                .unwrap();
            let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
                .await
                .unwrap();
            let escaped = receiver.await.unwrap();
            server.stop().await.unwrap();
            assert_eq!(active.load(Ordering::SeqCst), 0);
            assert!(escaped.send(Message::Text("gone".into())).await.is_err());
            assert!(socket.receive().await.is_err());
            listener.close().unwrap();
        })
        .unwrap();
}

struct Running(Arc<AtomicUsize>);
impl Drop for Running {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::SeqCst);
    }
}

#[test]
fn websocket_owner_stop_wakes_receive_in_another_task() {
    runtime()
        .run(async |scope| {
            let (listener, server) = server(scope).await;
            let port = listener.port().unwrap();
            let (sender, receiver) = oneshot::channel();
            let owner = scope
                .spawn(move |child| async move {
                    let socket = ws::connect("127.0.0.1", port, "/", &limits(), &child)
                        .await
                        .unwrap();
                    sender.send(socket).unwrap();
                    std::future::pending::<()>().await;
                })
                .unwrap();
            let socket = receiver.await.unwrap();
            let (started, ready) = oneshot::channel();
            let reader = scope
                .spawn_leaf(move || async move {
                    let receiving = socket.receive();
                    tokio::pin!(receiving);
                    assert!(futures_util::poll!(&mut receiving).is_pending());
                    started.send(()).unwrap();
                    tokio::time::timeout(Duration::from_millis(200), receiving)
                        .await
                        .unwrap()
                        .unwrap_err()
                })
                .unwrap();
            ready.await.unwrap();
            owner.stop().await.unwrap();
            assert!(reader.wait().await.unwrap().contains("WebSocket is closed"));
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_closing_message_stream_preserves_independent_socket_alias() {
    runtime()
        .run(async |scope| {
            let (listener, server) = server(scope).await;
            let socket = ws::connect("127.0.0.1", listener.port().unwrap(), "/", &limits(), scope)
                .await
                .unwrap();
            let stream = socket.messages(|message| message);
            stream.close();
            assert!(stream.next().await.unwrap().is_none());
            socket
                .send(Message::Binary(vec![1, 2].into()))
                .await
                .unwrap();
            assert_eq!(
                socket.receive().await.unwrap(),
                Some(Message::Binary(vec![1, 2].into()))
            );
            socket.close(1000, "").await.unwrap();
            listener.close().unwrap();
            server.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_cancelled_writer_closes_only_after_it_owns_the_direction() {
    runtime()
        .run(async |scope| {
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let port = listener.local_addr().unwrap().port();
            let (started, receiving) = oneshot::channel();
            let (release, released) = oneshot::channel();
            let peer = scope
                .spawn_leaf(move || async move {
                    let (stream, _) = listener.accept().await.unwrap();
                    let wire = tokio_tungstenite::accept_async(stream).await.unwrap();
                    wire.get_ref().readable().await.unwrap();
                    let _ = started.send(());
                    let _ = released.await;
                })
                .unwrap();
            let socket = ws::connect(
                "127.0.0.1",
                i64::from(port),
                "/",
                &ws::Limits {
                    message_bytes: 16 * 1024 * 1024,
                    write_ms: 1000,
                    ..limits()
                },
                scope,
            )
            .await
            .unwrap();
            let sender = socket.clone();
            let mut writing = scope
                .spawn_leaf(move || async move {
                    sender
                        .send(Message::Binary(vec![1; 16 * 1024 * 1024].into()))
                        .await
                })
                .unwrap();
            receiving.await.unwrap();
            let waiting = socket.clone();
            let waiter = scope
                .spawn_leaf(
                    move || async move { waiting.send(Message::Text("queued".into())).await },
                )
                .unwrap();
            tokio::time::sleep(Duration::from_millis(10)).await;
            waiter.stop().await.unwrap();
            assert!(
                tokio::time::timeout(Duration::from_millis(10), &mut writing)
                    .await
                    .is_err()
            );
            writing.stop().await.unwrap();
            assert!(socket.send(Message::Text("closed".into())).await.is_err());
            release.send(()).unwrap();
            peer.wait().await.unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_close_keeps_the_handshake_when_a_reader_is_waiting() {
    runtime()
        .run(async |scope| {
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let port = listener.local_addr().unwrap().port();
            let (replied, reply) = oneshot::channel();
            let peer = scope
                .spawn_leaf(move || async move {
                    let (stream, _) = listener.accept().await.unwrap();
                    let mut wire = tokio_tungstenite::accept_async(stream).await.unwrap();
                    assert!(matches!(
                        wire.next().await,
                        Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_)))
                    ));
                    wire.flush().await.unwrap();
                    replied.send(()).unwrap();
                })
                .unwrap();
            let socket = ws::connect("127.0.0.1", i64::from(port), "/", &limits(), scope)
                .await
                .unwrap();
            let receiving = socket.receive();
            tokio::pin!(receiving);
            assert!(futures_util::poll!(&mut receiving).is_pending());
            let closing = socket.close(1000, "");
            tokio::pin!(closing);
            assert!(futures_util::poll!(&mut closing).is_pending());
            reply.await.unwrap();
            tokio::time::sleep(Duration::from_millis(25)).await;
            let received = receiving.await;
            closing.await.unwrap();
            assert!(received.unwrap_err().contains("WebSocket is closing"));
            peer.wait().await.unwrap();
        })
        .unwrap();
}

#[test]
fn websocket_close_deadline_releases_all_aliases() {
    runtime()
        .run(async |scope| {
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let port = listener.local_addr().unwrap().port();
            let (release, released) = oneshot::channel();
            let peer = scope
                .spawn_leaf(move || async move {
                    let (stream, _) = listener.accept().await.unwrap();
                    let _wire = tokio_tungstenite::accept_async(stream).await.unwrap();
                    let _ = released.await;
                })
                .unwrap();
            let socket = ws::connect(
                "127.0.0.1",
                i64::from(port),
                "/",
                &ws::Limits {
                    write_ms: 50,
                    ..limits()
                },
                scope,
            )
            .await
            .unwrap();
            let alias = socket.clone();
            let receiving = alias.receive();
            tokio::pin!(receiving);
            assert!(futures_util::poll!(&mut receiving).is_pending());
            let (closed, received) = tokio::join!(socket.close(1000, ""), receiving);
            assert!(closed.unwrap_err().contains("timed out"));
            assert!(received.unwrap_err().contains("WebSocket is closing"));
            assert!(
                alias
                    .send(Message::Text("closed".into()))
                    .await
                    .unwrap_err()
                    .contains("WebSocket is closed")
            );
            release.send(()).unwrap();
            peer.wait().await.unwrap();
        })
        .unwrap();
}