solidb 1.0.1

use super::cluster::{collect_sysinfo, generate_cluster_status};
use super::system::AppState;
use crate::{server::handlers::auth::AuthParams, storage::StorageEngine};
use axum::{
    body::Body,
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query as AxumQuery, State,
    },
    http::HeaderMap,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use futures::{SinkExt, StreamExt};
use serde::Deserialize;
use std::sync::Arc;

/// Maximum WebSocket message size (1 MB) - prevents OOM attacks
const MAX_WS_MESSAGE_SIZE: usize = 1024 * 1024;

/// Validate the `Origin` header against `SOLIDB_CORS_ALLOWED_ORIGINS`.
/// Mirrors the HTTP CORS policy in `routes.rs`: empty allowlist = deny any
/// cross-origin request. Non-browser clients (no `Origin` header) are allowed.
/// Returns Ok(()) when the request may proceed, Err(()) to reject with 403.
fn validate_ws_origin(headers: &HeaderMap) -> Result<(), ()> {
    let origin = match headers.get("origin").and_then(|o| o.to_str().ok()) {
        Some(o) => o,
        None => return Ok(()), // No Origin header — non-browser client.
    };
    let allowed_raw = std::env::var("SOLIDB_CORS_ALLOWED_ORIGINS").unwrap_or_default();
    if allowed_raw == "*" {
        return Ok(());
    }
    if allowed_raw.is_empty() {
        tracing::warn!(
            "WebSocket: rejecting Origin '{}' — SOLIDB_CORS_ALLOWED_ORIGINS not set",
            origin
        );
        return Err(());
    }
    let allowed = allowed_raw
        .split(',')
        .map(str::trim)
        .any(|a| a == origin || a == "*");
    if allowed {
        Ok(())
    } else {
        tracing::warn!("WebSocket: rejecting disallowed Origin '{}'", origin);
        Err(())
    }
}

fn forbidden_response() -> Response {
    Response::builder()
        .status(StatusCode::FORBIDDEN)
        .body(Body::empty())
        .expect("Valid status code should not fail")
        .into_response()
}

// ==================== Cluster Status WebSocket ====================

/// WebSocket handler for real-time cluster status updates
pub async fn cluster_status_ws(
    ws: WebSocketUpgrade,
    AxumQuery(params): AxumQuery<AuthParams>,
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Response {
    let Ok(claims) = crate::server::auth::AuthService::validate_token(&params.token) else {
        return Response::builder()
            .status(StatusCode::UNAUTHORIZED)
            .body(Body::empty())
            .expect("Valid status code should not fail")
            .into_response();
    };
    if crate::server::authz_middleware::enforce(
        &claims,
        &state,
        crate::server::authorization::PermissionAction::Admin,
        None,
    )
    .await
    .is_err()
    {
        return forbidden_response();
    }

    if validate_ws_origin(&headers).is_err() {
        return forbidden_response();
    }

    ws.on_upgrade(|socket| handle_cluster_ws(socket, state))
}

/// Handle the WebSocket connection for cluster status
async fn handle_cluster_ws(mut socket: WebSocket, state: AppState) {
    use tokio::time::{interval, Duration};

    let mut ticker = interval(Duration::from_secs(1));

    // We use the shared system monitor from AppState to avoid expensive initialization
    // and to ensure CPU usage is calculated correctly (delta since last refresh).

    loop {
        tokio::select! {
            _ = ticker.tick() => {
                // Extract sysinfo under a short lock, then generate status without holding it
                let sysinfo = {
                    let mut sys = state.system_monitor.lock().unwrap();
                    collect_sysinfo(&mut sys)
                };
                let status = generate_cluster_status(&state, &sysinfo);

                let json = match serde_json::to_string(&status) {
                    Ok(j) => j,
                    Err(_) => continue,
                };

                if socket.send(Message::Text(json.into())).await.is_err() {
                    break; // Client disconnected
                }
            }
            msg = socket.recv() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    #[allow(clippy::collapsible_match)]
                    Some(Ok(Message::Ping(data))) => {
                        // Respond to ping with pong
                        if socket.send(Message::Pong(data)).await.is_err() {
                            break;
                        }
                    }
                    _ => {} // Ignore other messages
                }
            }
        }
    }
}

