//! The WS upgrade route (`/alk/channels`, ADR-067) and the channels
//! session it establishes — the server producer half.
//!
//! Per `docs/architecture/websocket.md`: bearer auth on the upgrade
//! request (`401` without a resolvable token — the same
//! `resolve_from_token` path as any HTTP request), then upgrade →
//! WS↔byte-stream adapter → `Connection::from_bidi(ws_stream,
//! b"alk/channels")` → alkcall `ChannelsAdapter` (the channels accept
//! path). The `install_channel_zero` hook constructs channel 0's
//! `CallConnection` (the identity rides the channels-layer
//! `Connection`) and runs the shared `Dispatcher::run_loop_single_stream`.
//!
//! ## Detached task lifetime semantics (WS-10)
//!
//! Two task families spawned here are **detached by design** and
//! outlive the session task that spawned them:
//!
//! - the WS pump tasks (the `WsPumps` pair behind the byte adapter,
//!   re-exported from `crate::websocket`), and
//! - the channel-0 dispatcher task spawned by `install_channel_zero`
//!   per accepted channels connection.
//!
//! If the channels session task (`run_channels_session`) dies —
//! upgrade-time early return, an adapter fault, or its own task being
//! cancelled — these tasks keep running **self-healing**: each ends on
//! its own when its stream half closes (peer disconnect, peer close
//! frame, read error, or the local teardown arms —
//! `WsSessions::abort`, the idle-read timeout, `AsyncWrite::shutdown`).
//! They are never leaked unconditionally: the leak window is bounded
//! by peer behavior (a peer that holds the socket open keeps the pump
//! and dispatcher tasks alive with it) and additionally by the WS-01
//! idle-read timeout when configured, and every session's pumps stay
//! force-evictable through the WS-08 registry
//! ([`WsSessions::abort`]) for the session's whole lifetime. The
//! dispatcher task specifically ends when channel 0's `BiStream` read
//! side hits EOF (its only exit condition), i.e. when the underlying
//! WS connection ends by any of the paths above.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use alkcall::channels::adapter::ChannelsAdapter;
use alkcall::channels::policy::{ChannelLifecyclePolicy, NoCap};
use alkcall::core::auth::{AuthContext, Identity};
use alkcall::core::types::{Connection, ProtocolHandler};
use alkcall::registry::registration::OperationRegistry;
use axum::extract::ws::WebSocketUpgrade;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use parking_lot::Mutex;

use super::byte_adapter::{
    split_ws_to_bytes_idle_with_write, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
};

/// Registry of live WS session pump handles (WS-08). The upgrade
/// handler retains each session's [`WsPumps`] handle here for the
/// session's lifetime so a stuck session is evictable in-crate:
/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket
/// halves close and the channels session unwinds). The handle for a
/// session is removed when the session task finishes (self-removing
/// guard), so the registry only holds live sessions. Assembly layers
/// share one instance via the adapter's `RouterState` (the upgrade
/// handler registers against it) or per-route request extensions (an
/// extension clone takes precedence).
///
/// Un-registered (default) — the upgrade runs fine and simply keeps no
/// eviction lever, matching the pre-WS-08 behavior.
#[derive(Clone, Default)]
pub struct WsSessions {
    counter: Arc<AtomicU64>,
    sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
}

/// Default bound on concurrent WS sessions (WS-09): the semaphore
/// initial capacity in [`SessionState`]; a deployment overrides it with
/// `HttpAdapter::with_ws_max_sessions`.
pub const DEFAULT_WS_MAX_SESSIONS: usize = 64;

/// The upgrade handler's state slice: what it needs beyond the request
/// itself. Axum lifts it via `FromRef` from either full router state —
/// the adapter's `RouterState` (carrying the shared [`WsSessions`]
/// instance and the configured session cap) or a bare
/// `Arc<OperationRegistry>` (custom upgrade routes / integration tests
/// get a handler-private registry and the default cap; eviction still
/// works in-crate, just not shared).
#[derive(Clone)]
pub struct SessionState {
    registry: Arc<OperationRegistry>,
    sessions: Arc<WsSessions>,
    /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller
    /// over the cap is rejected with 503.
    session_slots: Arc<tokio::sync::Semaphore>,
    /// Idle-read timeout (WS-01): `None` disables the knob.
    idle_timeout: Option<std::time::Duration>,
}

