//! WebSocket agent chat handler.
//!
//! Connect: `ws://host:port/ws/chat?session_id=ID&name=My+Session`
//!
//! Protocol:
//! ```text
//! Server -> Client: {"type":"session_start","session_id":"...","name":"...","resumed":true,"message_count":42}
//! Client -> Server: {"type":"message","content":"Hello"}
//! Server -> Client: {"type":"chunk","content":"Hi! "}
//! Server -> Client: {"type":"tool_call","name":"shell","args":{...}}
//! Server -> Client: {"type":"tool_result","name":"shell","output":"..."}
//! Server -> Client: {"type":"done","full_response":"..."}
//! ```
//!
//! Query params:
//! - `session_id` — resume or create a session (default: new UUID)
//! - `name` — optional human-readable label for the session
//! - `token` — bearer auth token (alternative to Authorization header)

use super::AppState;
use axum::{
    extract::{
        Query, State, WebSocketUpgrade,
        ws::{Message, WebSocket},
    },
    http::{HeaderMap, header},
    response::IntoResponse,
};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use tokio_util::sync::CancellationToken;
use tracing::debug;
use zeroclaw_api::media::MediaAttachment;

/// Optional connection parameters sent as the first WebSocket message.
///
/// If the first message after upgrade is `{"type":"connect",...}`, these
/// parameters are extracted and an acknowledgement is sent back. Old clients
/// that send `{"type":"message",...}` as the first frame still work — the
/// message is processed normally (backward-compatible).
#[derive(Debug, Deserialize)]
struct ConnectParams {
    #[serde(rename = "type")]
    msg_type: String,
    /// Client-chosen session ID for memory persistence
    #[serde(default)]
    session_id: Option<String>,
    /// Device name for device registry tracking
    #[serde(default)]
    device_name: Option<String>,
    /// Client capabilities
    #[serde(default)]
    capabilities: Vec<String>,
}

/// Supported sub-protocols for the chat WebSocket (primary + legacy).
const WS_PROTOCOLS: &[&str] = &["zeroclaw.v1", "kylinbot.v1"];

/// Prefix used in `Sec-WebSocket-Protocol` to carry a bearer token.
const BEARER_SUBPROTO_PREFIX: &str = "bearer.";

#[derive(Deserialize)]
pub struct WsQuery {
    pub token: Option<String>,
    pub session_id: Option<String>,
    /// Optional human-readable name for the session.
    pub name: Option<String>,
}

#[derive(Debug, PartialEq, Eq)]
enum WsInboundAction {
    Message {
        content: String,
        attachments: Vec<MediaAttachment>,
    },
    Cancel,
    SetModel {
        provider: Option<String>,
        model: Option<String>,
        api_key: Option<String>,
    },
    SetIdentity {
        name: Option<String>,
        role: Option<String>,
    },
    Ignore,
    Error(&'static str),
}

/// Extract a bearer token from WebSocket-compatible sources.
///
/// Precedence (first non-empty wins):
/// 1. `Authorization: Bearer <token>` header
/// 2. `Sec-WebSocket-Protocol: bearer.<token>` subprotocol
/// 3. `?token=<token>` query parameter
///
/// Browsers cannot set custom headers on `new WebSocket(url)`, so the query
/// parameter and subprotocol paths are required for browser-based clients.
fn extract_ws_token<'a>(headers: &'a HeaderMap, query_token: Option<&'a str>) -> Option<&'a str> {
    // 1. Authorization header
    if let Some(t) = headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|auth| auth.strip_prefix("Bearer "))
        && !t.is_empty()
    {
        return Some(t);
    }

    // 2. Sec-WebSocket-Protocol: bearer.<token>
    if let Some(t) = headers
        .get("sec-websocket-protocol")
        .and_then(|v| v.to_str().ok())
        .and_then(|protos| {
            protos
                .split(',')
                .map(|p| p.trim())
                .find_map(|p| p.strip_prefix(BEARER_SUBPROTO_PREFIX))
        })
        && !t.is_empty()
    {
        return Some(t);
    }

    // 3. ?token= query parameter
    if let Some(t) = query_token
        && !t.is_empty()
    {
        return Some(t);
    }

    None
}

fn parse_ws_inbound_action(raw: &str) -> WsInboundAction {
    let parsed: serde_json::Value = match serde_json::from_str(raw) {
        Ok(v) => v,
        Err(_) => return WsInboundAction::Error("Invalid JSON"),
    };

    match parsed["type"].as_str().unwrap_or("") {
        "message" => {
            let content = parsed["content"].as_str().unwrap_or("").trim().to_string();
            let attachments = match parse_inbound_attachments(&parsed) {
                Ok(attachments) => attachments,
                Err(message) => return WsInboundAction::Error(message),
            };
            if content.is_empty() && attachments.is_empty() {
                WsInboundAction::Error("Message content cannot be empty")
            } else {
                WsInboundAction::Message {
                    content,
                    attachments,
                }
            }
        }
        "cancel" => WsInboundAction::Cancel,
        "set_model" => WsInboundAction::SetModel {
            provider: parsed["provider"]
                .as_str()
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned),
            model: parsed["model"]
                .as_str()
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned),
            api_key: parsed["api_key"]
                .as_str()
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned),
        },
        "set_identity" => WsInboundAction::SetIdentity {
            name: parsed["name"]
                .as_str()
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned),
            role: parsed["role"]
                .as_str()
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned),
        },
        _ => WsInboundAction::Ignore,
    }
}

