//! WebSocket connection handling and event forwarding.
//!
//! Manages WebSocket lifecycle, message send/receive loops, and forwarding
//! agent events to connected clients.
//!
//! ## Binary frame envelope
//!
//! Audio messages use binary WebSocket frames with a self-contained envelope:
//!
//! ```text
//! [4 bytes LE: JSON header length]
//! [N bytes:    JSON header — same {type, data} adjacently-tagged schema]
//! [remaining:  raw binary payload (audio bytes), may be empty]
//! ```
//!
//! All other messages use JSON text frames as before.

#[cfg(feature = "remote")]
use crate::agent::utils::u32_from_usize;

use super::messages::{RoutingMode, UiClientMessage, UiServerMessage};
use super::session::{
    PRIMARY_AGENT_ID, active_profile_id, build_agent_list, list_profiles, mode_for_session,
    reasoning_effort_for_session, session_ref_for_profile,
};
use super::{ConnectionState, ServerState};
use axum::extract::ws::{Message, WebSocket};
use futures_util::{sink::SinkExt, stream::StreamExt as FuturesStreamExt};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::time::{Duration, sleep};
use uuid::Uuid;

// ── Binary frame helpers ───────────────────────────────────────────────────────

/// Parse a binary WebSocket frame into a JSON header string and a binary payload.
///
/// Layout: `[4 bytes LE: header_len][header_len bytes: JSON][remaining: payload]`
pub fn parse_binary_frame(data: &[u8]) -> Result<(&str, &[u8]), String> {
    if data.len() < 4 {
        return Err("Binary frame too short (need at least 4 bytes for header length)".into());
    }
    let header_len = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
    if data.len() < 4 + header_len {
        return Err(format!(
            "Binary frame truncated: header_len={header_len} but only {} bytes remain",
            data.len() - 4
        ));
    }
    let header_bytes = &data[4..4 + header_len];
    let header = std::str::from_utf8(header_bytes)
        .map_err(|e| format!("Binary frame header is not valid UTF-8: {e}"))?;
    let payload = &data[4 + header_len..];
    Ok((header, payload))
}

/// Encode a JSON header and binary payload into a binary WebSocket frame.
pub fn encode_binary_frame(header_json: &str, payload: &[u8]) -> Vec<u8> {
    let header_bytes = header_json.as_bytes();
    let header_len = header_bytes.len() as u32;
    let mut buf = Vec::with_capacity(4 + header_bytes.len() + payload.len());
    buf.extend_from_slice(&header_len.to_le_bytes());
    buf.extend_from_slice(header_bytes);
    buf.extend_from_slice(payload);
    buf
}

