//! 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`.

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, 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 [`crate::server::state::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>>>>,
}

/// 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 [`crate::server::state::RouterState`] (carrying the
/// shared [`WsSessions`] instance) or a bare `Arc<OperationRegistry>`
/// (custom upgrade routes / integration tests get a handler-private
/// registry; eviction still works in-crate, just not shared).
#[derive(Clone)]
pub struct SessionState {
    registry: Arc<OperationRegistry>,
    sessions: Arc<WsSessions>,
}

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).
    pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
        Self {
            registry: Arc::clone(registry),
            sessions: Arc::new(WsSessions::new()),
        }
    }

    /// From the adapter's router state: the shared [`WsSessions`]
    /// instance the assembly layer can hold for eviction.
    pub(crate) fn new(registry: Arc<OperationRegistry>, sessions: Arc<WsSessions>) -> Self {
        Self { registry, sessions }
    }

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

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

impl axum::extract::FromRef<SessionState> for Arc<OperationRegistry> {
    fn from_ref(state: &SessionState) -> Self {
        Arc::clone(&state.registry)
    }
}

/// `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 {
    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()
    }

    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`]).
pub async fn run_channels_session(
    socket: axum::extract::ws::WebSocket,
    registry: Arc<OperationRegistry>,
    identity: Identity,
    policy: Arc<dyn ChannelLifecyclePolicy>,
    sessions: Option<WsSessions>,
) {
    let (byte_stream, pumps) = split_ws_to_bytes(socket);
    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>);

/// 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).
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>>,
    ws_upgrade: WebSocketUpgrade,
) -> 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());
    ws_upgrade
        .max_frame_size(INBOUND_WS_FRAME_CAP)
        .max_message_size(INBOUND_WS_MESSAGE_CAP)
        .on_upgrade(move |socket| async move {
            run_channels_session(socket, registry, identity, policy, sessions).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.
    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 {
        pub fn new() -> Self {
            Self::default()
        }

        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        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 {
        pub fn new() -> Self {
            Self::default()
        }

        pub fn push(&mut self, bytes: &[u8]) {
            self.buf.extend_from_slice(bytes);
        }

        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 {
        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 }
        }

        pub async fn send_binary(&mut self, bytes: Vec<u8>) {
            self.send_binary_piece(&bytes).await;
        }

        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();
        }

        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),
            }
        }

        pub async fn close(&mut self) {
            use futures::SinkExt;
            let _ = self.sink.close().await;
        }
    }
}