fn build_identity_markdown(name: Option<&str>, role: Option<&str>) -> Option<String> {
    if name.is_none() && role.is_none() {
        return None;
    }

    let mut content = String::from("# IDENTITY.md — Who Am I?\n\n");
    if let Some(name) = name {
        content.push_str(&format!("I am {name}.\n\n"));
    }
    if let Some(role) = role {
        content.push_str("## Role\n");
        content.push_str(role);
        content.push('\n');
    }
    Some(content)
}

fn build_soul_identity_section(name: &str, role: Option<&str>) -> String {
    let role = role.unwrap_or("帮忙做很多事情");
    format!(
        "## Identity\n\n你的名字是**{name}**。你是一个智能助手。\n你可以**{role}**\n当被问到你叫什么名字时，回答“{name}”。"
    )
}

fn replace_identity_section(soul: &str, new_identity: &str) -> String {
    let mut result = String::new();
    let mut in_identity = false;
    let mut found_identity = false;

    for line in soul.lines() {
        if line.starts_with("## Identity") {
            in_identity = true;
            found_identity = true;
            result.push_str(new_identity);
            result.push('\n');
            continue;
        }
        if in_identity {
            if line.starts_with("## ") {
                in_identity = false;
            } else {
                continue;
            }
        }
        result.push_str(line);
        result.push('\n');
    }

    if !found_identity {
        if let Some(pos) = soul.find("## Core Truths") {
            let mut new_content = soul[..pos].to_string();
            new_content.push_str(new_identity);
            new_content.push_str("\n\n");
            new_content.push_str(&soul[pos..]);
            return new_content;
        }
        let mut new_content = soul.to_string();
        if !new_content.ends_with('\n') {
            new_content.push('\n');
        }
        new_content.push('\n');
        new_content.push_str(new_identity);
        return new_content;
    }

    result.trim_end().to_string()
}

fn parse_inbound_attachments(
    parsed: &serde_json::Value,
) -> Result<Vec<MediaAttachment>, &'static str> {
    let Some(raw_attachments) = parsed.get("attachments") else {
        return Ok(vec![]);
    };
    let Some(array) = raw_attachments.as_array() else {
        return Err("Invalid attachments payload");
    };

    let mut attachments = Vec::with_capacity(array.len());
    for item in array {
        let file_name = item
            .get("file_name")
            .and_then(serde_json::Value::as_str)
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .ok_or("Attachment file_name is required")?;
        let data_base64 = item
            .get("data_base64")
            .and_then(serde_json::Value::as_str)
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .ok_or("Attachment data_base64 is required")?;
        let data = STANDARD
            .decode(data_base64)
            .map_err(|_| "Attachment data_base64 is invalid")?;
        let mime_type = item
            .get("mime_type")
            .and_then(serde_json::Value::as_str)
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .map(ToOwned::to_owned);

        attachments.push(MediaAttachment {
            file_name: file_name.to_string(),
            data,
            mime_type,
        });
    }

    Ok(attachments)
}

/// Enrich the user message with textual descriptions of any attached media.
///
/// Uses the configured [`MediaPipeline`] to transcribe audio, describe images,
/// and summarise video. When the pipeline is disabled or there are no
/// attachments the original `content` is returned unchanged.
async fn enrich_with_attachments(
    state: &AppState,
    content: &str,
    attachments: Vec<MediaAttachment>,
) -> String {
    if attachments.is_empty() {
        return content.to_string();
    }
    let config = state.config.lock().clone();
    let pipeline = zeroclaw_channels::orchestrator::media_pipeline::MediaPipeline::new(
        &config.media_pipeline,
        &config.transcription,
        false,
    );
    pipeline.process(content, &attachments).await
}