/// Handle a new WebSocket connection.
pub async fn handle_websocket_connection(socket: WebSocket, state: ServerState) {
    let conn_id = Uuid::new_v4().to_string();
    let (mut ws_sender, mut ws_receiver) = socket.split();
    let (tx, mut rx) = mpsc::channel::<String>(100);
    // Parallel channel for binary frames (audio responses).
    // Keeps all existing `mpsc::Sender<String>` signatures untouched.
    let (bin_tx, mut bin_rx) = mpsc::channel::<Vec<u8>>(16);

    {
        let mut connections = state.connections.lock().await;
        connections.insert(
            conn_id.clone(),
            ConnectionState {
                routing_mode: RoutingMode::Single,
                active_agent_id: PRIMARY_AGENT_ID.to_string(),
                sessions: HashMap::new(),
                subscribed_sessions: HashSet::new(),
                session_cursors: HashMap::new(),
                current_workspace_root: None,
                file_index_forwarder: None,
            },
        );
    }
    {
        let mut senders = state.connection_senders.lock().await;
        senders.insert(conn_id.clone(), tx.clone());
    }

    spawn_event_forwarders(state.clone(), conn_id.clone(), tx.clone());
    #[cfg(feature = "remote")]
    spawn_peer_event_watcher(state.clone(), tx.clone());
    send_state(&state, &conn_id, &tx).await;

    let send_task = tokio::spawn(async move {
        loop {
            tokio::select! {
                msg = rx.recv() => match msg {
                    Some(json) => {
                        if ws_sender.send(Message::Text(json.into())).await.is_err() {
                            break;
                        }
                    }
                    None => break,
                },
                msg = bin_rx.recv() => match msg {
                    Some(bytes) => {
                        if ws_sender.send(Message::Binary(bytes.into())).await.is_err() {
                            break;
                        }
                    }
                    None => break,
                },
            }
        }
    });

    let state_for_receive = state.clone();
    let conn_id_for_receive = conn_id.clone();
    let tx_for_receive = tx.clone();
    let bin_tx_for_receive = bin_tx.clone();
    let receive_task = tokio::spawn(async move {
        while let Some(result) = FuturesStreamExt::next(&mut ws_receiver).await {
            match result {
                Ok(Message::Text(text)) => {
                    let msg = match serde_json::from_str::<UiClientMessage>(&text) {
                        Ok(msg) => msg,
                        Err(e) => {
                            let _ = send_error(&tx_for_receive, format!("Invalid message: {}", e))
                                .await;
                            continue;
                        }
                    };
                    super::handlers::handle_ui_message(
                        &state_for_receive,
                        &conn_id_for_receive,
                        &tx_for_receive,
                        &bin_tx_for_receive,
                        msg,
                    )
                    .await;
                }
                Ok(Message::Binary(data)) => {
                    let (header, payload) = match parse_binary_frame(&data) {
                        Ok(parsed) => parsed,
                        Err(e) => {
                            let _ =
                                send_error(&tx_for_receive, format!("Invalid binary frame: {e}"))
                                    .await;
                            continue;
                        }
                    };
                    let msg = match serde_json::from_str::<UiClientMessage>(header) {
                        Ok(msg) => msg,
                        Err(e) => {
                            let _ = send_error(
                                &tx_for_receive,
                                format!("Invalid binary frame header: {e}"),
                            )
                            .await;
                            continue;
                        }
                    };
                    super::handlers::handle_binary_message(
                        &state_for_receive,
                        &conn_id_for_receive,
                        &tx_for_receive,
                        &bin_tx_for_receive,
                        msg,
                        payload.to_vec(),
                    )
                    .await;
                }
                Ok(Message::Close(_)) => break,
                Ok(Message::Ping(_)) => {}
                Ok(_) => {}
                Err(e) => {
                    log::error!("UI WebSocket error: {}", e);
                    break;
                }
            }
        }
    });

    tokio::select! {
        _ = send_task => {},
        _ = receive_task => {},
    }

    {
        let mut connections = state.connections.lock().await;
        if let Some(conn) = connections.remove(&conn_id)
            && let Some(handle) = conn.file_index_forwarder
        {
            handle.abort();
        }
    }
    {
        let mut senders = state.connection_senders.lock().await;
        senders.remove(&conn_id);
    }

    // Clean up OAuth state (pending flows + callback listener) for this connection.
    super::handlers::stop_oauth_callback_listener_for_connection(&state, &conn_id).await;
}

/// Send current state to a client.
pub async fn send_state(state: &ServerState, conn_id: &str, tx: &mpsc::Sender<String>) {
    let (routing_mode, active_agent_id, active_session_id, sessions_by_agent) = {
        let connections = state.connections.lock().await;
        if let Some(conn) = connections.get(conn_id) {
            (
                conn.routing_mode,
                conn.active_agent_id.clone(),
                conn.sessions.get(&conn.active_agent_id).cloned(),
                conn.sessions.clone(),
            )
        } else {
            (
                RoutingMode::Single,
                PRIMARY_AGENT_ID.to_string(),
                None,
                HashMap::new(),
            )
        }
    };

    let agents = build_agent_list(state).await;
    let profiles = match list_profiles(state).await {
        Ok(profiles) => profiles,
        Err(message) => {
            let _ = send_error(tx, message).await;
            Vec::new()
        }
    };
    let active_profile_id = active_profile_id(state).await;
    let default_cwd = state
        .default_cwd
        .as_ref()
        .map(|path| path.to_string_lossy().to_string());

    let agent_mode = mode_for_session(state, active_session_id.as_deref()).await;
    let reasoning_effort = reasoning_effort_for_session(state, active_session_id.as_deref()).await;

    let _ = send_message(
        tx,
        UiServerMessage::State {
            routing_mode,
            active_agent_id,
            active_session_id,
            default_cwd,
            agents,
            profiles,
            active_profile_id,
            sessions_by_agent,
            agent_mode: agent_mode.as_str().to_string(),
            reasoning_effort: reasoning_effort.map(|e| e.to_string()),
        },
    )
    .await;
}

/// Send a message to the client.
pub async fn send_message(
    tx: &mpsc::Sender<String>,
    message: UiServerMessage,
) -> Result<(), String> {
    let message_type = message.type_name();

    match serde_json::to_string(&message) {
        Ok(json) => match tx.send(json).await {
            Ok(_) => Ok(()),
            Err(e) => {
                log::error!("send_message: failed to send {}: {}", message_type, e);
                Err(format!("Failed to send: {}", e))
            }
        },
        Err(err) => {
            log::error!(
                "send_message: failed to serialize {}: {}",
                message_type,
                err
            );
            Err(format!("Failed to serialize: {}", err))
        }
    }
}

/// Send an error message to the client.
pub async fn send_error(tx: &mpsc::Sender<String>, message: String) -> Result<(), String> {
    log::debug!("send_error: sending error message: {}", message);
    send_message(tx, UiServerMessage::Error { message }).await
}

