//! WebSocket JSON-RPC bridge for browser clients.
//!
//! Provides a lightweight JSON-RPC interface over WebSocket so that browsers
//! can interact with the server without a Cap'n Proto / QUIC stack.
//!
//! Security parity with the Cap'n Proto path:
//! - Rate limiting via `check_rate_limit()` on all mutating handlers
//! - DM channel membership verification on `send`
//! - Payload size limits (5 MB)
//! - Timing floor on `resolveUser` to mask lookup timing

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use base64::Engine;
use dashmap::DashMap;
use futures::stream::StreamExt;
use futures::SinkExt;
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::time::Instant;
use tokio_tungstenite::tungstenite::Message;

use crate::auth::{check_rate_limit, validate_token_raw, AuthConfig, AuthContext, RateEntry, SessionInfo};
use crate::storage::Store;

const B64: base64::engine::general_purpose::GeneralPurpose =
    base64::engine::general_purpose::STANDARD;

/// Maximum payload size for WS bridge (same as Cap'n Proto path).
const MAX_PAYLOAD_BYTES: usize = 5 * 1024 * 1024;

/// Minimum response time for resolveUser to mask DB lookup timing differences.
const RESOLVE_TIMING_FLOOR: Duration = Duration::from_millis(5);

// ── Shared state ────────────────────────────────────────────────────────────

/// Subset of server state needed by the WS bridge (all `Send + Sync`).
#[allow(dead_code)] // sealed_sender plumbed for future use
pub struct WsBridgeState {
    pub store: Arc<dyn Store>,
    pub waiters: Arc<DashMap<Vec<u8>, Arc<Notify>>>,
    pub auth_cfg: Arc<AuthConfig>,
    pub sessions: Arc<DashMap<Vec<u8>, SessionInfo>>,
    pub rate_limits: Arc<DashMap<Vec<u8>, RateEntry>>,
    pub sealed_sender: bool,
    pub allow_insecure_auth: bool,
}

// ── JSON-RPC types ──────────────────────────────────────────────────────────

#[derive(Deserialize)]
struct RpcRequest {
    id: serde_json::Value,
    method: String,
    #[serde(default)]
    params: serde_json::Value,
}

#[derive(Serialize)]
struct RpcResponse {
    id: serde_json::Value,
    ok: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

impl RpcResponse {
    fn success(id: serde_json::Value, result: serde_json::Value) -> Self {
        Self {
            id,
            ok: true,
            result: Some(result),
            error: None,
        }
    }