impl SessionState {
    /// From the plain registry state (custom upgrade routes /
    /// integration tests): sessions default to a fresh [`WsSessions`]
    /// private to the handler (eviction still functional in-crate but
    /// not shared with the assembly layer) and the default session cap.
    pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
        Self {
            registry: Arc::clone(registry),
            sessions: Arc::new(WsSessions::new()),
            session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)),
            idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
        }
    }

    /// From the adapter's router state: the shared [`WsSessions`]
    /// instance the assembly layer can hold for eviction, and the
    /// pre-built session-cap semaphore (one per `HttpAdapter`, shared
    /// across requests).
    pub(crate) fn new(
        registry: Arc<OperationRegistry>,
        sessions: Arc<WsSessions>,
        session_slots: Arc<tokio::sync::Semaphore>,
        idle_timeout: Option<std::time::Duration>,
    ) -> Self {
        Self {
            registry,
            sessions,
            session_slots,
            idle_timeout,
        }
    }

    pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
        &self.registry
    }

    pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
        &self.sessions
    }

    pub(crate) fn idle_timeout(&self) -> Option<std::time::Duration> {
        self.idle_timeout
    }
}

/// `FromRef` chain: a bare `Arc<OperationRegistry>` router state lifts
/// into the handler's [`SessionState`]; a full `RouterState` carries
/// the shared [`WsSessions`] instance and lifts through its own impl.
impl axum::extract::FromRef<Arc<OperationRegistry>> for SessionState {
    fn from_ref(registry: &Arc<OperationRegistry>) -> Self {
        SessionState::from_registry(registry)
    }
}

impl WsSessions {
    /// A fresh, empty session registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Abort every live session's pump tasks (forced teardown).
    pub fn abort(&self) {
        for (_, pumps) in self.sessions.lock().drain() {
            pumps.abort();
        }
    }

    /// Number of live sessions currently tracked.
    pub fn len(&self) -> usize {
        self.sessions.lock().len()
    }

    /// Whether no sessions are tracked.
    pub fn is_empty(&self) -> bool {
        self.sessions.lock().is_empty()
    }

    fn insert(&self, pumps: Arc<WsPumps>) -> u64 {
        let id = self.counter.fetch_add(1, Ordering::Relaxed);
        self.sessions.lock().insert(id, pumps);
        id
    }

    fn remove(&self, id: u64) {
        self.sessions.lock().remove(&id);
    }
}

/// The channels session for an upgraded socket: adapt → `Connection`
/// (identity attached) → `ChannelsAdapter::handle`. `policy` gates
/// data-channel opens (ADR-041). When `sessions` is `Some`, the
/// session's pump handle is registered for the session's lifetime —
/// the WS-08 eviction lever ([`WsSessions::abort`]). `idle_timeout`
/// bounds the read stall (WS-01): `None` disables the knob, `Some(d)`
/// closes the read with 1001 after `d` without inbound chunk
/// progress. `write_timeout` bounds one outbound WS send (WS-18):
/// `None` = the crate default window, `Some(d)` a deployment-set
/// window.
#[allow(clippy::too_many_arguments)]
pub async fn run_channels_session(
    socket: axum::extract::ws::WebSocket,
    registry: Arc<OperationRegistry>,
    identity: Identity,
    policy: Arc<dyn ChannelLifecyclePolicy>,
    sessions: Option<WsSessions>,
    idle_timeout: Option<std::time::Duration>,
    write_timeout: Option<std::time::Duration>,
) {
    let (byte_stream, pumps) =
        split_ws_to_bytes_idle_with_write(socket, idle_timeout, write_timeout);
    let pumps = Arc::new(pumps);

    let _guard = sessions.map(|sessions| {
        let id = sessions.insert(Arc::clone(&pumps));
        SessionGuard { sessions, id }
    });

    let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None);
    let _ = conn.set_identity(identity.clone());

    let adapter = ChannelsAdapter::new(install_channel_zero(registry), policy);
    let auth = AuthContext {
        identity: Some(identity),
        alpn: b"alk/channels".to_vec(),
        remote_addr: None,
        tls_client_fingerprint: None,
    };
    if let Err(e) = ProtocolHandler::handle(&adapter, conn, &auth).await {
        tracing::warn!(error = %e, "channels session ended");
    }
}

struct SessionGuard {
    sessions: WsSessions,
    id: u64,
}