/// Send a binary frame to the client via the binary channel.
///
/// `header_json` is the JSON header (same `{type, data}` schema as text messages).
/// `payload` is the raw binary data (e.g. audio bytes).
pub async fn send_binary(
    bin_tx: &mpsc::Sender<Vec<u8>>,
    header_json: &str,
    payload: &[u8],
) -> Result<(), String> {
    let frame = encode_binary_frame(header_json, payload);
    bin_tx
        .send(frame)
        .await
        .map_err(|e| format!("Failed to send binary frame: {e}"))
}

/// Spawn event forwarders for all event sources.
pub fn spawn_event_forwarders(state: ServerState, conn_id: String, tx: mpsc::Sender<String>) {
    let mut seen_sources = HashSet::new();
    for event_source in &state.event_sources {
        spawn_event_source_forwarder(
            state.clone(),
            conn_id.clone(),
            tx.clone(),
            event_source.clone(),
        );
        seen_sources.insert(Arc::as_ptr(event_source) as usize);
    }

    if let Some(profiles) = state.profiles.clone() {
        tokio::spawn(async move {
            loop {
                if !connection_is_open(&state, &conn_id, &tx).await {
                    break;
                }

                let runtimes = profiles.materialized_runtimes().await;
                for runtime in runtimes {
                    for event_source in
                        super::session::collect_event_sources(&runtime.agent().handle())
                    {
                        if seen_sources.insert(Arc::as_ptr(&event_source) as usize) {
                            spawn_event_source_forwarder(
                                state.clone(),
                                conn_id.clone(),
                                tx.clone(),
                                event_source,
                            );
                        }
                    }
                }

                sleep(Duration::from_millis(100)).await;
            }
        });
    }
}

async fn connection_is_open(state: &ServerState, conn_id: &str, tx: &mpsc::Sender<String>) -> bool {
    !tx.is_closed() && state.connections.lock().await.contains_key(conn_id)
}

fn spawn_event_source_forwarder(
    state: ServerState,
    conn_id: String,
    tx: mpsc::Sender<String>,
    event_source: Arc<crate::event_fanout::EventFanout>,
) {
    let mut events = event_source.subscribe();
    tokio::spawn(async move {
        while let Ok(event) = events.recv().await {
            if forward_event_to_ui(&state, &conn_id, &tx, event)
                .await
                .is_err()
            {
                break;
            }
        }
    });
}