// ==================== System Monitoring WebSocket ====================

pub async fn monitor_ws_handler(
    ws: WebSocketUpgrade,
    AxumQuery(params): AxumQuery<AuthParams>,
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Response {
    if crate::server::auth::AuthService::validate_token(&params.token).is_err() {
        return Response::builder()
            .status(StatusCode::UNAUTHORIZED)
            .body(Body::empty())
            .expect("Valid status code should not fail")
            .into_response();
    }

    if validate_ws_origin(&headers).is_err() {
        return forbidden_response();
    }

    ws.on_upgrade(|socket| handle_monitor_socket(socket, state))
}

async fn handle_monitor_socket(mut socket: WebSocket, state: AppState) {
    use std::sync::atomic::Ordering;

    tracing::info!("Monitor WS: Client connected");

    let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));

    loop {
        // Wait for next tick
        interval.tick().await;

        let stats = {
            let mut sys = state.system_monitor.lock().unwrap();

            // Refresh specific stats
            sys.refresh_cpu_all();
            sys.refresh_memory();

            let cpu = sys.global_cpu_usage();
            let mem_used = sys.used_memory();
            let mem_total = sys.total_memory();
            let up = sysinfo::System::uptime();
            let name = sysinfo::System::name().unwrap_or_else(|| "Unknown".to_string());
            let version =
                sysinfo::System::kernel_version().unwrap_or_else(|| "Unknown".to_string());
            let host = sysinfo::System::host_name().unwrap_or_else(|| "Unknown".to_string());
            let cores = sys.cpus().len();

            serde_json::json!({
                "cpu_usage": cpu,
                "memory_usage": mem_used,
                "memory_total": mem_total,
                "uptime": up,
                "os_name": name,
                "os_version": version,
                "hostname": host,
                "num_cpus": cores,
                "pid": std::process::id(),
                "active_scripts": state.script_stats.active_scripts.load(Ordering::Relaxed),
                "active_ws": state.script_stats.active_ws.load(Ordering::Relaxed)
            })
        };

        let msg = match serde_json::to_string(&stats) {
            Ok(s) => s,
            Err(_) => continue,
        };

        if socket.send(Message::Text(msg.into())).await.is_err() {
            // Client disconnected
            break;
        }
    }
}

// ==================== Real-time Changefeeds ====================

#[derive(Debug, Deserialize)]
pub struct ChangefeedRequest {
    #[serde(rename = "type")]
    pub type_: String,
    pub collection: Option<String>,
    pub database: Option<String>,
    pub key: Option<String>,
    pub local: Option<bool>,
    /// SDBQL query for live_query mode
    pub query: Option<String>,
    /// Optional Client ID to identify the subscription/query in responses
    pub id: Option<String>,
}

/// WebSocket handler for real-time changefeeds
pub async fn ws_changefeed_handler(
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    AxumQuery(params): AxumQuery<AuthParams>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    // Check for cluster-internal authentication (bypasses normal JWT validation)
    let is_cluster_internal = {
        let cluster_secret = state.cluster_secret();
        let provided_secret = headers
            .get("X-Cluster-Secret")
            .and_then(|h| h.to_str().ok())
            .unwrap_or("");

        // Use constant-time comparison to prevent timing attacks
        !cluster_secret.is_empty()
            && crate::server::auth::constant_time_eq(
                cluster_secret.as_bytes(),
                provided_secret.as_bytes(),
            )
    };

    // If not cluster-internal, validate the JWT token. Keep the claims:
    // each subscription is authorized against the database it targets.
    let claims = if is_cluster_internal {
        crate::server::auth::Claims {
            sub: "cluster-internal".to_string(),
            exp: usize::MAX,
            livequery: None,
            roles: Some(vec!["admin".to_string()]),
            scoped_databases: None,
        }
    } else {
        match crate::server::auth::AuthService::validate_token(&params.token) {
            Ok(claims) => claims,
            Err(_) => {
                return Response::builder()
                    .status(StatusCode::UNAUTHORIZED)
                    .body(Body::empty())
                    .expect("Valid status code should not fail")
                    .into_response();
            }
        }
    };

    if validate_ws_origin(&headers).is_err() {
        return forbidden_response();
    }

    // Check if HTMX mode is requested
    let use_htmx = params.htmx.map(|s| s == "true").unwrap_or(false);

    ws.on_upgrade(move |socket| handle_socket(socket, state, claims, use_htmx))
}