impl Drop for SessionGuard {
    fn drop(&mut self) {
        self.sessions.remove(self.id);
    }
}

/// The `install_channel_zero` hook: split channel 0's `BiStream` into
/// the shared writer + reader, construct the single-stream
/// `CallConnection`, and run the shared `Dispatcher`'s single-stream
/// loop (the alkcall `channels/client.rs` accept-side wiring pattern).
///
/// The dispatcher's token resolver is a no-op: the WS identity is
/// attached to the connection at upgrade time and
/// `Dispatcher::resolve_identity` falls back to it when the payload
/// carries no `auth_token` — the correct accept-side behavior (the
/// identity was established at upgrade time, not per-call).
fn install_channel_zero(
    registry: Arc<OperationRegistry>,
) -> alkcall::channels::adapter::InstallChannelZero {
    Arc::new(move |_manager, channel0_conn, auth| {
        let registry = Arc::clone(&registry);
        tokio::spawn(async move {
            // The WS identity rides the upgrade request; propagate it to
            // channel 0's `CallConnection` so the dispatcher's
            // `resolve_identity` (and thus `AccessControl::check` on
            // `services/list` and every operation) sees it. Without this
            // the freshly constructed channel-0 `Connection` carries no
            // identity and all ACL-restricted ops look unauthenticated.
            if let Some(identity) = auth.identity.clone() {
                let _ = channel0_conn.set_identity(identity);
            }
            let channel0_bidi = match channel0_conn.accept_bi().await {
                Ok(s) => s,
                Err(_) => return,
            };
            let (writer, reader) =
                alkcall::protocol::connection::split_single_stream(channel0_bidi);
            let call_connection = Arc::new(
                alkcall::protocol::connection::CallConnection::new_single_stream(
                    channel0_conn,
                    Arc::clone(&writer),
                ),
            );
            let dispatcher = alkcall::protocol::dispatch::Dispatcher::new(
                registry,
                std::sync::Arc::new(NoopProvider),
            );
            dispatcher
                .run_loop_single_stream(call_connection, reader, writer)
                .await;
        })
    })
}

#[allow(dead_code)]
pub(crate) fn adapter_install_channel_zero(
    registry: Arc<OperationRegistry>,
) -> alkcall::channels::adapter::InstallChannelZero {
    install_channel_zero(registry)
}

struct NoopProvider;

impl alkcall::core::auth::IdentityProvider for NoopProvider {
    fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
        None
    }
    fn resolve_from_token(&self, _: &alkcall::core::auth::AuthToken) -> Option<Identity> {
        None
    }
}

/// Extension wrapper for the channels-policy injection point (SRV-10):
/// a deployment inserts `ChannelsPolicy(Arc<dyn ChannelLifecyclePolicy>)`
/// into the request extensions (a route layer on the WS route) to gate
/// data-channel opens per identity (ADR-041). Without the extension the
/// upgrade defaults to [`NoCap`] — the crate default for the built-in
/// surface (POC/trusted-peer semantics); assembly layers that build
/// their own upgrade route pass a stricter policy directly to
/// [`run_channels_session`].
#[derive(Clone)]
pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);

/// Per-request WS pump timeouts (WS-17): a deployment inserts
/// `WsTimeouts` into the request extensions (a route layer on the WS
/// route, mirroring [`ChannelsPolicy`]) to set the pump knobs per
/// route instead of the router-state defaults. `idle` (WS-01) bounds
/// the read staleness (no completed inbound chunk for the window →
/// 1001 eviction); `write` (WS-18) bounds one outbound WS send (a peer
/// that stops reading is evicted once a single send outlasts it);
/// `None` disables a knob. Without the extension the
/// [`SessionState`] values apply — the built-in surface carries
/// `HttpAdapter::with_ws_idle_timeout` /
/// `DEFAULT_WS_IDLE_TIMEOUT` for the read side and
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] for the write side,
/// which is also what bare-registry routes (custom upgrade routes on
/// a plain `Arc<OperationRegistry>` state) get: 60 s idle + 60 s
/// write windows, 64-session semaphore, handler-private
/// [`WsSessions`].
#[derive(Clone, Copy, Debug)]
pub struct WsTimeouts {
    /// Idle-read window (WS-01); `None` disables the eviction.
    pub idle: Option<std::time::Duration>,
    /// Write-progress window (WS-18); `None` selects the crate
    /// default (`DEFAULT_WS_WRITE_TIMEOUT`).
    pub write: Option<std::time::Duration>,
}