async fn forward_event_to_ui(
    state: &ServerState,
    conn_id: &str,
    tx: &mpsc::Sender<String>,
    event: crate::events::EventEnvelope,
) -> Result<(), ()> {
    // Check if this connection is subscribed and if this event is newer than
    // the replay cursor (prevents replay/live overlap duplicates). Ephemeral
    // events (seq=0) are live-only and are never cursor-filtered.
    let should_forward = {
        let mut connections = state.connections.lock().await;
        if let Some(conn) = connections.get_mut(conn_id) {
            if !conn.subscribed_sessions.contains(event.session_id()) {
                false
            } else if event.seq() == 0 {
                true
            } else {
                let cursor = conn
                    .session_cursors
                    .get(event.session_id())
                    .cloned()
                    .unwrap_or_default();
                match event.origin() {
                    crate::events::EventOrigin::Local | crate::events::EventOrigin::Unknown(_) => {
                        if event.seq() <= cursor.local_seq {
                            log::debug!(
                                "ui forwarder: dropping local duplicate/overlap event conn={} session={} seq={} cursor_local_seq={}",
                                conn_id,
                                event.session_id(),
                                event.seq(),
                                cursor.local_seq
                            );
                            false
                        } else {
                            true
                        }
                    }
                    crate::events::EventOrigin::Remote => {
                        if let Some(source_node) = event.source_node() {
                            let source_cursor = cursor
                                .remote_seq_by_source
                                .get(source_node)
                                .copied()
                                .unwrap_or(0);
                            if event.seq() <= source_cursor {
                                log::debug!(
                                    "ui forwarder: dropping remote duplicate/overlap event conn={} session={} source={} seq={} cursor_remote_seq={}",
                                    conn_id,
                                    event.session_id(),
                                    source_node,
                                    event.seq(),
                                    source_cursor
                                );
                                false
                            } else {
                                true
                            }
                        } else {
                            log::warn!(
                                "ui forwarder: dropping remote event without source conn={} session={} seq={}",
                                conn_id,
                                event.session_id(),
                                event.seq(),
                            );
                            false
                        }
                    }
                }
            }
        } else {
            false
        }
    };

    if !should_forward {
        return Ok(());
    }

    // React to WorkspaceIndexReady from a remote session: push status +
    // fetch and push the file index so the UI gets it without polling.
    if let crate::events::AgentEventKind::WorkspaceIndexReady { .. } = event.kind() {
        let session_id = event.session_id().to_owned();

        let _ = send_message(
            tx,
            UiServerMessage::WorkspaceIndexStatus {
                session_id: session_id.clone(),
                status: "ready".to_string(),
                message: None,
            },
        )
        .await;

        let profile_id = profile_id_for_session(state, &session_id).await;
        if let Some(actor_ref) =
            session_ref_for_profile(state, profile_id.as_deref(), &session_id).await
        {
            let cwd = {
                let cwds = state.session_cwds.lock().await;
                cwds.get(&session_id).cloned()
            };
            match actor_ref.get_file_index().await {
                Ok(resp) => {
                    use super::mentions::filter_index_for_cwd_entries;
                    let root = std::path::PathBuf::from(&resp.workspace_root);
                    let files = match cwd.as_ref().and_then(|c| c.strip_prefix(&root).ok()) {
                        Some(relative_cwd) => {
                            filter_index_for_cwd_entries(&resp.files, relative_cwd)
                        }
                        None => resp.files,
                    };
                    let _ = send_message(
                        tx,
                        UiServerMessage::FileIndex {
                            files,
                            generated_at: resp.generated_at,
                        },
                    )
                    .await;
                }
                Err(e) => {
                    log::warn!(
                        "spawn_event_forwarders: WorkspaceIndexReady for {} but get_file_index failed: {}",
                        session_id,
                        e
                    );
                }
            }
        }
        return Ok(());
    }

    let agent_id = {
        let agents = state.session_agents.lock().await;
        agents
            .get(event.session_id())
            .cloned()
            .unwrap_or_else(|| "unknown".to_string())
    };
    let profile_id = profile_id_for_session(state, event.session_id()).await;

    if log::log_enabled!(log::Level::Trace)
        && matches!(
            event.kind(),
            crate::events::AgentEventKind::AssistantThinkingDelta { .. }
                | crate::events::AgentEventKind::AssistantContentDelta { .. }
                | crate::events::AgentEventKind::AssistantMessageStored { .. }
                | crate::events::AgentEventKind::LlmRequestStart { .. }
                | crate::events::AgentEventKind::LlmRequestEnd { .. }
        )
    {
        log::trace!(
            "ui forwarder: conn={} session={} agent={} seq={} kind={:?}",
            conn_id,
            event.session_id(),
            agent_id,
            event.seq(),
            event.kind()
        );
    }

    if send_message(
        tx,
        UiServerMessage::Event {
            agent_id,
            profile_id,
            session_id: event.session_id().to_owned(),
            event: event.clone(),
        },
    )
    .await
    .is_err()
    {
        return Err(());
    }

    update_live_cursor(state, conn_id, &event).await;
    Ok(())
}

async fn profile_id_for_session(state: &ServerState, session_id: &str) -> Option<String> {
    match &state.profiles {
        Some(profiles) => profiles
            .session_binding(session_id)
            .await
            .map(|binding| binding.profile_id),
        None => None,
    }
}

async fn update_live_cursor(
    state: &ServerState,
    conn_id: &str,
    event: &crate::events::EventEnvelope,
) {
    let mut connections = state.connections.lock().await;
    if let Some(conn) = connections.get_mut(conn_id) {
        let cursor = conn
            .session_cursors
            .entry(event.session_id().to_owned())
            .or_default();
        match event.origin() {
            crate::events::EventOrigin::Local | crate::events::EventOrigin::Unknown(_) => {
                cursor.local_seq = cursor.local_seq.max(event.seq());
            }
            crate::events::EventOrigin::Remote => {
                if let Some(source_node) = event.source_node() {
                    cursor
                        .remote_seq_by_source
                        .entry(source_node.to_owned())
                        .and_modify(|seq| *seq = (*seq).max(event.seq()))
                        .or_insert(event.seq());
                }
            }
        }
    }
}

