//! Server-side `/crdt-sync` WebSocket handler.

use bft_json_crdt::json_crdt::SignedOp;
use futures::{SinkExt, StreamExt};
use poem::handler;
use poem::http::StatusCode;
use poem::web::Query;
use poem::web::websocket::{Message as WsMessage, WebSocket};
use serde::Deserialize;

use crate::crdt_snapshot;
use crate::crdt_state;
use crate::crdt_wire;
use crate::slog;
use crate::slog_warn;

use super::auth::{REQUIRE_TOKEN, validate_join_token};
use super::dispatch::{handle_incoming_binary, handle_incoming_text};
use super::rpc::try_handle_rpc_text;
use super::wire::SyncMessage;
use super::{AUTH_TIMEOUT_SECS, PING_INTERVAL_SECS, PONG_TIMEOUT_SECS};

/// Query parameters accepted on the `/crdt-sync` WebSocket upgrade request.
#[derive(Deserialize)]
struct SyncQueryParams {
    /// Optional bearer token.  Required when the server is in token-required mode.
    token: Option<String>,
}

/// WebSocket handler for CRDT peer synchronisation.
///
/// Accepts an optional `?token=<bearer-token>` query parameter.  When the
/// server is configured with `crdt_require_token = true`, a valid token must
/// be supplied or the upgrade is rejected with HTTP 401.  When the server is
/// in open-access mode (the default), a token is optional but still validated
/// if present.
#[handler]