async fn handle_socket(
    socket: WebSocket,
    state: AppState,
    claims: crate::server::auth::Claims,
    use_htmx: bool,
) {
    // Split socket into sender and receiver
    let (mut sender, mut receiver) = socket.split();

    // Unified channel for sending messages to the client
    // All subscription tasks and live queries will send ready-to-emit Messages to this channel
    let (tx, mut rx) = tokio::sync::mpsc::channel::<Message>(1000);

    // Spawn writer task that forwards messages from the channel to the WebSocket
    let send_task = tokio::spawn(async move {
        // Heartbeat: Send a Ping every 30 seconds to keep the connection alive.
        // Start the interval 30s in the future — `tokio::time::interval` would
        // otherwise fire its first tick immediately, racing the first response
        // frame and surfacing a stray Ping to clients that only ever expect a
        // reply to what they just sent.
        let heartbeat = std::time::Duration::from_secs(30);
        let mut heartbeat_interval =
            tokio::time::interval_at(tokio::time::Instant::now() + heartbeat, heartbeat);

        loop {
            tokio::select! {
                // Send heartbeat
                _ = heartbeat_interval.tick() => {
                    if sender.send(Message::Ping(vec![].into())).await.is_err() {
                        tracing::debug!("[WS] Failed to send ping, closing writer");
                        break;
                    }
                }
                // Forward messages
                Some(msg) = rx.recv() => {
                    if sender.send(msg).await.is_err() {
                        tracing::debug!("[WS] Failed to send message, closing writer");
                        break;
                    }
                }
                else => break,
            }
        }
    });

    // Main Receiver Loop
    while let Some(Ok(msg)) = receiver.next().await {
        // Security: Check message size to prevent OOM attacks
        let msg_len = match &msg {
            Message::Text(text) => text.len(),
            Message::Binary(data) => data.len(),
            Message::Ping(data) => data.len(),
            Message::Pong(data) => data.len(),
            _ => 0,
        };

        if msg_len > MAX_WS_MESSAGE_SIZE {
            tracing::warn!(
                "[WS] Message size {} exceeds limit {}, closing connection",
                msg_len,
                MAX_WS_MESSAGE_SIZE
            );
            let _ = tx
                .send(Message::Text(
                    serde_json::json!({
                        "error": "Message too large"
                    })
                    .to_string()
                    .into(),
                ))
                .await;
            break;
        }

        match msg {
            Message::Text(text) => {
                let req_result = serde_json::from_str::<ChangefeedRequest>(&text);
                match req_result {
                    Ok(req) if req.type_ == "subscribe" => {
                        let tx_clone = tx.clone();
                        let state_clone = state.clone();
                        let claims_clone = claims.clone();

                        // Spawn a dedicated task for this subscription
                        tokio::spawn(async move {
                            handle_subscribe_request(
                                req,
                                state_clone,
                                claims_clone,
                                tx_clone,
                                use_htmx,
                            )
                            .await;
                        });
                    }
                    Ok(req) if req.type_ == "live_query" => {
                        let tx_clone = tx.clone();
                        let state_clone = state.clone();
                        let claims_clone = claims.clone();

                        // Spawn a dedicated task for this live query
                        tokio::spawn(async move {
                            handle_live_query_request(req, state_clone, claims_clone, tx_clone)
                                .await;
                        });
                    }
                    _ => {
                        let _ = tx
                            .send(Message::Text(
                                serde_json::json!({
                                    "error": "Invalid subscription request or unknown type"
                                })
                                .to_string()
                                .into(),
                            ))
                            .await;
                    }
                }
            }
            Message::Close(_) => break,
            Message::Ping(_) => {
                // Auto-replied with Pong by axum usually, but we can ignore
            }
            Message::Pong(_) => {
                // Heartbeat response, ignore
            }
            _ => {}
        }
    }

    // When log out, abort the sender task
    send_task.abort();
}