    fn error(id: serde_json::Value, msg: impl Into<String>) -> Self {
        Self {
            id,
            ok: false,
            result: None,
            error: Some(msg.into()),
        }
    }
}

// ── Auth helper ─────────────────────────────────────────────────────────────

/// Extract and validate the "token" field from params. In insecure-auth mode
/// with no token configured, an empty token is accepted as the bearer.
fn extract_auth(
    state: &WsBridgeState,
    params: &serde_json::Value,
) -> Result<AuthContext, String> {
    let token_str = params
        .get("token")
        .and_then(|v| v.as_str())
        .unwrap_or("");

    let token_bytes = token_str.as_bytes().to_vec();

    // In insecure-auth mode with no configured token, accept any (including empty)
    // token as the bearer token. This mirrors the Cap'n Proto path behaviour.
    if state.allow_insecure_auth && state.auth_cfg.required_token.is_none() {
        // Treat the request identity from params as the identity.
        return Ok(AuthContext {
            token: token_bytes,
            identity_key: None,
        });
    }

    validate_token_raw(&state.auth_cfg, &state.sessions, &token_bytes)
}

/// Resolve identity key: either from auth context (session-bound) or from
/// request params (insecure-auth mode). Returns the 32-byte identity key.
fn resolve_identity(
    state: &WsBridgeState,
    auth_ctx: &AuthContext,
    params: &serde_json::Value,
) -> Result<Vec<u8>, String> {
    // If auth context has an identity-bound session, use that.
    if let Some(ref ik) = auth_ctx.identity_key {
        return Ok(ik.clone());
    }

    // In insecure-auth mode, accept identity from params.
    if state.allow_insecure_auth {
        // Try base64-encoded identityKey first.
        if let Some(b64) = params.get("identityKey").and_then(|v| v.as_str()) {
            return B64
                .decode(b64)
                .map_err(|e| format!("bad base64 identityKey: {e}"));
        }
        // Try username lookup.
        if let Some(username) = params.get("username").and_then(|v| v.as_str()) {
            if let Ok(Some(ik)) = state.store.get_user_identity_key(username) {
                return Ok(ik);
            }
            return Err(format!("user not found: {username}"));
        }
    }

    Err("no identity: login required or pass identityKey/username in insecure mode".to_string())
}

/// Apply rate limiting using the auth token. Returns an error string on limit exceeded.
fn ws_check_rate_limit(state: &WsBridgeState, auth_ctx: &AuthContext) -> Result<(), String> {
    check_rate_limit(&state.rate_limits, &auth_ctx.token)
        .map_err(|e| format!("rate limit exceeded: {e}"))
}

// ── Dispatch ────────────────────────────────────────────────────────────────

async fn dispatch(state: &WsBridgeState, req: RpcRequest) -> RpcResponse {
    match req.method.as_str() {
        "health" => handle_health(req.id),
        "resolveUser" => handle_resolve_user(state, req.id, &req.params).await,
        "createChannel" => handle_create_channel(state, req.id, &req.params),
        "send" => handle_send(state, req.id, &req.params),
        "receive" => handle_receive(state, req.id, &req.params),
        "deleteAccount" => handle_delete_account(state, req.id, &req.params),
        "register" => handle_register(state, req.id, &req.params),
        _ => RpcResponse::error(req.id, format!("unknown method: {}", req.method)),
    }
}

// ── Handlers ────────────────────────────────────────────────────────────────

fn handle_health(id: serde_json::Value) -> RpcResponse {
    RpcResponse::success(id, serde_json::json!("ok"))
}

fn handle_register(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    // Only allow in insecure-auth mode (development/demo).
    if !state.allow_insecure_auth {
        return RpcResponse::error(id, "register is only available in --allow-insecure-auth mode");
    }

    // Rate limit.
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    // Validate username.
    let username = match params.get("username").and_then(|v| v.as_str()) {
        Some(u) if !u.is_empty() => u,
        _ => return RpcResponse::error(id, "missing or empty 'username' param"),
    };
    if username.len() > 32 {
        return RpcResponse::error(id, "username must be at most 32 characters");
    }
    if !username.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
        return RpcResponse::error(id, "username must be alphanumeric or underscore only");
    }

    // Validate identity key.
    let ik_b64 = match params.get("identityKey").and_then(|v| v.as_str()) {
        Some(s) if !s.is_empty() => s,
        _ => return RpcResponse::error(id, "missing or empty 'identityKey' param"),
    };
    let identity_key = match B64.decode(ik_b64) {
        Ok(k) => k,
        Err(e) => return RpcResponse::error(id, format!("bad base64 identityKey: {e}")),
    };
    if identity_key.len() != 32 {
        return RpcResponse::error(id, "identityKey must be 32 bytes");
    }

    // Check if username is already taken by a different key.
    match state.store.get_user_identity_key(username) {
        Ok(Some(existing)) if existing == identity_key => {
            // Idempotent: same key, return success.
            return RpcResponse::success(
                id,
                serde_json::json!({
                    "username": username,
                    "identityKey": B64.encode(&identity_key),
                }),
            );
        }
        Ok(Some(_)) => {
            return RpcResponse::error(id, "username already taken");
        }
        Ok(None) => {} // Available, proceed.
        Err(e) => return RpcResponse::error(id, format!("storage error: {e}")),
    }

    // Store the mapping.
    if let Err(e) = state.store.store_user_identity_key(username, identity_key.clone()) {
        return RpcResponse::error(id, format!("storage error: {e}"));
    }