impl Default for WsTimeouts {
    fn default() -> Self {
        Self {
            idle: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
            write: None,
        }
    }
}

/// The upgrade handler. Requires the resolved identity in request
/// extensions (stashed by [`ws_bearer_auth`]) — a WS session without
/// an identity cannot run `AccessControl::check`.
///
/// The channel lifecycle policy comes from the
/// [`ChannelsPolicy`] request extension when present, else `NoCap`.
/// The session pump handles are retained in the [`WsSessions`]
/// registry — the shared instance from the router state
/// (`RouterState::ws_sessions`) unless a request extension carries
/// one, so a stuck session stays evictable (WS-08).
///
/// The pump timeout knobs (WS-01 read / WS-18 write) come from the
/// [`WsTimeouts`] request extension when present, else from the
/// router state (`HttpAdapter::with_ws_idle_timeout` for the read;
/// the write side is fixed at
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] on the built-in
/// surface). Bare-registry routes (`Arc<OperationRegistry>` state)
/// carry no configurable state — they run the documented defaults
/// above; a deployment overriding them inserts the `WsTimeouts`
/// extension on its upgrade route.
///
/// Session cap (WS-09): one semaphore permit is acquired per upgrade,
/// post-auth and pre-upgrade; when the configured cap
/// (`HttpAdapter::with_ws_max_sessions`, default
/// [`DEFAULT_WS_MAX_SESSIONS`]) is exhausted the upgrade is rejected
/// with **503 Service Unavailable** — holding the permit for the
/// session's lifetime, so ended sessions free their slot.
pub async fn ws_upgrade_handler(
    sessions: Option<axum::Extension<WsSessions>>,
    axum::extract::State(state): axum::extract::State<SessionState>,
    axum::Extension(identity): axum::Extension<Identity>,
    policy: Option<axum::Extension<ChannelsPolicy>>,
    timeouts: Option<axum::Extension<WsTimeouts>>,
    ws_upgrade: WebSocketUpgrade,
) -> Response {
    let Ok(permit) = state.session_slots.clone().try_acquire_owned() else {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            "503 Service Unavailable: WS session cap reached",
        )
            .into_response();
    };
    let sessions = Some(
        sessions
            .map(|axum::Extension(s)| s)
            .unwrap_or_else(|| WsSessions::clone(state.sessions())),
    );
    let policy = policy
        .map(|axum::Extension(p)| p.0)
        .unwrap_or_else(|| Arc::new(NoCap));
    let registry = Arc::clone(state.registry());
    let idle_timeout = match timeouts {
        Some(axum::Extension(t)) => t.idle,
        None => state.idle_timeout(),
    };
    let write_timeout = timeouts.and_then(|t| t.write);
    ws_upgrade
        .max_frame_size(INBOUND_WS_FRAME_CAP)
        .max_message_size(INBOUND_WS_MESSAGE_CAP)
        .on_upgrade(move |socket| async move {
            let _permit = permit;
            run_channels_session(
                socket,
                registry,
                identity,
                policy,
                sessions,
                idle_timeout,
                Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
            )
            .await
        })
}