/// Handle a single subscription request
async fn handle_subscribe_request(
    req: ChangefeedRequest,
    state: AppState,
    claims: crate::server::auth::Claims,
    tx: tokio::sync::mpsc::Sender<Message>,
    use_htmx: bool,
) {
    let db_name = req.database.clone().unwrap_or("_system".to_string());

    // A changefeed exposes every document change in the collection; require
    // read permission on the target database before subscribing.
    if let Err(e) = crate::server::authz_middleware::enforce(
        &claims,
        &state,
        crate::server::authorization::PermissionAction::Read,
        Some(&db_name),
    )
    .await
    {
        let mut response = serde_json::json!({ "error": e.to_string() });
        if let Some(req_id) = &req.id {
            response["id"] = serde_json::Value::String(req_id.clone());
        }
        let _ = tx.send(Message::Text(response.to_string().into())).await;
        return;
    }

    let coll_name = match req.collection.clone() {
        Some(c) => c,
        None => {
            // Try to infer from SDBQL query
            if let Some(query_str) = &req.query {
                if let Ok(query_ast) = crate::sdbql::parser::parse(query_str) {
                    // Check explicit FOR clauses first
                    if let Some(first_for) = query_ast.for_clauses.first() {
                        first_for.collection.clone()
                    } else {
                        // Check body clauses
                        query_ast
                            .body_clauses
                            .iter()
                            .find_map(|c| {
                                if let crate::sdbql::ast::BodyClause::For(f) = c {
                                    Some(f.collection.clone())
                                } else {
                                    None
                                }
                            })
                            .unwrap_or_default()
                    }
                } else {
                    "".to_string()
                }
            } else {
                "".to_string()
            }
        }
    };

    if coll_name.is_empty() {
        let _ = tx
            .send(Message::Text(
                serde_json::json!({
                    "error": "Collection required for subscribe mode (could not infer from query)"
                })
                .to_string()
                .into(),
            ))
            .await;
        return;
    }

    // Try to get collection from specific database or fallback
    let collection_result = state
        .storage
        .get_database(&db_name)
        .and_then(|db| db.get_collection(&coll_name));

    match collection_result {
        Ok(collection) => {
            // Send confirmation
            let msg = if use_htmx {
                format!(
                    r#"<div id="connection-status" hx-swap-oob="innerHTML" class="inline-flex items-center gap-2 px-3 py-1.5 rounded-full text-sm bg-success/10 text-success">
                    <span class="w-2 h-2 rounded-full bg-success animate-pulse"></span>
                    <span>Connected: {}</span>
                </div>
                <div id="no-subscriptions" hx-swap-oob="true" class="hidden"></div>
                <div id="subscriptions-list" hx-swap-oob="beforeend">
                    <div class="px-4 py-3 border-b border-border/20 last:border-0 flex items-center justify-between">
                        <div class="flex items-center gap-3">
                        <span class="w-2 h-2 rounded-full bg-success animate-pulse"></span>
                        <div>
                            <span class="font-medium text-text">{}</span>
                        </div>
                        </div>
                    </div>
                </div>"#,
                    coll_name, coll_name
                )
            } else {
                serde_json::json!({
                    "type": "subscribed",
                    "collection": coll_name
                })
                .to_string()
            };
            if tx.send(Message::Text(msg.into())).await.is_err() {
                return;
            }

            // Set up our OWN internal channel to aggregate events for THIS subscription
            // Then we format them and send to the main `tx`
            let (sub_tx, mut sub_rx) =
                tokio::sync::mpsc::channel::<crate::storage::collection::ChangeEvent>(1000);
            let req_key = req.key.clone();

            // 1. Subscribe to local logical collection
            let mut local_rx = collection.change_sender.subscribe();
            let sub_tx_local = sub_tx.clone();
            let req_key_local = req_key.clone();

            tokio::spawn(async move {
                loop {
                    match local_rx.recv().await {
                        Ok(event) => {
                            if let Some(ref target_key) = req_key_local {
                                if &event.key != target_key {
                                    continue;
                                }
                            }
                            if sub_tx_local.send(event).await.is_err() {
                                break;
                            }
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                        Err(_) => break,
                    }
                }
            });

            // 2. Subscribe to PHYSICAL SHARDS (if sharded)
            if let Some(shard_config) = collection.get_shard_config() {
                if shard_config.num_shards > 0 {
                    if let Ok(database) = state.storage.get_database(&db_name) {
                        for shard_id in 0..shard_config.num_shards {
                            let physical_name = format!("{}_s{}", coll_name, shard_id);
                            if let Ok(physical_coll) = database.get_collection(&physical_name) {
                                let mut shard_rx = physical_coll.change_sender.subscribe();
                                let sub_tx_shard = sub_tx.clone();
                                let req_key_shard = req_key.clone();

                                tokio::spawn(async move {
                                    loop {
                                        match shard_rx.recv().await {
                                            Ok(event) => {
                                                if let Some(ref target_key) = req_key_shard {
                                                    if &event.key != target_key {
                                                        continue;
                                                    }
                                                }
                                                if sub_tx_shard.send(event).await.is_err() {
                                                    break;
                                                }
                                            }
                                            Err(
                                                tokio::sync::broadcast::error::RecvError::Lagged(_),
                                            ) => continue,
                                            Err(_) => break,
                                        }
                                    }
                                });
                            }
                        }
                    }
                }
            }

            // 3. Connect to REMOTE nodes
            let is_local_only = req.local.unwrap_or(false);

            if !is_local_only {
                if let Some(shard_config) = collection.get_shard_config() {
                    if let Some(coordinator) = &state.shard_coordinator {
                        let my_addr = coordinator.my_address();
                        let all_nodes = coordinator.get_collection_nodes(&shard_config);
                        let cluster_secret = state.cluster_secret();

                        let mut remote_nodes = std::collections::HashSet::new();
                        for node_addr in all_nodes {
                            if node_addr != my_addr {
                                remote_nodes.insert(node_addr);
                            }
                        }

                        for node_addr in remote_nodes {
                            let sub_tx_remote = sub_tx.clone();
                            let db_name_remote = db_name.clone();
                            let coll_name_remote = coll_name.clone();
                            let node_addr_clone = node_addr.clone();
                            let secret_clone = cluster_secret.clone();

                            tokio::spawn(async move {
                                use crate::cluster::ClusterWebsocketClient;
                                if let Ok(stream) = ClusterWebsocketClient::connect(
                                    &node_addr_clone,
                                    &db_name_remote,
                                    &coll_name_remote,
                                    true,
                                    &secret_clone,
                                )
                                .await
                                {
                                    tokio::pin!(stream);
                                    while let Some(result) = stream.next().await {
                                        match result {
                                            Ok(event) => {
                                                if sub_tx_remote.send(event).await.is_err() {
                                                    break;
                                                }
                                            }
                                            Err(_) => break,
                                        }
                                    }
                                }
                            });
                        }
                    }
                }
            }

            // Drop original sub_tx so we don't hold the channel open forever if all producers die
            drop(sub_tx);

            // Forward aggregated events to the main socket channel
            while let Some(event) = sub_rx.recv().await {
                // Double check filter (especially for remote events)
                if let Some(ref target_key) = req.key {
                    if &event.key != target_key {
                        continue;
                    }
                }

                // Format message
                let msg_text = if use_htmx {
                    use crate::storage::collection::ChangeType;
                    let op_type = match event.type_ {
                        ChangeType::Insert => "INSERT",
                        ChangeType::Update => "UPDATE",
                        ChangeType::Delete => "DELETE",
                        ChangeType::Truncate => "TRUNCATE",
                    };
                    let status_class = match event.type_ {
                        ChangeType::Insert => "bg-success/10 text-success",
                        ChangeType::Update => "bg-warning/10 text-warning",
                        ChangeType::Delete => "bg-error/10 text-error",
                        ChangeType::Truncate => "bg-error/10 text-error",
                    };
                    let data_str = event
                        .data
                        .as_ref()
                        .map(|v| v.to_string())
                        .unwrap_or_default();

                    format!(
                        r#"<div hx-swap-oob="afterbegin:#events-container">
                        <div class="px-4 py-2 border-b border-border/10 last:border-0 font-mono text-sm hover:bg-white/5 transition-colors">
                            <div class="flex items-center gap-2 mb-1">
                                <span class="px-1.5 py-0.5 rounded text-xs {}">{}</span>
                                <span class="text-text-dim text-xs">{}</span>
                                <span class="text-text-dim text-xs ml-auto">{}</span>
                            </div>
                            <pre class="text-text-muted text-xs overflow-x-auto">{}</pre>
                        </div>
                    </div>"#,
                        status_class,
                        op_type,
                        coll_name,
                        chrono::Local::now().format("%H:%M:%S"),
                        data_str
                    )
                } else {
                    serde_json::json!({
                        "operation": event.type_,
                        "collection": coll_name,
                        "key": event.key,
                        "data": event.data
                    })
                    .to_string()
                };

                if tx.send(Message::Text(msg_text.into())).await.is_err() {
                    break;
                }
            }
        }
        Err(_) => {
            let _ = tx
                .send(Message::Text(
                    serde_json::json!({
                        "error": format!("Collection '{}' not found", coll_name)
                    })
                    .to_string()
                    .into(),
                ))
                .await;
        }
    }
}

/// Handle a live query request
async fn handle_live_query_request(
    req: ChangefeedRequest,
    state: AppState,
    claims: crate::server::auth::Claims,
    tx: tokio::sync::mpsc::Sender<Message>,
) {
    if let Some(query_str) = req.query {
        let db_name = req.database.clone().unwrap_or("_system".to_string());

        // Live queries re-execute against the database on every change;
        // require read permission before registering the subscription.
        if let Err(e) = crate::server::authz_middleware::enforce(
            &claims,
            &state,
            crate::server::authorization::PermissionAction::Read,
            Some(&db_name),
        )
        .await
        {
            let mut response = serde_json::json!({ "error": e.to_string() });
            if let Some(req_id) = &req.id {
                response["id"] = serde_json::Value::String(req_id.clone());
            }
            let _ = tx.send(Message::Text(response.to_string().into())).await;
            return;
        }

        // 1. Parse query to identify dependencies
        match crate::sdbql::parser::parse(&query_str) {
            Ok(query) => {
                let mut dependencies = std::collections::HashSet::new();
                for clause in &query.for_clauses {
                    dependencies.insert(clause.collection.clone());
                }

                if dependencies.is_empty() {
                    let _ = tx
                        .send(Message::Text(
                            serde_json::json!({
                                "error": "Live query must reference at least one collection"
                            })
                            .to_string()
                            .into(),
                        ))
                        .await;
                    return;
                }

                // Send confirmation
                let mut response = serde_json::json!({
                    "type": "subscribed",
                    "mode": "live_query",
                    "collections": dependencies
                });
                if let Some(req_id) = &req.id {
                    response["id"] = serde_json::Value::String(req_id.clone());
                }
                if tx
                    .send(Message::Text(response.to_string().into()))
                    .await
                    .is_err()
                {
                    return;
                }

                // 2. Setup aggregated change channel for dependencies
                let (dep_tx, mut dep_rx) =
                    tokio::sync::mpsc::channel::<crate::storage::collection::ChangeEvent>(1000);

                // 3. Subscribe to ALL dependencies
                for coll_name in &dependencies {
                    let coll_name = coll_name.clone();

                    if let Ok(collection) = state
                        .storage
                        .get_database(&db_name)
                        .and_then(|db| db.get_collection(&coll_name))
                    {
                        // A. Subscribe to local logical
                        let mut local_rx = collection.change_sender.subscribe();
                        let tx_local = dep_tx.clone();
                        tokio::spawn(async move {
                            loop {
                                match local_rx.recv().await {
                                    Ok(event) => {
                                        if tx_local.send(event).await.is_err() {
                                            break;
                                        }
                                    }
                                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
                                        continue
                                    }
                                    Err(_) => break,
                                }
                            }
                        });

                        // B. Subscribe to local physical shards
                        if let Some(shard_config) = collection.get_shard_config() {
                            if shard_config.num_shards > 0 {
                                if let Ok(database) = state.storage.get_database(&db_name) {
                                    for shard_id in 0..shard_config.num_shards {
                                        let physical_name = format!("{}_s{}", coll_name, shard_id);
                                        if let Ok(physical_coll) =
                                            database.get_collection(&physical_name)
                                        {
                                            let mut shard_rx =
                                                physical_coll.change_sender.subscribe();
                                            let tx_shard = dep_tx.clone();
                                            tokio::spawn(async move {
                                                loop {
                                                    match shard_rx.recv().await {
                                                        Ok(event) => { if tx_shard.send(event).await.is_err() { break; } },
                                                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                                                        Err(_) => break,
                                                    }
                                                }
                                            });
                                        }
                                    }
                                }
                            }
                        }

                        // C. Subscribe to REMOTE nodes
                        let is_local_only = req.local.unwrap_or(false);
                        if !is_local_only {
                            if let Some(shard_config) = collection.get_shard_config() {
                                if let Some(coordinator) = &state.shard_coordinator {
                                    let my_addr = coordinator.my_address();
                                    let all_nodes = coordinator.get_collection_nodes(&shard_config);
                                    let cluster_secret = state.cluster_secret();
                                    let mut remote_nodes = std::collections::HashSet::new();
                                    for node_addr in all_nodes {
                                        if node_addr != my_addr {
                                            remote_nodes.insert(node_addr);
                                        }
                                    }

                                    for node_addr in remote_nodes {
                                        let tx_remote = dep_tx.clone();
                                        let db_remote = db_name.clone();
                                        let c_remote = coll_name.clone();
                                        let n_addr = node_addr.clone();
                                        let secret_clone = cluster_secret.clone();

                                        tokio::spawn(async move {
                                            use crate::cluster::ClusterWebsocketClient;
                                            if let Ok(stream) = ClusterWebsocketClient::connect(
                                                &n_addr,
                                                &db_remote,
                                                &c_remote,
                                                true,
                                                &secret_clone,
                                            )
                                            .await
                                            {
                                                tokio::pin!(stream);
                                                while let Some(result) = stream.next().await {
                                                    if let Ok(event) = result {
                                                        if tx_remote.send(event).await.is_err() {
                                                            break;
                                                        }
                                                    } else {
                                                        break;
                                                    }
                                                }
                                            }
                                        });
                                    }
                                }
                            }
                        }
                    }
                }

                drop(dep_tx); // Close original sender

                // 5. Initial Execution
                if !execute_live_query_step(
                    &tx,
                    state.storage.clone(),
                    query_str.clone(),
                    db_name.clone(),
                    state.shard_coordinator.clone(),
                    req.id.clone(),
                )
                .await
                {
                    return;
                }

                // 6. Reactive Loop
                // Coalesce change bursts: a bulk write of N documents used to
                // re-run the query N times per subscriber. After the first
                // event, keep absorbing events until the stream is quiet for
                // DEBOUNCE — capped at MAX_DELAY from the first event so a
                // continuous write stream can't starve the subscriber of
                // updates. Events arriving while the query re-runs stay
                // buffered in the channel and start the next cycle.
                const DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(150);
                const MAX_DELAY: std::time::Duration = std::time::Duration::from_millis(500);
                'reactive: while dep_rx.recv().await.is_some() {
                    let deadline = tokio::time::Instant::now() + MAX_DELAY;
                    loop {
                        tokio::select! {
                            more = dep_rx.recv() => {
                                if more.is_none() {
                                    break 'reactive; // all forwarders gone
                                }
                                if tokio::time::Instant::now() >= deadline {
                                    break; // burst still going — run anyway
                                }
                            }
                            _ = tokio::time::sleep(DEBOUNCE) => break,
                        }
                    }
                    // If the client is gone, the send returns an error and we
                    // bail so the forwarder tasks above can shut down (they
                    // break once dep_rx is dropped here).
                    if !execute_live_query_step(
                        &tx,
                        state.storage.clone(),
                        query_str.clone(),
                        db_name.clone(),
                        state.shard_coordinator.clone(),
                        req.id.clone(),
                    )
                    .await
                    {
                        break;
                    }
                }
            }
            Err(e) => {
                let _ = tx
                    .send(Message::Text(
                        serde_json::json!({
                            "error": format!("Invalid SDBQL query: {}", e)
                        })
                        .to_string()
                        .into(),
                    ))
                    .await;
            }
        }
    } else {
        let _ = tx
            .send(Message::Text(
                serde_json::json!({
                    "error": "Missing 'query' field for live_query"
                })
                .to_string()
                .into(),
            ))
            .await;
    }
}

