//! WebSocket workflow for live chat delivery and AI responses.
//!
//! This module does two jobs:
//! - fan out database-backed room events to subscribed browser sockets
//! - turn incoming user chat messages into stored messages and optional AI replies

use axum::{
    extract::{
        ws::{Message, WebSocket},
        Query, State, WebSocketUpgrade,
    },
    response::IntoResponse,
};
use futures::{SinkExt, StreamExt};
use std::sync::Arc;
use uuid::Uuid;

use crate::{
    middleware::auth::decode_token,
    models::{BroadcastEvent, MessagePayload, WsClientMessage, WsServerMessage},
    services::{fetch, openrouter, search},
    AppState,
};

/// Maximum number of tool call rounds before forcing a text response.
const MAX_TOOL_ROUNDS: usize = 5;

#[derive(serde::Deserialize)]
pub struct WsQuery {
    token: String,
}

/// Upgrade an authenticated request into a WebSocket connection.
pub async fn ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Query(query): Query<WsQuery>,
) -> impl IntoResponse {
    // Authenticate before upgrading
    let claims = match decode_token(&query.token, &state.jwt_secret) {
        Ok(c) => c,
        Err(_) => {
            return axum::http::StatusCode::UNAUTHORIZED.into_response();
        }
    };

    ws.on_upgrade(move |socket| {
        handle_socket(socket, state, claims.sub, claims.display_name, claims.email)
    })
}

/// Drive a single WebSocket connection until either the send or receive side ends.
async fn handle_socket(
    socket: WebSocket,
    state: Arc<AppState>,
    user_id: String,
    display_name: String,
    email: String,
) {
    let (mut ws_tx, mut ws_rx) = socket.split();
    let mut broadcast_rx = state.tx.subscribe();

    // Track which rooms this user is watching
    let subscribed_rooms: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>> =
        Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new()));

    let rooms_clone = subscribed_rooms.clone();

    // Task 1: forward room events from the shared broadcast channel into this
    // specific socket, but only for rooms the browser subscribed to.
    let mut send_task = tokio::spawn(async move {
        loop {
            match broadcast_rx.recv().await {
                Ok(event) => {
                    let rooms = rooms_clone.lock().await;
                    if rooms.contains(&event.room_id) {
                        let msg = serde_json::to_string(&event.message).unwrap();
                        if ws_tx.send(Message::Text(msg.into())).await.is_err() {
                            break;
                        }
                    }
                }
                Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                    tracing::warn!("WS subscriber lagged, skipped {} messages", n);
                    // Continue receiving — don't drop the connection
                    continue;
                }
                Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                    break;
                }
            }
        }
    });

    let state_clone = state.clone();
    let user_id_clone = user_id.clone();
    let display_name_clone = display_name.clone();
    let email_clone = email.clone();
    let rooms_clone2 = subscribed_rooms.clone();

    // Task 2: receive commands from the browser and translate them into
    // database writes, broadcasts, or AI work.
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = ws_rx.next().await {
            let text = match msg {
                Message::Text(t) => t.to_string(),
                Message::Close(_) => break,
                _ => continue,
            };

            let client_msg: WsClientMessage = match serde_json::from_str(&text) {
                Ok(m) => m,
                Err(e) => {
                    let _ = state_clone.tx.send(BroadcastEvent {
                        room_id: String::new(),
                        message: WsServerMessage::Error {
                            message: format!("Invalid message: {}", e),
                        },
                    });
                    continue;
                }
            };

            match client_msg {
                WsClientMessage::JoinRoom { room_id } => {
                    tracing::info!("User {} joined room {}", user_id_clone, room_id);
                    rooms_clone2.lock().await.insert(room_id.clone());
                }
                WsClientMessage::Typing { room_id } => {
                    let _ = state_clone.tx.send(BroadcastEvent {
                        room_id: room_id.clone(),
                        message: WsServerMessage::UserTyping {
                            room_id,
                            user_id: user_id_clone.clone(),
                            display_name: display_name_clone.clone(),
                        },
                    });
                }
                WsClientMessage::SendMessage {
                    room_id,
                    content,
                    mentions,
                    image_url,
                } => {
                    tracing::info!("User {} sending message to room {}", user_id_clone, room_id);
                    handle_send_message(
                        &state_clone,
                        &user_id_clone,
                        &display_name_clone,
                        &email_clone,
                        &room_id,
                        &content,
                        &mentions,
                        image_url.as_deref(),
                    )
                    .await;
                }
            }
        }
    });

    // If either half of the connection ends, stop the companion task too.
    tokio::select! {
        _ = &mut send_task => recv_task.abort(),
        _ = &mut recv_task => send_task.abort(),
    }

    tracing::info!("WebSocket disconnected: {}", user_id);
}