pub async fn crdt_sync_handler(
    ws: WebSocket,
    remote_addr: &poem::web::RemoteAddr,
    Query(params): Query<SyncQueryParams>,
) -> poem::Response {
    // ── Bearer-token check (pre-upgrade) ────────────────────────────
    let require_token = REQUIRE_TOKEN.get().copied().unwrap_or(false);
    match &params.token {
        Some(t) if !validate_join_token(t) => {
            slog!("[crdt-sync] Rejected connection: invalid or expired token");
            return poem::Response::builder()
                .status(StatusCode::UNAUTHORIZED)
                .body("invalid or expired token");
        }
        None if require_token => {
            slog!("[crdt-sync] Rejected connection: token required but not provided");
            return poem::Response::builder()
                .status(StatusCode::UNAUTHORIZED)
                .body("token required");
        }
        _ => {}
    }

    // ── WebSocket upgrade ────────────────────────────────────────────
    use poem::IntoResponse as _;
    let peer_addr = remote_addr.to_string();
    ws.on_upgrade(move |socket| async move {
        let (mut sink, mut stream) = socket.split();

        slog!("[crdt-sync] Peer connected, starting auth handshake");

        let _auth_msg = match super::handshake::perform_auth_handshake(&mut sink, &mut stream).await {
            Some(m) => m,
            None => return,
        };

        // v2 protocol: send our vector clock so the peer can compute the delta.
        let our_clock = crdt_state::our_vector_clock().unwrap_or_default();
        let clock_msg = SyncMessage::Clock { clock: our_clock };
        if let Ok(json) = serde_json::to_string(&clock_msg)
            && sink.send(WsMessage::Text(json)).await.is_err()
        {
            return;
        }

        // Wait for the peer's first sync message to determine protocol version.
        let first_msg = tokio::time::timeout(
            std::time::Duration::from_secs(AUTH_TIMEOUT_SECS),
            wait_for_sync_text(&mut stream, &mut sink),
        )
        .await;

        match first_msg {
            Ok(Some(SyncMessage::Clock { clock: peer_clock })) => {
                // v2 peer — if we have a snapshot and the peer has an empty
                // clock (new node), send the snapshot first for onboarding.
                if peer_clock.is_empty()
                    && let Some(snapshot) = crdt_snapshot::latest_snapshot()
                {
                    let snap_msg = crdt_snapshot::SnapshotMessage::Snapshot(snapshot);
                    if let Ok(json) = serde_json::to_string(&snap_msg) {
                        if sink.send(WsMessage::Text(json)).await.is_err() {
                            return;
                        }
                        slog!("[crdt-sync] Sent snapshot to new node for onboarding");
                    }
                }

                // Send only the ops the peer is missing.
                let delta = crdt_state::ops_since(&peer_clock).unwrap_or_default();
                slog!(
                    "[crdt-sync] v2 delta sync: sending {} ops (peer missing)",
                    delta.len()
                );
                let msg = SyncMessage::Bulk { ops: delta };
                if let Ok(json) = serde_json::to_string(&msg)
                    && sink.send(WsMessage::Text(json)).await.is_err()
                {
                    return;
                }
            }
            Ok(Some(SyncMessage::Bulk { ops })) => {
                // v1 peer — apply their bulk and send our full bulk.
                let mut applied = 0u64;
                for op_json in &ops {
                    if let Ok(signed_op) = serde_json::from_str::<SignedOp>(op_json)
                        && crdt_state::apply_remote_op(signed_op)
                    {
                        applied += 1;
                    }
                }
                slog!(
                    "[crdt-sync] v1 bulk sync: received {} ops, applied {applied}",
                    ops.len()
                );
                if let Some(all) = crdt_state::all_ops_json() {
                    let msg = SyncMessage::Bulk { ops: all };
                    if let Ok(json) = serde_json::to_string(&msg)
                        && sink.send(WsMessage::Text(json)).await.is_err()
                    {
                        return;
                    }
                }
            }
            Ok(Some(SyncMessage::Op { op })) => {
                // Single op before negotiation — treat as v1.
                if let Ok(signed_op) = serde_json::from_str::<SignedOp>(&op) {
                    crdt_state::apply_remote_op(signed_op);
                }
                if let Some(all) = crdt_state::all_ops_json() {
                    let msg = SyncMessage::Bulk { ops: all };
                    if let Ok(json) = serde_json::to_string(&msg)
                        && sink.send(WsMessage::Text(json)).await.is_err()
                    {
                        return;
                    }
                }
            }
            _ => {
                // Timeout or error — send full bulk as fallback.
                slog!("[crdt-sync] No sync message from peer; sending full bulk as fallback");
                if let Some(all) = crdt_state::all_ops_json() {
                    let msg = SyncMessage::Bulk { ops: all };
                    if let Ok(json) = serde_json::to_string(&msg)
                        && sink.send(WsMessage::Text(json)).await.is_err()
                    {
                        return;
                    }
                }
            }
        }

        // Bulk-delta phase complete — signal the peer that we are ready for
        // real-time op streaming.
        if let Ok(json) = serde_json::to_string(&SyncMessage::Ready)
            && sink.send(WsMessage::Text(json)).await.is_err()
        {
            return;
        }

        // Subscribe to new local ops.
        let Some(mut op_rx) = crdt_state::subscribe_ops() else {
            return;
        };

        // Buffer for locally-generated ops produced before the peer's `ready`
        // arrives.  Flushed in-order once the peer signals catch-up.
        let mut peer_ready = false;
        let mut op_buffer: Vec<bft_json_crdt::json_crdt::SignedOp> = Vec::new();

        // ── Keepalive state ───────────────────────────────────────────
        let mut pong_deadline = tokio::time::Instant::now()
            + std::time::Duration::from_secs(PONG_TIMEOUT_SECS);
        let mut ping_ticker = tokio::time::interval_at(
            tokio::time::Instant::now()
                + std::time::Duration::from_secs(PING_INTERVAL_SECS),
            std::time::Duration::from_secs(PING_INTERVAL_SECS),
        );

        loop {
            tokio::select! {
                // Send periodic Ping and enforce Pong timeout.
                _ = ping_ticker.tick() => {
                    if tokio::time::Instant::now() >= pong_deadline {
                        slog_warn!(
                            "[crdt-sync] No pong from peer {} in {}s; disconnecting",
                            peer_addr,
                            PONG_TIMEOUT_SECS
                        );
                        break;
                    }
                    if sink.send(WsMessage::Ping(vec![])).await.is_err() {
                        break;
                    }
                }
                // Forward new local ops to the peer encoded via the wire codec.
                result = op_rx.recv() => {
                    match result {
                        Ok(signed_op) => {
                            if peer_ready {
                                let bytes = crdt_wire::encode(&signed_op);
                                if sink.send(WsMessage::Binary(bytes)).await.is_err() {
                                    break;
                                }
                            } else {
                                // Buffer until the peer signals ready.
                                op_buffer.push(signed_op);
                            }
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                            // The peer cannot keep up; disconnect so it can
                            // reconnect and receive a fresh bulk state dump.
                            slog!("[crdt-sync] Slow peer lagged {n} ops; disconnecting");
                            break;
                        }
                        Err(_) => break,
                    }
                }
                // Receive ops from the peer.
                frame = stream.next() => {
                    match frame {
                        Some(Ok(WsMessage::Pong(_))) => {
                            // Reset the pong deadline on every Pong received.
                            pong_deadline = tokio::time::Instant::now()
                                + std::time::Duration::from_secs(PONG_TIMEOUT_SECS);
                        }
                        Some(Ok(WsMessage::Ping(data))) => {
                            // Respond to peer's Ping so the peer's keepalive passes.
                            let _ = sink.send(WsMessage::Pong(data)).await;
                        }
                        Some(Ok(WsMessage::Text(text))) => {
                            // Check for the ready signal before other text frames.
                            if let Ok(SyncMessage::Ready) = serde_json::from_str(&text) {
                                peer_ready = true;
                                slog!("[crdt-sync] Peer ready; flushing {} buffered ops", op_buffer.len());
                                let mut flush_ok = true;
                                for op in op_buffer.drain(..) {
                                    let bytes = crdt_wire::encode(&op);
                                    if sink.send(WsMessage::Binary(bytes)).await.is_err() {
                                        flush_ok = false;
                                        break;
                                    }
                                }
                                if !flush_ok {
                                    break;
                                }
                            } else if let Some(rpc_resp) = try_handle_rpc_text(&text).await {
                                // RPC request — dispatch to registry and send response.
                                if let Ok(json) = serde_json::to_string(&rpc_resp)
                                    && sink.send(WsMessage::Text(json)).await.is_err()
                                {
                                    break;
                                }
                            } else {
                                // Bulk state dump, legacy op frame, or clock frame.
                                handle_incoming_text(&text);
                            }
                        }
                        Some(Ok(WsMessage::Binary(bytes))) => {
                            // Real-time op encoded via wire codec — applied immediately
                            // regardless of our own ready state.
                            handle_incoming_binary(&bytes);
                        }
                        Some(Ok(WsMessage::Close(_))) | None => break,
                        _ => {}
                    }
                }
            }
        }

        slog!("[crdt-sync] Peer disconnected");
    })
    .into_response()
}

/// Wait for the next text-frame sync message from the peer, handling Ping/Pong
/// transparently.
///
/// Returns `None` on connection close or read error.
async fn wait_for_sync_text(
    stream: &mut futures::stream::SplitStream<poem::web::websocket::WebSocketStream>,
    sink: &mut futures::stream::SplitSink<poem::web::websocket::WebSocketStream, WsMessage>,
) -> Option<SyncMessage> {
    loop {
        match stream.next().await {
            Some(Ok(WsMessage::Text(text))) => {
                return serde_json::from_str(&text).ok();
            }
            Some(Ok(WsMessage::Ping(data))) => {
                let _ = sink.send(WsMessage::Pong(data)).await;
            }
            Some(Ok(WsMessage::Pong(_))) => continue,
            _ => return None,
        }
    }
}

#[cfg(test)]
mod tests_keepalive;
#[cfg(test)]
mod tests_sync;