// Helper for live query execution. Returns false if the client channel has
// been closed (so the caller should stop the reactive loop and let the
// forwarder tasks shut down).
async fn execute_live_query_step(
    tx: &tokio::sync::mpsc::Sender<Message>,
    storage: Arc<StorageEngine>,
    query_str: String,
    db_name: String,
    shard_coordinator: Option<Arc<crate::sharding::ShardCoordinator>>,
    req_id: Option<String>,
) -> bool {
    // Execute SDBQL
    let exec_result = tokio::task::spawn_blocking(move || {
        match crate::sdbql::parser::parse(&query_str) {
            Ok(parsed) => {
                // Security check
                for clause in &parsed.body_clauses {
                    match clause {
                        crate::sdbql::BodyClause::Insert(_)
                        | crate::sdbql::BodyClause::Update(_)
                        | crate::sdbql::BodyClause::Remove(_) => {
                            return Err(crate::error::DbError::ExecutionError(
                                "Live queries are read-only".to_string(),
                            ));
                        }
                        _ => {}
                    }
                }

                let mut executor =
                    crate::sdbql::executor::QueryExecutor::with_database(&storage, db_name);
                if let Some(coord) = shard_coordinator {
                    executor = executor.with_shard_coordinator(coord);
                }
                executor.execute(&parsed)
            }
            Err(e) => Err(crate::error::DbError::ParseError(e.to_string())),
        }
    })
    .await
    .unwrap();

    match exec_result {
        Ok(results) => {
            let mut response = serde_json::json!({
                "type": "query_result",
                "result": results
            });
            if let Some(id) = req_id {
                response["id"] = serde_json::Value::String(id);
            }
            tx.send(Message::Text(response.to_string().into()))
                .await
                .is_ok()
        }
        Err(e) => {
            let mut response = serde_json::json!({
                "type": "error",
                "error": e.to_string()
            });
            if let Some(id) = req_id {
                response["id"] = serde_json::Value::String(id);
            }
            tx.send(Message::Text(response.to_string().into()))
                .await
                .is_ok()
        }
    }
}