    tracing::info!(
        username = %username,
        key_prefix = %hex::encode(&identity_key[..4]),
        "audit: ws_bridge register"
    );

    RpcResponse::success(
        id,
        serde_json::json!({
            "username": username,
            "identityKey": B64.encode(&identity_key),
        }),
    )
}

async fn handle_resolve_user(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Rate limit resolve requests to prevent bulk enumeration.
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    let username = match params.get("username").and_then(|v| v.as_str()) {
        Some(u) if !u.is_empty() => u,
        _ => return RpcResponse::error(id, "missing or empty 'username' param"),
    };

    // Timing floor: mask DB-lookup timing differences between existing and
    // non-existing usernames (same as Cap'n Proto resolveUser handler).
    let deadline = Instant::now() + RESOLVE_TIMING_FLOOR;

    let response = match state.store.get_user_identity_key(username) {
        Ok(Some(key)) => {
            RpcResponse::success(id, serde_json::json!({ "identityKey": B64.encode(&key) }))
        }
        Ok(None) => RpcResponse::success(id, serde_json::json!({ "identityKey": null })),
        Err(e) => RpcResponse::error(id, format!("storage error: {e}")),
    };

    // Pad to timing floor before responding.
    tokio::time::sleep_until(deadline).await;

    response
}

fn handle_create_channel(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Rate limit.
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    let my_key = match resolve_identity(state, &auth_ctx, params) {
        Ok(k) => k,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Accept peer key as base64 or resolve from username.
    let peer_key = if let Some(b64) = params.get("peerKey").and_then(|v| v.as_str()) {
        match B64.decode(b64) {
            Ok(k) => k,
            Err(e) => return RpcResponse::error(id, format!("bad base64 peerKey: {e}")),
        }
    } else if let Some(username) = params.get("peerUsername").and_then(|v| v.as_str()) {
        match state.store.get_user_identity_key(username) {
            Ok(Some(k)) => k,
            Ok(None) => return RpcResponse::error(id, format!("peer user not found: {username}")),
            Err(e) => return RpcResponse::error(id, format!("storage error: {e}")),
        }
    } else {
        return RpcResponse::error(id, "missing 'peerKey' (base64) or 'peerUsername'");
    };

    if peer_key.len() != 32 {
        return RpcResponse::error(id, "peerKey must be 32 bytes");
    }

    match state.store.create_channel(&my_key, &peer_key) {
        Ok((channel_id, was_new)) => RpcResponse::success(
            id,
            serde_json::json!({
                "channelId": B64.encode(&channel_id),
                "wasNew": was_new,
            }),
        ),
        Err(e) => RpcResponse::error(id, format!("storage error: {e}")),
    }
}

fn handle_send(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Rate limit (parity with Cap'n Proto enqueue path).
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    let sender_key = match resolve_identity(state, &auth_ctx, params) {
        Ok(k) => k,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Resolve recipient: base64 key or username.
    let recipient_key =
        if let Some(b64) = params.get("recipientKey").and_then(|v| v.as_str()) {
            match B64.decode(b64) {
                Ok(k) => k,
                Err(e) => {
                    return RpcResponse::error(id, format!("bad base64 recipientKey: {e}"))
                }
            }
        } else if let Some(username) = params.get("recipient").and_then(|v| v.as_str()) {
            match state.store.get_user_identity_key(username) {
                Ok(Some(k)) => k,
                Ok(None) => {
                    return RpcResponse::error(id, format!("recipient not found: {username}"))
                }
                Err(e) => return RpcResponse::error(id, format!("storage error: {e}")),
            }
        } else {
            return RpcResponse::error(id, "missing 'recipientKey' (base64) or 'recipient' (username)");
        };

    if recipient_key.len() != 32 {
        return RpcResponse::error(id, "recipientKey must be 32 bytes");
    }

    // Payload: base64-encoded binary or plain text message.
    let payload = if let Some(b64) = params.get("payload").and_then(|v| v.as_str()) {
        match B64.decode(b64) {
            Ok(p) => p,
            Err(e) => return RpcResponse::error(id, format!("bad base64 payload: {e}")),
        }
    } else if let Some(msg) = params.get("message").and_then(|v| v.as_str()) {
        msg.as_bytes().to_vec()
    } else {
        return RpcResponse::error(id, "missing 'payload' (base64) or 'message' (text)");
    };

    if payload.is_empty() {
        return RpcResponse::error(id, "payload must not be empty");
    }

    // Payload size limit (same as Cap'n Proto path: 5 MB).
    if payload.len() > MAX_PAYLOAD_BYTES {
        return RpcResponse::error(
            id,
            format!("payload exceeds max size ({MAX_PAYLOAD_BYTES} bytes)"),
        );
    }

    // Create or look up the DM channel between sender and recipient.
    let channel_id = match state.store.create_channel(&sender_key, &recipient_key) {
        Ok((ch, _)) => ch,
        Err(e) => return RpcResponse::error(id, format!("channel error: {e}")),
    };

    // DM channel membership verification (parity with Cap'n Proto enqueue path).
    if channel_id.len() == 16 {
        let members = match state.store.get_channel_members(&channel_id) {
            Ok(Some(m)) => m,
            Ok(None) => return RpcResponse::error(id, "channel not found"),
            Err(e) => return RpcResponse::error(id, format!("storage error: {e}")),
        };
        let (a, b) = &members;
        let caller_in = sender_key == *a || sender_key == *b;
        let recipient_other = (recipient_key == *a && sender_key == *b)
            || (recipient_key == *b && sender_key == *a);
        if !caller_in || !recipient_other {
            return RpcResponse::error(id, "caller or recipient not a member of this channel");
        }
    }

    match state
        .store
        .enqueue(&recipient_key, &channel_id, payload, None)
    {
        Ok(seq) => {
            // Notify any waiting long-poll fetchers.
            if let Some(notify) = state.waiters.get(&recipient_key) {
                notify.notify_waiters();
            }

            // Audit logging (no secrets: no payload, no full keys).
            tracing::info!(
                recipient_prefix = %hex::encode(&recipient_key[..std::cmp::min(4, recipient_key.len())]),
                seq = seq,
                "audit: ws_bridge enqueue"
            );

            RpcResponse::success(id, serde_json::json!({ "seq": seq }))
        }
        Err(e) => RpcResponse::error(id, format!("enqueue error: {e}")),
    }
}

fn handle_receive(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Rate limit.
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    let my_key = match resolve_identity(state, &auth_ctx, params) {
        Ok(k) => k,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Resolve sender/peer: base64 key or username (needed to find the channel).
    let peer_key = if let Some(b64) = params.get("recipientKey").and_then(|v| v.as_str()) {
        match B64.decode(b64) {
            Ok(k) => k,
            Err(e) => return RpcResponse::error(id, format!("bad base64 recipientKey: {e}")),
        }
    } else if let Some(username) = params.get("recipient").and_then(|v| v.as_str()) {
        match state.store.get_user_identity_key(username) {
            Ok(Some(k)) => k,
            Ok(None) => return RpcResponse::error(id, format!("user not found: {username}")),
            Err(e) => return RpcResponse::error(id, format!("storage error: {e}")),
        }
    } else {
        return RpcResponse::error(id, "missing 'recipientKey' (base64) or 'recipient' (username)");
    };

    // Find the channel between me and the peer.
    let channel_id = match state.store.create_channel(&my_key, &peer_key) {
        Ok((ch, _)) => ch,
        Err(e) => return RpcResponse::error(id, format!("channel error: {e}")),
    };

    // Fetch (drain) all messages for me in this channel.
    match state.store.fetch(&my_key, &channel_id) {
        Ok(messages) => {
            let items: Vec<serde_json::Value> = messages
                .into_iter()
                .map(|(seq, data)| {
                    // Try to decode as UTF-8 text; fall back to base64.
                    let text = String::from_utf8(data.clone()).ok();
                    serde_json::json!({
                        "seq": seq,
                        "data": B64.encode(&data),
                        "text": text,
                    })
                })
                .collect();
            RpcResponse::success(id, serde_json::json!(items))
        }
        Err(e) => RpcResponse::error(id, format!("fetch error: {e}")),
    }
}

fn handle_delete_account(
    state: &WsBridgeState,
    id: serde_json::Value,
    params: &serde_json::Value,
) -> RpcResponse {
    let auth_ctx = match extract_auth(state, params) {
        Ok(ctx) => ctx,
        Err(e) => return RpcResponse::error(id, e),
    };

    // Rate limit.
    if let Err(e) = ws_check_rate_limit(state, &auth_ctx) {
        return RpcResponse::error(id, e);
    }

    let identity_key = match resolve_identity(state, &auth_ctx, params) {
        Ok(k) => k,
        Err(e) => return RpcResponse::error(id, e),
    };

    match state.store.delete_account(&identity_key) {
        Ok(()) => {
            // Invalidate sessions for this identity.
            let tokens_to_remove: Vec<Vec<u8>> = state
                .sessions
                .iter()
                .filter(|entry| entry.value().identity_key == identity_key)
                .map(|entry| entry.key().clone())
                .collect();
            for token in &tokens_to_remove {
                state.sessions.remove(token);
            }
            RpcResponse::success(id, serde_json::json!({ "deleted": true }))
        }
        Err(e) => RpcResponse::error(id, format!("delete failed: {e}")),
    }
}

// ── WebSocket listener ──────────────────────────────────────────────────────

/// Spawn the WebSocket JSON-RPC bridge as a background tokio task.
pub fn spawn_ws_bridge(addr: SocketAddr, state: Arc<WsBridgeState>) {
    tokio::spawn(async move {
        let listener = match TcpListener::bind(addr).await {
            Ok(l) => l,
            Err(e) => {
                tracing::error!(addr = %addr, error = %e, "ws_bridge: failed to bind");
                return;
            }
        };
        tracing::info!(addr = %addr, "ws_bridge: accepting WebSocket connections");

        loop {
            let (stream, peer) = match listener.accept().await {
                Ok(pair) => pair,
                Err(e) => {
                    tracing::warn!(error = %e, "ws_bridge: accept error");
                    continue;
                }
            };

            let state = Arc::clone(&state);
            tokio::spawn(async move {
                let ws = match tokio_tungstenite::accept_async(stream).await {
                    Ok(ws) => ws,
                    Err(e) => {
                        tracing::debug!(peer = %peer, error = %e, "ws_bridge: handshake failed");
                        return;
                    }
                };
                tracing::debug!(peer = %peer, "ws_bridge: client connected");

                let (mut sink, mut stream) = ws.split();

                while let Some(msg) = stream.next().await {
                    let msg = match msg {
                        Ok(m) => m,
                        Err(e) => {
                            tracing::debug!(peer = %peer, error = %e, "ws_bridge: read error");
                            break;
                        }
                    };

                    let text = match msg {
                        Message::Text(t) => t,
                        Message::Close(_) => break,
                        Message::Ping(_) | Message::Pong(_) => continue,
                        _ => continue,
                    };

                    let req: RpcRequest = match serde_json::from_str(&text) {
                        Ok(r) => r,
                        Err(e) => {
                            let resp = RpcResponse::error(
                                serde_json::Value::Null,
                                format!("invalid JSON: {e}"),
                            );
                            let json = serde_json::to_string(&resp).unwrap_or_default();
                            if sink.send(Message::Text(json.into())).await.is_err() {
                                break;
                            }
                            continue;
                        }
                    };

                    let resp = dispatch(&state, req).await;
                    let json = serde_json::to_string(&resp).unwrap_or_default();
                    if sink.send(Message::Text(json.into())).await.is_err() {
                        break;
                    }
                }

                tracing::debug!(peer = %peer, "ws_bridge: client disconnected");
            });
        }
    });
}