/// Persist a user message, broadcast it, and optionally generate an AI reply.
async fn handle_send_message(
    state: &Arc<AppState>,
    user_id: &str,
    display_name: &str,
    email: &str,
    room_id: &str,
    content: &str,
    mentions: &[String],
    image_url: Option<&str>,
) {
    let msg_id = Uuid::new_v4().to_string();
    let mentions_json = serde_json::to_string(mentions).unwrap_or_else(|_| "[]".to_string());
    let now = chrono::Utc::now().to_rfc3339();

    // Compute integrity hash from timestamp + content
    let hash = crate::models::message_hash(&now, content);

    // Store in database (with image_url)
    let _ = sqlx::query(
        "INSERT INTO messages (id, room_id, sender_id, sender_name, content, mentions, is_ai, created_at, image_url, hash) VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?, ?)",
    )
    .bind(&msg_id)
    .bind(room_id)
    .bind(user_id)
    .bind(display_name)
    .bind(content)
    .bind(&mentions_json)
    .bind(&now)
    .bind(image_url)
    .bind(&hash)
    .execute(&state.db)
    .await;

    // Look up the sender's custom avatar (if any) for the message payload
    let avatar_url: Option<String> =
        sqlx::query_scalar("SELECT avatar_url FROM users WHERE id = ?")
            .bind(user_id)
            .fetch_optional(&state.db)
            .await
            .ok()
            .flatten()
            .flatten();

    // Broadcast human message
    let payload = MessagePayload {
        id: msg_id,
        room_id: room_id.to_string(),
        sender_id: user_id.to_string(),
        sender_name: display_name.to_string(),
        content: content.to_string(),
        mentions: mentions.to_vec(),
        is_ai: false,
        created_at: now,
        ai_meta: None,
        avatar_hash: crate::models::gravatar_hash(email),
        avatar_url,
        image_url: image_url.map(String::from),
        hash: Some(hash),
    };

    let _ = state.tx.send(BroadcastEvent {
        room_id: room_id.to_string(),
        message: WsServerMessage::NewMessage { message: payload },
    });

    // The AI only replies when explicitly mentioned or when the room is set to
    // auto-reply to every message.
    let ai_user_id = "ai-assistant";
    let should_respond = mentions.contains(&ai_user_id.to_string());

    // Also check room settings for ai_always_respond
    let room = sqlx::query_as::<_, (String, bool, String, String)>(
        "SELECT model_id, ai_always_respond, system_prompt, ai_name FROM rooms WHERE id = ? AND deleted_at IS NULL",
    )
    .bind(room_id)
    .fetch_optional(&state.db)
    .await;

    let (model_id, always_respond, system_prompt, ai_name) = match room {
        Ok(Some(r)) => r,
        _ => return,
    };

    if !should_respond && !always_respond {
        return;
    }

    // Signal AI is typing
    let _ = state.tx.send(BroadcastEvent {
        room_id: room_id.to_string(),
        message: WsServerMessage::AiTyping {
            room_id: room_id.to_string(),
        },
    });

    // Fetch recent history (including image_url for multimodal support)
    let recent_messages = sqlx::query_as::<_, (String, String, bool, Option<String>)>(
        "SELECT sender_name, content, is_ai, image_url FROM messages WHERE room_id = ? ORDER BY created_at DESC LIMIT 50",
    )
    .bind(room_id)
    .fetch_all(&state.db)
    .await
    .unwrap_or_default();

    // OpenRouter accepts image inputs as data URLs, so local uploads need to be
    // loaded from disk and encoded before they are sent upstream.
    let mut history: Vec<(String, String, bool, Option<String>)> = Vec::new();
    for (sender_name, msg_content, is_ai, msg_image_url) in recent_messages.into_iter().rev() {
        let image_data_url = match &msg_image_url {
            Some(url) if !is_ai => encode_image_as_data_url(url).await,
            _ => None,
        };
        history.push((sender_name, msg_content, is_ai, image_data_url));
    }

    let mut chat_history = openrouter::build_chat_history(&system_prompt, &history);

    // Build tools for AI
    let tools = openrouter::build_tools();

    // Pre-generate AI message ID so we can reference it in stream chunks
    let ai_msg_id = Uuid::new_v4().to_string();

    // Run the AI in a loop because the model may first request tools, then need
    // follow-up rounds after those tool results are added to history.
    let mut total_prompt_tokens: u32 = 0;
    let mut total_completion_tokens: u32 = 0;
    let mut total_response_ms: u64 = 0;
    let mut final_model = model_id.clone();
    let mut ai_response = String::new();
    let mut had_error = false;
    let mut collected_tool_results: Vec<crate::models::ToolResult> = vec![];

    'tool_loop: for round in 0..MAX_TOOL_ROUNDS {
        let mut stream_rx = openrouter::chat_completion_stream(
            chat_history.clone(),
            &model_id,
            &state.openrouter_key,
            Some(tools.clone()),
        )
        .await;

        while let Some(event) = stream_rx.recv().await {
            match event {
                openrouter::StreamEvent::Delta(text) => {
                    // Broadcast each content chunk to clients
                    let _ = state.tx.send(BroadcastEvent {
                        room_id: room_id.to_string(),
                        message: WsServerMessage::AiStreamChunk {
                            room_id: room_id.to_string(),
                            message_id: ai_msg_id.clone(),
                            delta: text.clone(),
                        },
                    });
                    ai_response.push_str(&text);
                }
                openrouter::StreamEvent::ToolCalls(assistant_msg, stats) => {
                    total_prompt_tokens += stats.prompt_tokens;
                    total_completion_tokens += stats.completion_tokens;
                    total_response_ms += stats.response_ms;
                    final_model = stats.model.clone();

                    tracing::info!(
                        "AI requesting tool calls (round {}): {:?}",
                        round + 1,
                        assistant_msg
                            .tool_calls
                            .as_ref()
                            .map(|tc| tc.iter().map(|t| &t.function.name).collect::<Vec<_>>())
                    );

                    // Preserve the assistant tool-call message so the next round
                    // has the same context the model produced.
                    let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default();
                    chat_history.push(assistant_msg);

                    // Tool results are fed back into the conversation as
                    // synthetic `tool` messages, matching the upstream API.
                    for tool_call in &tool_calls {
                        let tool_input = extract_tool_input(
                            &tool_call.function.name,
                            &tool_call.function.arguments,
                        );

                        // Broadcast real-time tool usage event
                        let _ = state.tx.send(BroadcastEvent {
                            room_id: room_id.to_string(),
                            message: WsServerMessage::AiToolUsage {
                                room_id: room_id.to_string(),
                                tool_name: tool_call.function.name.clone(),
                                status: "calling".to_string(),
                            },
                        });

                        let tool_result = execute_tool(
                            &tool_call.function.name,
                            &tool_call.function.arguments,
                            state.search_provider,
                            state.tavily_api_key.as_deref(),
                            state.brave_api_key.as_deref(),
                        )
                        .await;

                        tracing::info!(
                            "Tool {} result: {} chars",
                            tool_call.function.name,
                            tool_result.len()
                        );

                        collected_tool_results.push(crate::models::ToolResult {
                            tool: tool_call.function.name.clone(),
                            input: tool_input,
                            result: tool_result.clone(),
                        });

                        chat_history.push(openrouter::ChatMessage {
                            role: "tool".into(),
                            content: Some(openrouter::Content::Text(tool_result)),
                            tool_calls: None,
                            tool_call_id: Some(tool_call.id.clone()),
                        });
                    }
                    // Ask the model to continue now that tool output exists.
                    continue 'tool_loop;
                }
                openrouter::StreamEvent::Done(stats) => {
                    total_prompt_tokens += stats.prompt_tokens;
                    total_completion_tokens += stats.completion_tokens;
                    total_response_ms += stats.response_ms;
                    final_model = stats.model;
                    break 'tool_loop;
                }
                openrouter::StreamEvent::Error(e) => {
                    tracing::error!("OpenRouter stream error (round {}): {}", round + 1, e);
                    ai_response = format!("*Sorry, I encountered an error: {}*", e);
                    had_error = true;
                    break 'tool_loop;
                }
            }
        }
    }

    // Guardrail: if the model never produced final prose, store a clear fallback
    // instead of leaving the client waiting indefinitely.
    if ai_response.is_empty() && !had_error {
        ai_response =
            "*I used several tools but couldn't formulate a final response. Please try again.*"
                .to_string();
    }

    // Signal stream end so client can finalize rendering
    let _ = state.tx.send(BroadcastEvent {
        room_id: room_id.to_string(),
        message: WsServerMessage::AiStreamEnd {
            room_id: room_id.to_string(),
            message_id: ai_msg_id.clone(),
        },
    });

    let ai_meta = if !had_error {
        Some(crate::models::AiMeta {
            model: final_model,
            prompt_tokens: total_prompt_tokens,
            completion_tokens: total_completion_tokens,
            total_tokens: total_prompt_tokens + total_completion_tokens,
            response_ms: total_response_ms,
            tool_results: if collected_tool_results.is_empty() {
                None
            } else {
                Some(collected_tool_results)
            },
        })
    } else {
        None
    };

    // Store AI response
    let ai_now = chrono::Utc::now().to_rfc3339();

    // Serialize ai_meta for database storage
    let ai_meta_json = ai_meta.as_ref().and_then(|m| serde_json::to_string(m).ok());

    // Compute integrity hash from timestamp + content
    let ai_hash = crate::models::message_hash(&ai_now, &ai_response);

    let _ = sqlx::query(
        "INSERT INTO messages (id, room_id, sender_id, sender_name, content, mentions, is_ai, created_at, ai_meta, hash) VALUES (?, ?, ?, ?, ?, '[]', 1, ?, ?, ?)",
    )
    .bind(&ai_msg_id)
    .bind(room_id)
    .bind(ai_user_id)
    .bind(&ai_name)
    .bind(&ai_response)
    .bind(&ai_now)
    .bind(&ai_meta_json)
    .bind(&ai_hash)
    .execute(&state.db)
    .await;

    // Broadcast final AI message (includes full content + ai_meta)
    let ai_payload = MessagePayload {
        id: ai_msg_id,
        room_id: room_id.to_string(),
        sender_id: ai_user_id.to_string(),
        sender_name: ai_name.clone(),
        content: ai_response,
        mentions: vec![],
        is_ai: true,
        created_at: ai_now,
        ai_meta,
        avatar_hash: String::new(),
        avatar_url: None,
        image_url: None,
        hash: Some(ai_hash),
    };

    let _ = state.tx.send(BroadcastEvent {
        room_id: room_id.to_string(),
        message: WsServerMessage::NewMessage {
            message: ai_payload,
        },
    });
}