/// Spawn a task that watches for mDNS peer discovery / expiry events and pushes
/// an updated `remote_nodes` list to this WebSocket client whenever the mesh
/// topology changes.
///
/// This replaces polling: the kameo swarm event loop emits a `PeerEvent` the
/// instant mDNS fires, and we react immediately by re-querying the DHT and
/// pushing the fresh node list to the client.
#[cfg(feature = "remote")]
pub fn spawn_peer_event_watcher(state: ServerState, tx: mpsc::Sender<String>) {
    use super::messages::RemoteNodeInfo;
    use crate::agent::remote::PeerEvent;

    let Some(mesh) = state.agent.mesh() else {
        return;
    };

    let mut rx = mesh.subscribe_peer_events();

    tokio::spawn(async move {
        loop {
            match rx.recv().await {
                Ok(event) => {
                    if tx.is_closed() {
                        break;
                    }

                    // On Expired we can push the update immediately because
                    // the node is gone and the DHT doesn't need time to settle.
                    // On Discovered we use a retry loop: the remote node needs
                    // time to bootstrap and register its actors in the DHT, so
                    // we poll with exponential back-off rather than a fixed sleep.
                    let needs_retry = match event {
                        PeerEvent::Discovered(peer_id) => {
                            log::debug!(
                                "spawn_peer_event_watcher: peer discovered {peer_id}, will poll until actor is visible"
                            );
                            // Invalidate cache — topology changed.
                            state.invalidate_remote_node_cache().await;
                            // Also invalidate remote model inventory so model list
                            // is refreshed when the peer becomes visible.
                            state.agent.model_inventory.invalidate_remote().await;
                            true
                        }
                        PeerEvent::Expired(peer_id) => {
                            log::debug!(
                                "spawn_peer_event_watcher: peer expired {peer_id}, refreshing remote nodes"
                            );
                            // Invalidate cache — topology changed.
                            state.invalidate_remote_node_cache().await;
                            // Also invalidate remote model inventory so stale
                            // remote models are removed from the list.
                            state.agent.model_inventory.invalidate_remote().await;
                            false
                        }
                        _ => {
                            continue;
                        }
                    };

                    if tx.is_closed() {
                        break;
                    }

                    if needs_retry {
                        // Exponential back-off: 500 ms, 1 s, 2 s, 4 s, 8 s.
                        // We stop as soon as the newly discovered peer shows up
                        // in list_remote_nodes(), giving up after ~15 s total.
                        //
                        // Invalidate cache before each retry so we always get a
                        // fresh DHT snapshot.
                        const DELAYS_MS: &[u64] = &[500, 1_000, 2_000, 4_000, 8_000];
                        let mut pushed = false;
                        for &delay_ms in DELAYS_MS {
                            tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;

                            if tx.is_closed() {
                                break;
                            }

                            state.invalidate_remote_node_cache().await;
                            let nodes = state.get_remote_nodes_cached().await;

                            let msg = UiServerMessage::RemoteNodes {
                                nodes: nodes
                                    .iter()
                                    .map(|n| RemoteNodeInfo {
                                        id: n.node_id.to_string(),
                                        label: n.hostname.clone(),
                                        capabilities: n.capabilities.clone(),
                                        active_sessions: u32_from_usize(
                                            n.active_sessions,
                                            "active_sessions",
                                            None,
                                        ),
                                    })
                                    .collect(),
                            };

                            if send_message(&tx, msg).await.is_err() {
                                break;
                            }

                            // Stop retrying once we can see at least one remote node
                            // — the UI has been updated and the peer is reachable.
                            if !nodes.is_empty() {
                                pushed = true;
                                break;
                            }

                            log::debug!(
                                "spawn_peer_event_watcher: no remote nodes yet, retrying in {}ms",
                                delay_ms * 2
                            );
                        }

                        if !pushed {
                            log::warn!(
                                "spawn_peer_event_watcher: peer discovered but no remote nodes \
                                 visible after all retries — peer may not have registered its actors yet"
                            );
                        }
                    } else {
                        let nodes = state.get_remote_nodes_cached().await;

                        // Clear provider_node_id for any session that was using the
                        // expired peer. We detect stale sessions by comparing the
                        // remote_sessions() peer_label against the set of nodes
                        // still visible after expiry.
                        let available_hostnames: HashSet<&str> =
                            nodes.iter().map(|n| n.hostname.as_str()).collect();
                        let remote_sessions = {
                            let registry = state.agent.registry.lock().await;
                            registry.remote_sessions()
                        };

                        for (session_id, peer_label, _) in remote_sessions {
                            if !available_hostnames.contains(peer_label.as_str()) {
                                match state.agent.get_session_llm_config(&session_id).await {
                                    Ok(Some(config)) => {
                                        state.agent.emit_event(
                                            &session_id,
                                            crate::events::AgentEventKind::ProviderChanged {
                                                provider: config.provider,
                                                model: config.model,
                                                config_id: config.id,
                                                context_limit: None,
                                                provider_node_id: None,
                                            },
                                        );
                                        log::info!(
                                            "spawn_peer_event_watcher: cleared provider_node_id \
                                             for session {} (peer {} expired)",
                                            session_id,
                                            peer_label
                                        );
                                    }
                                    Ok(None) => {
                                        log::debug!(
                                            "spawn_peer_event_watcher: no LLM config for \
                                             session {} (peer {} expired)",
                                            session_id,
                                            peer_label
                                        );
                                    }
                                    Err(e) => {
                                        log::warn!(
                                            "spawn_peer_event_watcher: failed to get LLM \
                                             config for session {} (peer {} expired): {}",
                                            session_id,
                                            peer_label,
                                            e
                                        );
                                    }
                                }
                            }
                        }

                        let msg = UiServerMessage::RemoteNodes {
                            nodes: nodes
                                .into_iter()
                                .map(|n| RemoteNodeInfo {
                                    id: n.node_id.to_string(),
                                    label: n.hostname,
                                    capabilities: n.capabilities,
                                    active_sessions: u32_from_usize(
                                        n.active_sessions,
                                        "active_sessions",
                                        None,
                                    ),
                                })
                                .collect(),
                        };

                        if send_message(&tx, msg).await.is_err() {
                            break;
                        }

                        // Schedule a delayed re-check to catch any stale DHT
                        // records that resolved between the immediate push and
                        // the DHT eventually purging the expired peer.
                        let state_delayed = state.clone();
                        let tx_delayed = tx.clone();
                        tokio::spawn(async move {
                            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
                            if tx_delayed.is_closed() {
                                return;
                            }
                            let nodes = state_delayed.get_remote_nodes_cached().await;
                            let msg = UiServerMessage::RemoteNodes {
                                nodes: nodes
                                    .into_iter()
                                    .map(|n| RemoteNodeInfo {
                                        id: n.node_id.to_string(),
                                        label: n.hostname,
                                        capabilities: n.capabilities,
                                        active_sessions: u32_from_usize(
                                            n.active_sessions,
                                            "active_sessions",
                                            None,
                                        ),
                                    })
                                    .collect(),
                            };
                            let _ = send_message(&tx_delayed, msg).await;
                        });
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                    // We missed some events — re-query immediately to catch up.
                    log::warn!("spawn_peer_event_watcher: lagged by {n} events, re-querying");
                    state.invalidate_remote_node_cache().await;
                    state.agent.model_inventory.invalidate_remote().await;
                    let nodes = state.get_remote_nodes_cached().await;
                    let msg = UiServerMessage::RemoteNodes {
                        nodes: nodes
                            .into_iter()
                            .map(|n| RemoteNodeInfo {
                                id: n.node_id.to_string(),
                                label: n.hostname,
                                capabilities: n.capabilities,
                                active_sessions: u32_from_usize(
                                    n.active_sessions,
                                    "active_sessions",
                                    None,
                                ),
                            })
                            .collect(),
                    };
                    if send_message(&tx, msg).await.is_err() {
                        break;
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                    // The swarm is shutting down — nothing to do.
                    break;
                }
            }
        }
    });
}

/// Subscribe to file index updates for a workspace and forward them to the WebSocket client.
///
/// This spawns a long-lived task that listens to file index broadcasts from the FileIndexWatcher
/// and sends FileIndex messages to the client whenever files are created/modified/deleted.
pub async fn subscribe_to_file_index(
    state: ServerState,
    conn_id: String,
    tx: mpsc::Sender<String>,
    workspace_root: std::path::PathBuf,
) {
    // Get the workspace and subscribe to its file index updates
    let manager = &state.workspace_manager;
    let handle =
        match crate::index::get_or_create_workspace_with_timeout(manager, workspace_root.clone())
            .await
        {
            Ok(handle) => handle,
            Err(err) => {
                log::error!(
                    "Failed to get workspace for file index subscription: {}",
                    err
                );
                return;
            }
        };

    // Ensure each connection has at most one active forwarder task.
    let existing_forwarder = {
        let mut connections = state.connections.lock().await;
        let Some(conn) = connections.get_mut(&conn_id) else {
            return;
        };

        let already_subscribed = conn.current_workspace_root.as_ref() == Some(&workspace_root)
            && conn
                .file_index_forwarder
                .as_ref()
                .is_some_and(|handle| !handle.is_finished());
        if already_subscribed {
            log::debug!(
                "Connection {} is already subscribed to workspace {:?}",
                conn_id,
                workspace_root
            );
            return;
        }

        conn.current_workspace_root = Some(workspace_root.clone());
        conn.file_index_forwarder.take()
    };

    if let Some(handle) = existing_forwarder {
        handle.abort();
    }

    let mut index_rx = handle.file_watcher.subscribe_index();

    let state_for_forward = state.clone();
    let conn_id_for_forward = conn_id.clone();
    let workspace_root_for_forward = workspace_root.clone();
    let tx_for_forward = tx.clone();

    // Spawn a task to forward file index updates to the client
    let forwarder = tokio::spawn(async move {
        while let Ok(index) = index_rx.recv().await {
            // Get the current session's cwd to filter the index appropriately
            let session_id = {
                let connections = state_for_forward.connections.lock().await;
                let conn = match connections.get(&conn_id_for_forward) {
                    Some(conn) => conn,
                    None => {
                        // Connection closed, stop forwarding
                        log::debug!(
                            "Connection {} closed, stopping file index forwarding",
                            conn_id_for_forward
                        );
                        break;
                    }
                };

                conn.sessions.get(&conn.active_agent_id).cloned()
            };

            let session_id = match session_id {
                Some(id) => id,
                None => continue, // No active session, skip this update
            };

            let cwd = {
                let cwds = state_for_forward.session_cwds.lock().await;
                match cwds.get(&session_id).cloned() {
                    Some(cwd) => cwd,
                    None => continue, // No cwd for this session, skip
                }
            };

            // Filter the index to the session's cwd
            let relative_cwd = match cwd.strip_prefix(&workspace_root_for_forward) {
                Ok(relative) => relative,
                Err(_) => continue, // cwd outside workspace root, skip
            };

            let files = super::mentions::filter_index_for_cwd(&index, relative_cwd);

            // Send the filtered index to the client
            if send_message(
                &tx_for_forward,
                UiServerMessage::FileIndex {
                    files,
                    generated_at: index.generated_at,
                },
            )
            .await
            .is_err()
            {
                log::debug!(
                    "Failed to send file index to connection {}, stopping forwarding",
                    conn_id_for_forward
                );
                break;
            }

            log::debug!(
                "Pushed file index update to connection {}",
                conn_id_for_forward
            );
        }
    });

    {
        let mut connections = state.connections.lock().await;
        if let Some(conn) = connections.get_mut(&conn_id) {
            conn.file_index_forwarder = Some(forwarder);
        } else {
            forwarder.abort();
            return;
        }
    }

    log::info!(
        "Subscribed connection {} to file index updates for workspace {:?}",
        conn_id,
        workspace_root
    );

    // IMPORTANT: Send the current index immediately to avoid subscription race condition.
    // When subscribing to a broadcast channel, we only receive FUTURE messages.
    // If the workspace already existed (cached), the initial index was sent before we subscribed.
    // This ensures new subscribers get the current state immediately.
    if let Some(current_index) = handle.file_index() {
        // Get current session's cwd to filter the index
        let session_id = {
            let connections = state.connections.lock().await;
            connections
                .get(&conn_id)
                .and_then(|conn| conn.sessions.get(&conn.active_agent_id).cloned())
        };
        let cwd = if let Some(session_id) = session_id {
            let cwds = state.session_cwds.lock().await;
            cwds.get(&session_id).cloned()
        } else {
            None
        };

        if let Some(cwd) = cwd
            && let Ok(relative_cwd) = cwd.strip_prefix(&workspace_root)
        {
            let files = super::mentions::filter_index_for_cwd(&current_index, relative_cwd);

            // Send initial index
            if let Err(err) = send_message(
                &tx,
                UiServerMessage::FileIndex {
                    files,
                    generated_at: current_index.generated_at,
                },
            )
            .await
            {
                log::warn!(
                    "Failed to send initial file index to connection {}: {}",
                    conn_id,
                    err
                );
            } else {
                log::debug!("Sent initial file index to connection {}", conn_id);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::api::AgentInfra;
    use crate::profiles::{LocalProfileCatalog, ProfileCatalog, ProfileRuntimeManager};
    use crate::test_utils::{TestServerState, empty_plugin_registry};
    use std::path::Path;
    use std::sync::Arc;

    fn write_quorum_profile(dir: &Path) {
        std::fs::write(
            dir.join("delegates.toml"),
            r#"
[profile]
id = "delegates"
name = "Delegates"

[quorum]

[planner]
provider = "test"
model = "planner-model"
system = "plan"

[[delegates]]
id = "coder"
provider = "test"
model = "coder-model"
description = "Writes code."
capabilities = ["coding"]
system = "code"

[[delegates]]
id = "explorer"
provider = "test"
model = "explorer-model"
description = "Explores code."
capabilities = ["exploration"]
system = "explore"
"#,
        )
        .expect("profile should be written");
    }

    // ── Binary frame envelope tests ─────────────────────────────────────────

    #[test]
    fn binary_frame_round_trip() {
        let header = r#"{"type":"transcribe","data":{"provider":"izwi","model":"Qwen3-ASR-0.6B"}}"#;
        let payload = b"fake audio bytes";

        let encoded = encode_binary_frame(header, payload);
        let (decoded_header, decoded_payload) = parse_binary_frame(&encoded).expect("should parse");

        assert_eq!(decoded_header, header);
        assert_eq!(decoded_payload, payload);
    }

    #[test]
    fn binary_frame_round_trip_empty_payload() {
        let header =
            r#"{"type":"speech","data":{"provider":"izwi","model":"Kokoro-82M","text":"hi"}}"#;
        let payload: &[u8] = b"";

        let encoded = encode_binary_frame(header, payload);
        let (decoded_header, decoded_payload) = parse_binary_frame(&encoded).expect("should parse");

        assert_eq!(decoded_header, header);
        assert!(decoded_payload.is_empty());
    }

    #[test]
    fn binary_frame_round_trip_large_payload() {
        let header = r#"{"type":"transcribe","data":{"provider":"izwi","model":"Qwen3-ASR-0.6B"}}"#;
        // Simulate a ~320KB audio clip (10s @ 16kHz mono 16-bit)
        let payload = vec![0xABu8; 320_000];

        let encoded = encode_binary_frame(header, &payload);
        let (decoded_header, decoded_payload) = parse_binary_frame(&encoded).expect("should parse");

        assert_eq!(decoded_header, header);
        assert_eq!(decoded_payload.len(), 320_000);
        assert_eq!(decoded_payload[0], 0xAB);
    }

    #[test]
    fn parse_binary_frame_rejects_too_short() {
        let result = parse_binary_frame(&[0, 1, 2]);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("too short"));
    }

    #[test]
    fn parse_binary_frame_rejects_truncated_header() {
        // header_len says 100 bytes but only 4 bytes of actual data follow
        let mut buf = Vec::new();
        buf.extend_from_slice(&100u32.to_le_bytes());
        buf.extend_from_slice(b"tiny");
        let result = parse_binary_frame(&buf);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("truncated"));
    }

    #[test]
    fn parse_binary_frame_rejects_invalid_utf8_header() {
        let bad_header: &[u8] = &[0xFF, 0xFE, 0xFD];
        let mut buf = Vec::new();
        buf.extend_from_slice(&(bad_header.len() as u32).to_le_bytes());
        buf.extend_from_slice(bad_header);
        let result = parse_binary_frame(&buf);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("UTF-8"));
    }

    #[test]
    fn binary_frame_header_deserializes_as_transcribe() {
        let header = r#"{"type":"transcribe","data":{"provider":"izwi","model":"Qwen3-ASR-0.6B","mime_type":"audio/wav"}}"#;
        let payload = b"wav-bytes";
        let encoded = encode_binary_frame(header, payload);
        let (decoded_header, _) = parse_binary_frame(&encoded).expect("should parse");

        let msg: UiClientMessage =
            serde_json::from_str(decoded_header).expect("should deserialize header");

        match msg {
            UiClientMessage::Transcribe {
                provider,
                model,
                mime_type,
            } => {
                assert_eq!(provider, "izwi");
                assert_eq!(model, "Qwen3-ASR-0.6B");
                assert_eq!(mime_type.as_deref(), Some("audio/wav"));
            }
            _ => panic!("expected Transcribe variant"),
        }
    }

    #[test]
    fn encode_binary_frame_layout() {
        let header = "hello";
        let payload = b"world";

        let encoded = encode_binary_frame(header, payload);

        // First 4 bytes: header length as u32 LE
        assert_eq!(&encoded[0..4], &5u32.to_le_bytes());
        // Next 5 bytes: header
        assert_eq!(&encoded[4..9], b"hello");
        // Remaining: payload
        assert_eq!(&encoded[9..], b"world");
    }

    #[tokio::test]
    async fn send_state_reports_delegates_from_active_profile_runtime() {
        let mut fixture = TestServerState::new().await;
        let dir = tempfile::TempDir::new().expect("temp profile dir");
        write_quorum_profile(dir.path());
        let catalog: Arc<dyn ProfileCatalog> = Arc::new(
            LocalProfileCatalog::builder()
                .include_embedded_default(false)
                .local_dir(dir.path())
                .build(),
        );
        let (registry, _registry_dir) = empty_plugin_registry().expect("empty plugin registry");
        let infra = AgentInfra {
            plugin_registry: Arc::new(registry),
            storage: Some(fixture.agent.storage.clone()),
            session_mcp_attachment_source: None,
        };
        fixture.state.profiles = Some(Arc::new(ProfileRuntimeManager::with_infra_boxed(
            catalog,
            "delegates",
            infra,
        )));
        let conn_id = "conn-state-profile-agents";
        let (tx, mut rx) = fixture.add_connection(conn_id).await;

        send_state(&fixture.state, conn_id, &tx).await;

        let raw = rx.recv().await.expect("state message should be sent");
        let value: serde_json::Value = serde_json::from_str(&raw).expect("state should be JSON");
        let mut agent_ids: Vec<String> = value["data"]["agents"]
            .as_array()
            .expect("agents should be an array")
            .iter()
            .map(|agent| agent["id"].as_str().expect("agent id").to_string())
            .collect();

        assert_eq!(agent_ids.remove(0), "primary");
        agent_ids.sort();
        assert_eq!(agent_ids, vec!["coder", "explorer"]);
    }
}