async fn rebuild_agent_with_override(
    state: &AppState,
    agent: &mut zeroclaw_runtime::agent::Agent,
    session_id: &str,
    session_key: &str,
    provider_override: Option<String>,
    model_override: Option<String>,
    api_key_override: Option<String>,
) -> anyhow::Result<(String, String)> {
    let mut effective_config = state.config.lock().clone();
    if let Some(provider) = provider_override {
        effective_config.providers.fallback = Some(provider);
    }
    if let Some(model) = model_override {
        effective_config.ensure_fallback_provider().model = Some(model);
    }
    if let Some(api_key) = api_key_override {
        effective_config.ensure_fallback_provider().api_key = Some(api_key);
    }

    let provider_label = effective_config
        .providers
        .fallback
        .clone()
        .unwrap_or_else(|| "unknown".to_string());
    let model_label = effective_config
        .providers
        .fallback_provider()
        .and_then(|entry| entry.model.clone())
        .unwrap_or_else(|| state.model.clone());

    let mut next_agent = zeroclaw_runtime::agent::Agent::from_config(&effective_config).await?;
    next_agent.set_memory_session_id(Some(session_id.to_string()));

    if let Some(ref backend) = state.session_backend {
        let messages = backend.load(session_key);
        if !messages.is_empty() {
            next_agent.seed_history(&messages);
        }
    }

    *agent = next_agent;
    Ok((provider_label, model_label))
}

/// GET /ws/chat — WebSocket upgrade for agent chat
pub async fn handle_ws_chat(
    State(state): State<AppState>,
    Query(params): Query<WsQuery>,
    headers: HeaderMap,
    ws: WebSocketUpgrade,
) -> impl IntoResponse {
    // Auth: check header, subprotocol, then query param (precedence order)
    if state.pairing.require_pairing() {
        let token = extract_ws_token(&headers, params.token.as_deref()).unwrap_or("");
        if !state.pairing.is_authenticated(token) {
            return (
                axum::http::StatusCode::UNAUTHORIZED,
                "Unauthorized — provide Authorization header, Sec-WebSocket-Protocol bearer, or ?token= query param",
            )
                .into_response();
        }
    }

    // Echo Sec-WebSocket-Protocol if the client requests one of our supported sub-protocols.
    let client_protos: Vec<&str> = headers
        .get("sec-websocket-protocol")
        .and_then(|v| v.to_str().ok())
        .map(|s| s.split(',').map(|p| p.trim()).collect())
        .unwrap_or_default();
    let matched = client_protos
        .iter()
        .find_map(|cp| WS_PROTOCOLS.iter().find(|sp| **sp == *cp).copied());
    let ws = if let Some(proto) = matched {
        ws.protocols([proto])
    } else {
        ws
    };

    let session_id = params.session_id;
    let session_name = params.name;
    ws.on_upgrade(move |socket| handle_socket(socket, state, session_id, session_name))
        .into_response()
}

/// Gateway session key prefix to avoid collisions with channel sessions.
const GW_SESSION_PREFIX: &str = "gw_";