/// Bearer-auth middleware for the WS upgrade route: resolves the token
/// via the shared [`crate::server::auth`] path and stashes the identity
/// for the upgrade handler. No token / unresolvable token → `401`
/// before the upgrade.
pub async fn ws_bearer_auth(
    axum::extract::State(provider): axum::extract::State<
        Arc<dyn alkcall::core::auth::IdentityProvider>,
    >,
    mut req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Response {
    let identity = crate::server::auth::extract_bearer_identity(&req, provider.as_ref());
    match identity {
        Some(identity) => {
            req.extensions_mut().insert(identity);
            next.run(req).await
        }
        None => (StatusCode::UNAUTHORIZED, "401 Unauthorized").into_response(),
    }
}

/// Test support: a minimal tokio-tungstenite WS client speaking raw
/// channels framing (`frame_channel0_chunk` / `ChunkAssembler` /
/// `FrameAssembler`). Gated behind the `test-support` feature so it
/// never ships in release builds; used by this crate's integration
/// tests and by downstream consumers testing their deployments.
#[cfg(any(test, feature = "test-support"))]
pub mod test_support {
    use alkcall::protocol::wire::EventEnvelope;
    use futures::StreamExt;
    use tokio_tungstenite::tungstenite::client::IntoClientRequest;
    use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;

    /// Frame one `EventEnvelope` as a channel-0 chunk (8-byte chunk
    /// header + 4-byte length prefix + JSON body) — the client-side
    /// framing channel 0 uses over any transport.
    ///
    /// # Panics
    ///
    /// Panics only if `serde_json` cannot serialize the envelope —
    /// unreachable for the acyclic wire type (no non-string map keys,
    /// no untagged ambiguities), which is why this returns `Vec<u8>`
    /// rather than `Result`: a test helper returning `Result` for an
    /// impossible case is worse ergonomics than a documented panic
    /// (review 001 HY-04, kept-as-is decision — the item ships behind
    /// the opt-in `test-support` feature, the crate's documented
    /// exception to no-panics-in-library-code).
    pub fn frame_channel0_chunk(envelope: &EventEnvelope) -> Vec<u8> {
        let body = serde_json::to_vec(envelope).unwrap();
        let mut out = Vec::with_capacity(8 + 4 + body.len());
        out.extend_from_slice(&0u32.to_be_bytes());
        out.extend_from_slice(&((body.len() + 4) as u32).to_be_bytes());
        out.extend_from_slice(&(body.len() as u32).to_be_bytes());
        out.extend_from_slice(&body);
        out
    }

    /// Accumulate WS binary messages into bytes and extract complete
    /// chunks: (channel_id, payload).
    #[derive(Default)]
    pub struct ChunkAssembler {
        buf: Vec<u8>,
    }

    impl ChunkAssembler {
        /// A fresh, empty assembler.
        pub fn new() -> Self {
            Self::default()
        }

        /// Append raw bytes to the assembly buffer.
        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        /// Extract the next complete `(channel_id, payload)` chunk, if
        /// a full one is buffered (8-byte header + declared payload
        /// length).
        pub fn next_chunk(&mut self) -> Option<(u32, Vec<u8>)> {
            if self.buf.len() < 8 {
                return None;
            }
            let channel_id =
                u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]);
            let len =
                u32::from_be_bytes([self.buf[4], self.buf[5], self.buf[6], self.buf[7]]) as usize;
            if self.buf.len() < 8 + len {
                return None;
            }
            let payload = self.buf.drain(..8 + len).skip(8).collect();
            Some((channel_id, payload))
        }
    }

    /// Reassemble length-prefixed frames from the concatenated byte
    /// stream of channel-0 chunk payloads. POC finding (OQ-01): one
    /// call frame may arrive as multiple chunks (`write_frame`'s
    /// prefix and body surface as separate mux payloads), so frame
    /// parsing must run over the reassembled byte stream — never over
    /// individual chunks.
    #[derive(Default)]
    pub struct FrameAssembler {
        buf: Vec<u8>,
    }

    impl FrameAssembler {
        /// A fresh, empty assembler.
        pub fn new() -> Self {
            Self::default()
        }

        /// Append raw bytes to the assembly buffer.
        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        /// Extract the next complete `EventEnvelope` frame, if a full
        /// length-prefixed frame is buffered and parses. Unparseable
        /// frames are dropped (test-only surface).
        pub fn next_frame(&mut self) -> Option<EventEnvelope> {
            if self.buf.len() < 4 {
                return None;
            }
            let len =
                u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]) as usize;
            if self.buf.len() < 4 + len {
                return None;
            }
            let frame: Vec<u8> = self.buf.drain(..4 + len).collect();
            serde_json::from_slice(&frame[4..]).ok()
        }
    }

    /// Minimal WS client for tests: connect with/without a bearer
    /// token, send/recv binary + text, await the close frame.
    pub struct WsClient {
        sink: futures::stream::SplitSink<
            tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
            tokio_tungstenite::tungstenite::Message,
        >,
        stream: futures::stream::SplitStream<
            tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
        >,
    }

    impl WsClient {
        /// Connect with a `Bearer` token header; the full WS stream is
        /// returned (upgrade succeeded).
        pub async fn connect_authorized(url: &str, token: &str) -> Result<Self, String> {
            let mut request = url
                .into_client_request()
                .map_err(|e| format!("bad url: {e}"))?;
            request.headers_mut().insert(
                http::header::AUTHORIZATION,
                http::HeaderValue::from_str(&format!("Bearer {token}"))
                    .map_err(|e| format!("bad token: {e}"))?,
            );
            let (stream, _resp): (
                tokio_tungstenite::WebSocketStream<
                    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
                >,
                _,
            ) = tokio_tungstenite::connect_async(request)
                .await
                .map_err(|e| format!("connect failed: {e}"))?;
            Ok(Self::from_stream(stream))
        }

        /// Connect and return just the HTTP status (for negative tests).
        pub async fn connect_status(url: &str, token: Option<&str>) -> Option<u16> {
            let mut request = url.into_client_request().ok()?;
            if let Some(t) = token {
                request.headers_mut().insert(
                    http::header::AUTHORIZATION,
                    http::HeaderValue::from_str(&format!("Bearer {t}")).ok()?,
                );
            }
            type WsStream = tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >;
            type WsConnectResult = Result<
                (
                    WsStream,
                    tokio_tungstenite::tungstenite::http::Response<Option<Vec<u8>>>,
                ),
                tokio_tungstenite::tungstenite::Error,
            >;
            let result: WsConnectResult = tokio_tungstenite::connect_async(request).await;
            match result {
                Ok((stream, resp)) => {
                    drop(stream);
                    Some(resp.status().as_u16())
                }
                Err(tokio_tungstenite::tungstenite::Error::Http(resp)) => {
                    Some(resp.status().as_u16())
                }
                Err(_) => None,
            }
        }

        fn from_stream(
            stream: tokio_tungstenite::WebSocketStream<
                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
            >,
        ) -> Self {
            let (sink, stream) = stream.split();
            Self { sink, stream }
        }

        /// Send one binary WS message.
        ///
        /// # Panics
        ///
        /// Panics on a socket write failure — a test client that cannot
        /// send has a broken test, not a recoverable runtime state.
        pub async fn send_binary(&mut self, bytes: Vec<u8>) {
            self.send_binary_piece(&bytes).await;
        }

        /// Send one binary WS message, in pieces (for split-frame
        /// tests). Same panic contract as [`Self::send_binary`].
        pub async fn send_binary_piece(&mut self, bytes: &[u8]) {
            use futures::SinkExt;
            self.sink
                .send(tokio_tungstenite::tungstenite::Message::Binary(
                    bytes.to_vec().into(),
                ))
                .await
                .unwrap();
        }

        /// Send one text WS message. Same panic contract as
        /// [`Self::send_binary`].
        pub async fn send_text(&mut self, text: &str) {
            use futures::SinkExt;
            self.sink
                .send(tokio_tungstenite::tungstenite::Message::Text(
                    text.to_string().into(),
                ))
                .await
                .unwrap();
        }

        /// Next binary message, or `None` on timeout/close/error.
        pub async fn next_binary(&mut self, timeout: std::time::Duration) -> Option<Vec<u8>> {
            use futures::StreamExt;
            loop {
                match tokio::time::timeout(timeout, self.stream.next()).await {
                    Err(_) => return None,
                    Ok(None) => return None,
                    Ok(Some(Err(_))) => return None,
                    Ok(Some(Ok(m))) => match m {
                        tokio_tungstenite::tungstenite::Message::Binary(b) => {
                            return Some(b.to_vec())
                        }
                        tokio_tungstenite::tungstenite::Message::Close(_) => return None,
                        _ => continue,
                    },
                }
            }
        }

        /// Await the close frame: `Some(Some(code))` close with code,
        /// `Some(None)` stream ended without a close frame, `None`
        /// timed out.
        pub async fn next_close(&mut self, timeout: std::time::Duration) -> Option<Option<u16>> {
            use futures::StreamExt;
            match tokio::time::timeout(timeout, self.stream.next()).await {
                Err(_) => None,
                Ok(None) => Some(None),
                Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf)))) => {
                    Some(cf.map(|f| match f.code {
                        CloseCode::Error => 1011,
                        other => other.into(),
                    }))
                }
                Ok(Some(Ok(_))) => Box::pin(self.next_close(timeout)).await,
                Ok(Some(Err(_))) => Some(None),
            }
        }

        /// Close the WS with a normal-close frame.
        pub async fn close(&mut self) {
            use futures::SinkExt;
            let _ = self.sink.close().await;
        }
    }
}