/// Read a local image file and encode it as a base64 data URL for the OpenRouter API.
async fn encode_image_as_data_url(url: &str) -> Option<String> {
    // url is like "/uploads/chat-images/uuid.png"
    let file_path = url.strip_prefix('/').unwrap_or(url);
    let data = tokio::fs::read(file_path).await.ok()?;

    let mime = if file_path.ends_with(".png") {
        "image/png"
    } else if file_path.ends_with(".jpg") || file_path.ends_with(".jpeg") {
        "image/jpeg"
    } else if file_path.ends_with(".gif") {
        "image/gif"
    } else if file_path.ends_with(".webp") {
        "image/webp"
    } else {
        "application/octet-stream"
    };

    use base64::Engine;
    let b64 = base64::engine::general_purpose::STANDARD.encode(&data);
    Some(format!("data:{};base64,{}", mime, b64))
}

/// Extract a human-readable input string from tool arguments (for UI display).
fn extract_tool_input(tool_name: &str, arguments: &str) -> String {
    let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
    match tool_name {
        "web_search" | "brave_search" => args["query"].as_str().unwrap_or("").to_string(),
        "web_fetch" => args["url"].as_str().unwrap_or("").to_string(),
        _ => arguments.to_string(),
    }
}

/// Execute a tool call by name, returning the result as a string.
async fn execute_tool(
    name: &str,
    arguments: &str,
    search_provider: search::SearchProvider,
    tavily_api_key: Option<&str>,
    brave_api_key: Option<&str>,
) -> String {
    match name {
        "web_search" | "brave_search" => {
            let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
            let query = args["query"].as_str().unwrap_or("").to_string();
            let count = args["count"].as_u64().unwrap_or(5) as u8;

            if query.is_empty() {
                return "Error: search query is required".into();
            }

            match search::search(
                search_provider,
                &query,
                tavily_api_key,
                brave_api_key,
                count,
            )
            .await
            {
                Ok(results) => search::format_results(&results),
                Err(e) => format!("Search error: {}", e),
            }
        }
        "web_fetch" => {
            let args: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
            let url = args["url"].as_str().unwrap_or("").to_string();

            if url.is_empty() {
                return "Error: URL is required".into();
            }

            match fetch::fetch_url(&url, 8000).await {
                Ok(result) => fetch::format_result(&result),
                Err(e) => format!("Fetch error: {}", e),
            }
        }
        _ => format!("Unknown tool: {}", name),
    }
}