async fn handle_socket(
    socket: WebSocket,
    state: AppState,
    session_id: Option<String>,
    session_name: Option<String>,
) {
    let (mut sender, mut receiver) = socket.split();
    let default_provider_label = state
        .config
        .lock()
        .providers
        .fallback
        .clone()
        .unwrap_or_else(|| "unknown".to_string());
    let default_model_label = state.model.clone();
    let mut pending_model_override: Option<(Option<String>, Option<String>, Option<String>)> = None;

    // Resolve session ID: use provided or generate a new UUID
    let session_id = session_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
    let session_key = format!("{GW_SESSION_PREFIX}{session_id}");

    // Build a persistent Agent for this connection so history is maintained across turns.
    let config = state.config.lock().clone();
    let mut agent = match zeroclaw_runtime::agent::Agent::from_config(&config).await {
        Ok(a) => a,
        Err(e) => {
            tracing::error!(error = %e, "Agent initialization failed");
            let err = serde_json::json!({
                "type": "error",
                "message": format!("Failed to initialise agent: {e}"),
                "code": "AGENT_INIT_FAILED"
            });
            let _ = sender.send(Message::Text(err.to_string().into())).await;
            let _ = sender
                .send(Message::Close(Some(axum::extract::ws::CloseFrame {
                    code: 1011,
                    reason: axum::extract::ws::Utf8Bytes::from_static(
                        "Agent initialization failed",
                    ),
                })))
                .await;
            return;
        }
    };
    agent.set_memory_session_id(Some(session_id.clone()));

    // Hydrate agent from persisted session (if available)
    let mut resumed = false;
    let mut message_count: usize = 0;
    let mut effective_name: Option<String> = None;
    if let Some(ref backend) = state.session_backend {
        let messages = backend.load(&session_key);
        if !messages.is_empty() {
            message_count = messages.len();
            agent.seed_history(&messages);
            resumed = true;
        }
        // Set session name if provided (non-empty) on connect
        if let Some(ref name) = session_name
            && !name.is_empty()
        {
            let _ = backend.set_session_name(&session_key, name);
            effective_name = Some(name.clone());
        }
        // If no name was provided via query param, load the stored name
        if effective_name.is_none() {
            effective_name = backend.get_session_name(&session_key).unwrap_or(None);
        }
    }

    // Send session_start message to client
    let mut session_start = serde_json::json!({
        "type": "session_start",
        "session_id": session_id,
        "resumed": resumed,
        "message_count": message_count,
    });
    if let Some(ref name) = effective_name {
        session_start["name"] = serde_json::Value::String(name.clone());
    }
    let _ = sender
        .send(Message::Text(session_start.to_string().into()))
        .await;

    // ── Optional connect handshake ──────────────────────────────────
    // The first message may be a `{"type":"connect",...}` frame carrying
    // connection parameters.  If it is, we extract the params, send an
    // ack, and proceed to the normal message loop.  If the first message
    // is a regular `{"type":"message",...}` frame, we fall through and
    // process it immediately (backward-compatible).
    let mut first_msg_fallback: Option<String> = None;

    if let Some(first) = receiver.next().await {
        match first {
            Ok(Message::Text(text)) => {
                if let Ok(cp) = serde_json::from_str::<ConnectParams>(&text) {
                    if cp.msg_type == "connect" {
                        debug!(
                            session_id = ?cp.session_id,
                            device_name = ?cp.device_name,
                            capabilities = ?cp.capabilities,
                            "WebSocket connect params received"
                        );
                        // Override session_id if provided in connect params
                        if let Some(sid) = &cp.session_id {
                            agent.set_memory_session_id(Some(sid.clone()));
                        }
                        let ack = serde_json::json!({
                            "type": "connected",
                            "message": "Connection established"
                        });
                        let _ = sender.send(Message::Text(ack.to_string().into())).await;
                    } else {
                        // Not a connect message — fall through to normal processing
                        first_msg_fallback = Some(text.to_string());
                    }
                } else {
                    // Not parseable as ConnectParams — fall through
                    first_msg_fallback = Some(text.to_string());
                }
            }
            Ok(Message::Close(_)) | Err(_) => return,
            _ => {}
        }
    }

    // Process the first message if it was not a connect frame
    if let Some(ref text) = first_msg_fallback {
        match parse_ws_inbound_action(text) {
            WsInboundAction::Message {
                content,
                attachments,
            } => {
                let mut provider_label = default_provider_label.clone();
                let mut model_label = default_model_label.clone();
                let content = enrich_with_attachments(&state, &content, attachments).await;
                if let Some((provider_override, model_override, api_key_override)) =
                    pending_model_override.take()
                {
                    match rebuild_agent_with_override(
                        &state,
                        &mut agent,
                        &session_id,
                        &session_key,
                        provider_override,
                        model_override,
                        api_key_override,
                    )
                    .await
                    {
                        Ok((next_provider_label, next_model_label)) => {
                            provider_label = next_provider_label;
                            model_label = next_model_label;
                        }
                        Err(e) => {
                            let err = serde_json::json!({
                                "type": "error",
                                "message": format!("Failed to switch model: {e}"),
                                "code": "MODEL_SWITCH_FAILED"
                            });
                            let _ = sender.send(Message::Text(err.to_string().into())).await;
                            return;
                        }
                    }
                }

                if let Some(ref backend) = state.session_backend {
                    let user_msg = zeroclaw_providers::ChatMessage::user(&content);
                    let _ = backend.append(&session_key, &user_msg);
                }

                process_chat_message(
                    &state,
                    &mut agent,
                    &mut sender,
                    &mut receiver,
                    &content,
                    &session_key,
                    &provider_label,
                    &model_label,
                )
                .await;
            }
            WsInboundAction::Cancel => {
                let cancelled = serde_json::json!({
                    "type": "message_cancelled",
                    "session_id": session_id,
                });
                let _ = sender
                    .send(Message::Text(cancelled.to_string().into()))
                    .await;
            }
            WsInboundAction::SetModel {
                provider,
                model,
                api_key,
            } => {
                pending_model_override = Some((provider, model, api_key));
                let ack = serde_json::json!({ "type": "model_set", "session_id": session_id });
                let _ = sender.send(Message::Text(ack.to_string().into())).await;
            }
            WsInboundAction::SetIdentity { name, role } => {
                match persist_identity_update(&state, name, role).await {
                    Ok(()) => {
                        let ack =
                            serde_json::json!({ "type": "identity_set", "session_id": session_id });
                        let _ = sender.send(Message::Text(ack.to_string().into())).await;
                    }
                    Err(e) => {
                        let err = serde_json::json!({
                            "type": "error",
                            "message": format!("Failed to save identity: {e}"),
                            "code": "IDENTITY_SAVE_FAILED"
                        });
                        let _ = sender.send(Message::Text(err.to_string().into())).await;
                    }
                }
            }
            WsInboundAction::Ignore => {
                let err = serde_json::json!({
                    "type": "error",
                    "message": "Unsupported message type. Send {\"type\":\"message\",\"content\":\"your text\"}"
                });
                let _ = sender.send(Message::Text(err.to_string().into())).await;
            }
            WsInboundAction::Error(message) => {
                let err = serde_json::json!({ "type": "error", "message": message });
                let _ = sender.send(Message::Text(err.to_string().into())).await;
            }
        }
    }

    // Subscribe to the shared broadcast channel so cron/heartbeat events
    // are forwarded to this WebSocket client.
    let mut broadcast_rx = state.event_tx.subscribe();

    loop {
        tokio::select! {
            // ── Client message ────────────────────────────────────────
            client_msg = receiver.next() => {
                let Some(msg) = client_msg else { break };
                let msg = match msg {
                    Ok(Message::Text(text)) => text,
                    Ok(Message::Close(_)) | Err(_) => break,
                    _ => continue,
                };

                // Voice duplex text frames (gated by feature + runtime config).
                #[cfg(feature = "gateway-voice-duplex")]
                {
                    let duplex_enabled = state
                        .config
                        .lock()
                        .channels
                        .voice_duplex
                        .as_ref()
                        .is_some_and(|v| v.enabled);
                    if duplex_enabled {
                        if let Some(voice_event) = crate::voice_duplex::try_parse_voice_event(&msg)
                        {
                            if let Some(error_frame) =
                                crate::voice_duplex::handle_voice_event(voice_event)
                            {
                                let _ = sender
                                    .send(Message::Text(error_frame.to_string().into()))
                                    .await;
                            }
                            continue;
                        }
                    }
                }

                match parse_ws_inbound_action(&msg) {
                    WsInboundAction::Message { content, attachments } => {
                        let mut provider_label = default_provider_label.clone();
                        let mut model_label = default_model_label.clone();
                        let content = enrich_with_attachments(&state, &content, attachments).await;

                        let _session_guard = match state.session_queue.acquire(&session_key).await {
                            Ok(guard) => guard,
                            Err(e) => {
                                let err = serde_json::json!({
                                    "type": "error",
                                    "message": e.to_string(),
                                    "code": "SESSION_BUSY"
                                });
                                let _ = sender.send(Message::Text(err.to_string().into())).await;
                                continue;
                            }
                        };

                        if let Some((provider_override, model_override, api_key_override)) =
                            pending_model_override.take()
                        {
                            match rebuild_agent_with_override(
                                &state,
                                &mut agent,
                                &session_id,
                                &session_key,
                                provider_override,
                                model_override,
                                api_key_override,
                            )
                            .await
                            {
                                Ok((next_provider_label, next_model_label)) => {
                                    provider_label = next_provider_label;
                                    model_label = next_model_label;
                                }
                                Err(e) => {
                                    let err = serde_json::json!({
                                        "type": "error",
                                        "message": format!("Failed to switch model: {e}"),
                                        "code": "MODEL_SWITCH_FAILED"
                                    });
                                    let _ = sender.send(Message::Text(err.to_string().into())).await;
                                    continue;
                                }
                            }
                        }

                        if let Some(ref backend) = state.session_backend {
                            let user_msg = zeroclaw_providers::ChatMessage::user(&content);
                            let _ = backend.append(&session_key, &user_msg);
                        }

                        process_chat_message(
                            &state,
                            &mut agent,
                            &mut sender,
                            &mut receiver,
                            &content,
                            &session_key,
                            &provider_label,
                            &model_label,
                        )
                        .await;
                    }
                    WsInboundAction::Cancel => {
                        let cancelled = serde_json::json!({
                            "type": "message_cancelled",
                            "session_id": session_id
                        });
                        let _ = sender.send(Message::Text(cancelled.to_string().into())).await;
                    }
                    WsInboundAction::SetModel {
                        provider,
                        model,
                        api_key,
                    } => {
                        pending_model_override = Some((provider, model, api_key));
                        let ack = serde_json::json!({ "type": "model_set", "session_id": session_id });
                        let _ = sender.send(Message::Text(ack.to_string().into())).await;
                    }
                    WsInboundAction::SetIdentity { name, role } => {
                        match persist_identity_update(&state, name, role).await {
                            Ok(()) => {
                                let ack = serde_json::json!({
                                    "type": "identity_set",
                                    "session_id": session_id
                                });
                                let _ = sender.send(Message::Text(ack.to_string().into())).await;
                            }
                            Err(e) => {
                                let err = serde_json::json!({
                                    "type": "error",
                                    "message": format!("Failed to save identity: {e}"),
                                    "code": "IDENTITY_SAVE_FAILED"
                                });
                                let _ = sender.send(Message::Text(err.to_string().into())).await;
                            }
                        }
                    }
                    WsInboundAction::Ignore => {
                        let err = serde_json::json!({
                            "type": "error",
                            "message": "Unsupported message type. Send {\"type\":\"message\",\"content\":\"your text\"}",
                            "code": "UNKNOWN_MESSAGE_TYPE"
                        });
                        let _ = sender.send(Message::Text(err.to_string().into())).await;
                    }
                    WsInboundAction::Error(message) => {
                        let err = serde_json::json!({
                            "type": "error",
                            "message": message,
                            "code": if message == "Invalid JSON" { "INVALID_JSON" } else { "EMPTY_CONTENT" }
                        });
                        let _ = sender.send(Message::Text(err.to_string().into())).await;
                    }
                }
            }

            // ── Broadcast event (cron/heartbeat results) ──────────────
            event = broadcast_rx.recv() => {
                if let Ok(event) = event {
                    let _ = sender.send(Message::Text(event.to_string().into())).await;
                }
            }
        }
    }
}

async fn persist_identity_update(
    state: &AppState,
    name: Option<String>,
    role: Option<String>,
) -> anyhow::Result<()> {
    let mut config = state.config.lock().clone();
    config.identity.name = name.clone();
    config.identity.role = role.clone();
    let workspace_dir = config.workspace_dir.clone();
    config.save().await?;
    *state.config.lock() = config;

    if let Some(identity_markdown) = build_identity_markdown(name.as_deref(), role.as_deref()) {
        tokio::fs::write(workspace_dir.join("IDENTITY.md"), identity_markdown).await?;
    }
    if let Some(name) = name.as_deref() {
        let soul_path = workspace_dir.join("SOUL.md");
        if let Ok(soul_content) = tokio::fs::read_to_string(&soul_path).await {
            let next_identity = build_soul_identity_section(name, role.as_deref());
            let next_soul = replace_identity_section(&soul_content, &next_identity);
            tokio::fs::write(soul_path, next_soul).await?;
        }
    }

    Ok(())
}

/// Process a single chat message through the agent and send the response.
///
/// Uses [`Agent::turn_streamed`] so that intermediate text chunks, tool calls,
/// and tool results are forwarded to the WebSocket client in real time.
async fn process_chat_message(
    state: &AppState,
    agent: &mut zeroclaw_runtime::agent::Agent,
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    receiver: &mut futures_util::stream::SplitStream<WebSocket>,
    content: &str,
    session_key: &str,
    provider_label: &str,
    model_label: &str,
) {
    use zeroclaw_runtime::agent::TurnEvent;

    // Broadcast agent_start event
    let _ = state.event_tx.send(serde_json::json!({
        "type": "agent_start",
        "provider": provider_label,
        "model": model_label,
    }));

    // Set session state to running
    let turn_id = uuid::Uuid::new_v4().to_string();
    if let Some(ref backend) = state.session_backend {
        let _ = backend.set_session_state(session_key, "running", Some(&turn_id));
    }

    // ── Cancellation token lifecycle ─────────────────────────────
    // Create a token before the turn starts so the abort endpoint
    // can cancel it. Remove it after the turn completes regardless
    // of outcome (normal, error, or cancelled).
    let cancel_token = CancellationToken::new();
    {
        state
            .cancel_tokens
            .lock()
            .expect("cancel_tokens lock poisoned")
            .insert(session_key.to_string(), cancel_token.clone());
    }

    // Channel for streaming turn events from the agent.
    let (event_tx, mut event_rx) = tokio::sync::mpsc::channel::<TurnEvent>(64);

    let content_owned = content.to_string();
    let turn_fut = async {
        agent
            .turn_streamed(&content_owned, event_tx, Some(cancel_token.clone()))
            .await
    };
    tokio::pin!(turn_fut);

    let mut accumulated_text = String::new();
    let mut partial_saved = false;
    let mut last_partial_save_for_persist = std::time::Instant::now();
    let partial_save_interval = std::time::Duration::from_millis(500);

    let result = loop {
        tokio::select! {
            result = &mut turn_fut => break result,
            maybe_event = event_rx.recv() => {
                let Some(event) = maybe_event else {
                    continue;
                };

                let ws_msg = match event {
                    TurnEvent::Chunk { delta } => {
                        accumulated_text.push_str(&delta);

                        if last_partial_save_for_persist.elapsed() >= partial_save_interval {
                            if let Some(ref backend) = state.session_backend {
                                let partial =
                                    zeroclaw_providers::ChatMessage::assistant(&accumulated_text);
                                if partial_saved {
                                    let _ = backend.update_last(session_key, &partial);
                                } else {
                                    let _ = backend.append(session_key, &partial);
                                    partial_saved = true;
                                }
                            }
                            last_partial_save_for_persist = std::time::Instant::now();
                        }

                        serde_json::json!({ "type": "chunk", "content": delta })
                    }
                    TurnEvent::Thinking { delta } => {
                        serde_json::json!({ "type": "thinking", "content": delta })
                    }
                    TurnEvent::ToolCall { name, args } => {
                        serde_json::json!({ "type": "tool_call", "name": name, "args": args })
                    }
                    TurnEvent::ToolResult { name, output } => {
                        serde_json::json!({ "type": "tool_result", "name": name, "output": output })
                    }
                };
                let _ = sender.send(Message::Text(ws_msg.to_string().into())).await;
            }
            inbound = receiver.next() => {
                let Some(inbound) = inbound else {
                    cancel_token.cancel();
                    break Err(zeroclaw_runtime::agent::loop_::ToolLoopCancelled.into());
                };

                match inbound {
                    Ok(Message::Text(text)) => match parse_ws_inbound_action(&text) {
                        WsInboundAction::Cancel => {
                            cancel_token.cancel();
                        }
                        WsInboundAction::Ignore => {}
                        WsInboundAction::Error(message) => {
                            let err = serde_json::json!({
                                "type": "error",
                                "message": message,
                                "code": if message == "Invalid JSON" { "INVALID_JSON" } else { "EMPTY_CONTENT" }
                            });
                            let _ = sender.send(Message::Text(err.to_string().into())).await;
                        }
                        WsInboundAction::Message { .. }
                        | WsInboundAction::SetModel { .. }
                        | WsInboundAction::SetIdentity { .. } => {
                            let err = serde_json::json!({
                                "type": "error",
                                "message": "A turn is already in progress",
                                "code": "SESSION_BUSY"
                            });
                            let _ = sender.send(Message::Text(err.to_string().into())).await;
                        }
                    },
                    Ok(Message::Close(_)) | Err(_) => {
                        cancel_token.cancel();
                        break Err(zeroclaw_runtime::agent::loop_::ToolLoopCancelled.into());
                    }
                    _ => {}
                }
            }
        }
    };

    // ── Remove cancel token (turn finished) ──────────────────────
    {
        state
            .cancel_tokens
            .lock()
            .expect("cancel_tokens lock poisoned")
            .remove(session_key);
    }

    let was_cancelled = match &result {
        Err(e) => zeroclaw_runtime::agent::loop_::is_tool_loop_cancelled(e),
        Ok(_) => false,
    };

    if was_cancelled {
        let truncated = if accumulated_text.is_empty() {
            "[interrupted by user]".to_string()
        } else {
            format!("{accumulated_text}\n\n[interrupted by user]")
        };

        if let Some(ref backend) = state.session_backend {
            let assistant_msg = zeroclaw_providers::ChatMessage::assistant(&truncated);
            if partial_saved {
                let _ = backend.update_last(session_key, &assistant_msg);
            } else {
                let _ = backend.append(session_key, &assistant_msg);
            }
        }

        let aborted = serde_json::json!({ "type": "aborted" });
        let _ = sender.send(Message::Text(aborted.to_string().into())).await;

        if let Some(ref backend) = state.session_backend {
            let _ = backend.set_session_state(session_key, "idle", None);
        }

        let _ = state.event_tx.send(serde_json::json!({
            "type": "agent_end",
            "provider": provider_label,
            "model": model_label,
        }));

        return;
    }

    match result {
        Ok(response) => {
            // Persist final assistant response. If we saved partial content
            // during streaming, update it in-place; otherwise append fresh.
            if let Some(ref backend) = state.session_backend {
                let assistant_msg = zeroclaw_providers::ChatMessage::assistant(&response);
                if partial_saved {
                    let _ = backend.update_last(session_key, &assistant_msg);
                } else {
                    let _ = backend.append(session_key, &assistant_msg);
                }
            }

            // Fire-and-forget memory consolidation so facts from WS sessions
            // are extracted to long-term memory (Daily + Core categories).
            if state.auto_save {
                let mem = state.mem.clone();
                let provider = state.provider.clone();
                let model = model_label.to_string();
                let user_msg = content.to_string();
                let assistant_resp = response.clone();
                tokio::spawn(async move {
                    if let Err(e) = zeroclaw_memory::consolidation::consolidate_turn(
                        provider.as_ref(),
                        &model,
                        mem.as_ref(),
                        &user_msg,
                        &assistant_resp,
                    )
                    .await
                    {
                        tracing::debug!("WS memory consolidation skipped: {e}");
                    }
                });
            }

            // Send chunk_reset so the client clears any accumulated draft
            // before the authoritative done message.
            let reset = serde_json::json!({ "type": "chunk_reset" });
            let _ = sender.send(Message::Text(reset.to_string().into())).await;

            let done = serde_json::json!({
                "type": "done",
                "full_response": response,
            });
            let _ = sender.send(Message::Text(done.to_string().into())).await;

            // Set session state to idle
            if let Some(ref backend) = state.session_backend {
                let _ = backend.set_session_state(session_key, "idle", None);
            }

            // Broadcast agent_end event
            let _ = state.event_tx.send(serde_json::json!({
                "type": "agent_end",
                "provider": provider_label,
                "model": model_label,
            }));
        }
        Err(e) => {
            if zeroclaw_runtime::agent::loop_::is_tool_loop_cancelled(&e) {
                if let Some(ref backend) = state.session_backend {
                    let _ = backend.set_session_state(session_key, "idle", None);
                }
                return;
            }

            // Set session state to error
            if let Some(ref backend) = state.session_backend {
                let _ = backend.set_session_state(session_key, "error", Some(&turn_id));
            }

            tracing::error!(error = %e, "Agent turn failed");
            let sanitized = zeroclaw_providers::sanitize_api_error(&e.to_string());
            let error_code = if sanitized.to_lowercase().contains("api key")
                || sanitized.to_lowercase().contains("authentication")
                || sanitized.to_lowercase().contains("unauthorized")
            {
                "AUTH_ERROR"
            } else if sanitized.to_lowercase().contains("provider")
                || sanitized.to_lowercase().contains("model")
            {
                "PROVIDER_ERROR"
            } else {
                "AGENT_ERROR"
            };
            let err = serde_json::json!({
                "type": "error",
                "message": sanitized,
                "code": error_code,
            });
            let _ = sender.send(Message::Text(err.to_string().into())).await;

            // Broadcast error event
            let _ = state.event_tx.send(serde_json::json!({
                "type": "error",
                "component": "ws_chat",
                "message": sanitized,
            }));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::HeaderMap;

    #[test]
    fn extract_ws_token_from_authorization_header() {
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer zc_test123".parse().unwrap());
        assert_eq!(extract_ws_token(&headers, None), Some("zc_test123"));
    }

    #[test]
    fn extract_ws_token_from_subprotocol() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "sec-websocket-protocol",
            "zeroclaw.v1, bearer.zc_sub456".parse().unwrap(),
        );
        assert_eq!(extract_ws_token(&headers, None), Some("zc_sub456"));
    }

    #[test]
    fn extract_ws_token_from_query_param() {
        let headers = HeaderMap::new();
        assert_eq!(
            extract_ws_token(&headers, Some("zc_query789")),
            Some("zc_query789")
        );
    }

    #[test]
    fn extract_ws_token_precedence_header_over_subprotocol() {
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer zc_header".parse().unwrap());
        headers.insert("sec-websocket-protocol", "bearer.zc_sub".parse().unwrap());
        assert_eq!(
            extract_ws_token(&headers, Some("zc_query")),
            Some("zc_header")
        );
    }

    #[test]
    fn extract_ws_token_precedence_subprotocol_over_query() {
        let mut headers = HeaderMap::new();
        headers.insert("sec-websocket-protocol", "bearer.zc_sub".parse().unwrap());
        assert_eq!(extract_ws_token(&headers, Some("zc_query")), Some("zc_sub"));
    }

    #[test]
    fn extract_ws_token_returns_none_when_empty() {
        let headers = HeaderMap::new();
        assert_eq!(extract_ws_token(&headers, None), None);
    }

    #[test]
    fn extract_ws_token_skips_empty_header_value() {
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer ".parse().unwrap());
        assert_eq!(
            extract_ws_token(&headers, Some("zc_fallback")),
            Some("zc_fallback")
        );
    }

    #[test]
    fn extract_ws_token_skips_empty_query_param() {
        let headers = HeaderMap::new();
        assert_eq!(extract_ws_token(&headers, Some("")), None);
    }

    #[test]
    fn extract_ws_token_subprotocol_with_multiple_entries() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "sec-websocket-protocol",
            "zeroclaw.v1, bearer.zc_tok, other".parse().unwrap(),
        );
        assert_eq!(extract_ws_token(&headers, None), Some("zc_tok"));
    }

    #[test]
    fn parse_ws_inbound_action_parses_set_model_with_both_fields() {
        let action = parse_ws_inbound_action(
            r#"{"type":"set_model","provider":"kylin-proxy","model":"qwen-plus"}"#,
        );
        assert_eq!(
            action,
            WsInboundAction::SetModel {
                provider: Some("kylin-proxy".to_string()),
                model: Some("qwen-plus".to_string()),
                api_key: None,
            }
        );
    }

    #[test]
    fn parse_ws_inbound_action_parses_set_model_reset() {
        assert_eq!(
            parse_ws_inbound_action(r#"{"type":"set_model"}"#),
            WsInboundAction::SetModel {
                provider: None,
                model: None,
                api_key: None,
            }
        );
    }

    #[test]
    fn parse_ws_inbound_action_parses_cancel() {
        assert_eq!(
            parse_ws_inbound_action(r#"{"type":"cancel"}"#),
            WsInboundAction::Cancel
        );
    }

    #[test]
    fn parse_ws_inbound_action_parses_message_attachments() {
        let action = parse_ws_inbound_action(
            r#"{"type":"message","content":"","attachments":[{"file_name":"a.png","data_base64":"YWI=","mime_type":"image/png"}]}"#,
        );
        match action {
            WsInboundAction::Message {
                content,
                attachments,
            } => {
                assert!(content.is_empty());
                assert_eq!(attachments.len(), 1);
                assert_eq!(attachments[0].file_name, "a.png");
                assert_eq!(attachments[0].data, b"ab");
            }
            _ => panic!("expected message action"),
        }
    }
}
