use actix_web::{Error, HttpResponse, http::header, web};
use async_openai::{
    Client,
    config::OpenAIConfig,
    types::chat::{
        ChatCompletionMessageToolCall, ChatCompletionMessageToolCalls,
        ChatCompletionRequestAssistantMessage, ChatCompletionRequestMessage,
        ChatCompletionRequestSystemMessage, ChatCompletionRequestToolMessage,
        ChatCompletionRequestUserMessage, ChatCompletionStreamOptions,
        CreateChatCompletionRequestArgs, FinishReason,
    },
};
use futures::stream::{Stream, StreamExt};
use log::{debug, warn};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, broadcast, oneshot};
use tokio_stream::wrappers::BroadcastStream;

use crate::{config, llm, logger, session, template, tokens, tools};

// Tool approval state management
#[derive(Debug)]
pub struct ApprovalState {
    pub tool_name: String,
    pub tool_args: Value,
    pub tool_call_id: String,
    pub agent_id: String,
    pub sender: oneshot::Sender<bool>,
    pub created_at: Instant,
}

pub type ApprovalStateMap = Arc<Mutex<HashMap<String, ApprovalState>>>;

// Session update SSE broadcaster
static SESSION_UPDATE_BROADCASTER: OnceLock<broadcast::Sender<SessionUpdateEvent>> =
    OnceLock::new();

#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum SessionUpdateEvent {
    #[serde(rename = "update")]
    Update { session: SessionListItem },
    #[serde(rename = "deleted")]
    Deleted { session_id: String },
}

/// Initialize the session update broadcaster (called at server startup)
pub fn init_session_broadcaster() {
    let (tx, _rx) = broadcast::channel(100);
    let _ = SESSION_UPDATE_BROADCASTER.set(tx);
}

/// Broadcast a session update event to all SSE listeners
pub fn broadcast_session_update(event: SessionUpdateEvent) {
    if let Some(tx) = SESSION_UPDATE_BROADCASTER.get() {
        // Ignore send errors (no active listeners)
        let _ = tx.send(event);
    }
}

/// Get a human-readable description for a tool
fn get_tool_description(tool_name: &str) -> String {
    match tool_name {
        "read_file" => "Read the contents of a file from the filesystem".to_string(),
        "write_file" => "Write content to a file on the filesystem".to_string(),
        "grep" => "Search for a pattern in files using regex".to_string(),
        "bash" => "Execute a bash command (safe, read-only commands only)".to_string(),
        "demo_tool" => {
            "A demo tool for testing the approval workflow (safe, read-only)".to_string()
        }
        _ => format!("Execute tool: {}", tool_name),
    }
}

// ========================================
// Helper Functions
// ========================================

/// Build a SessionListItem from a session
fn build_session_list_item(session: &session::ChatSession) -> SessionListItem {
    let preview = session
        .messages
        .iter()
        .find(|msg| msg.role == "user")
        .map(|msg| {
            let content = &msg.content;
            if content.len() > 100 {
                format!("{}...", &content[..100])
            } else {
                content.clone()
            }
        });

    SessionListItem {
        session_id: session.id.clone(),
        message_count: session.messages.len(),
        created_at: session.created_at,
        updated_at: session.updated_at,
        preview,
        title: session.title.clone(),
        agent_id: session.agent_id.clone(),
        token_usage: TokenUsageResponse {
            total_tokens: session.token_usage.total_tokens,
            input_tokens: session.token_usage.input_tokens,
            output_tokens: session.token_usage.output_tokens,
            reasoning_tokens: session.token_usage.reasoning_tokens,
            cache_tokens: session.token_usage.cache_tokens,
            context_window: session.token_usage.context_window,
            context_utilization: session.token_usage.context_utilization,
        },
        cost_usd: session.cost_usd,
        is_readonly: session.is_readonly,
    }
}

/// Broadcast a session update with the given session
fn broadcast_session_update_for_session(
    session_manager: &session::SessionManager,
    session_id: &str,
) {
    if let Some(session) = session_manager.get_session(session_id) {
        let session_item = build_session_list_item(&session);
        broadcast_session_update(SessionUpdateEvent::Update {
            session: session_item,
        });
    }
}

/// Remove <think>...</think> and <tool_call>...</tool_call> tags from content
fn sanitize_assistant_content(content: &str) -> String {
    let mut result = content.to_string();

    // Remove all <think>...</think> tags
    while let Some(think_start) = result.find("<think>") {
        if let Some(think_end) = result.find("</think>") {
            if think_end > think_start {
                result = format!("{}{}", &result[..think_start], &result[think_end + 8..]);
            } else {
                break;
            }
        } else {
            break;
        }
    }

    // Remove all <tool_call>...</tool_call> tags
    while let Some(tool_start) = result.find("<tool_call>") {
        if let Some(tool_end) = result.find("</tool_call>") {
            if tool_end > tool_start {
                result = format!("{}{}", &result[..tool_start], &result[tool_end + 12..]);
            } else {
                break;
            }
        } else {
            break;
        }
    }

    result
}

/// Reconstruct chat messages from session history (excluding last message)
fn build_messages_from_history(
    session: &session::ChatSession,
    system_message: String,
) -> Vec<ChatCompletionRequestMessage> {
    let mut messages: Vec<ChatCompletionRequestMessage> = vec![
        ChatCompletionRequestSystemMessage {
            content: system_message.into(),
            ..Default::default()
        }
        .into(),
    ];

    // Add conversation history (excluding the last user message)
    for (i, msg) in session.messages.iter().enumerate() {
        if i == session.messages.len() - 1 {
            break; // Skip the last message as it will be added with full context
        }

        if msg.role == "user" {
            messages.push(
                ChatCompletionRequestUserMessage {
                    content: msg.content.clone().into(),
                    ..Default::default()
                }
                .into(),
            );
        } else if msg.role == "assistant" {
            // Check if this message has tool invocations in thinking steps
            if let Some(thinking_steps) = &msg.thinking_steps {
                let tool_steps: Vec<_> = thinking_steps
                    .iter()
                    .filter(|s| s.step_type == "tool")
                    .collect();

                if !tool_steps.is_empty() {
                    // Reconstruct tool calls from thinking steps
                    let assistant_tool_calls: Vec<ChatCompletionMessageToolCalls> = tool_steps
                        .iter()
                        .enumerate()
                        .map(|(idx, step)| {
                            let mut tool_call = ChatCompletionMessageToolCall {
                                id: format!("call_{}", idx),
                                function: Default::default(),
                            };
                            tool_call.function.name = step.tool_name.clone().unwrap_or_default();
                            tool_call.function.arguments =
                                serde_json::to_string(&step.tool_arguments).unwrap_or_default();
                            ChatCompletionMessageToolCalls::Function(tool_call)
                        })
                        .collect();

                    // Add assistant message with tool calls
                    let filtered_content = llm::strip_reasoning_blocks(&msg.content);
                    messages.push(
                        ChatCompletionRequestAssistantMessage {
                            content: if filtered_content.is_empty() {
                                None
                            } else {
                                Some(filtered_content.into())
                            },
                            tool_calls: Some(assistant_tool_calls),
                            ..Default::default()
                        }
                        .into(),
                    );

                    // Add tool result messages
                    for (idx, step) in tool_steps.iter().enumerate() {
                        let result_content = if let Some(error) = &step.tool_error {
                            json!({"error": error}).to_string()
                        } else if let Some(result) = &step.tool_result {
                            result.clone()
                        } else {
                            json!({"message": "Tool executed"}).to_string()
                        };

                        messages.push(
                            ChatCompletionRequestToolMessage {
                                content: result_content.into(),
                                tool_call_id: format!("call_{}", idx),
                            }
                            .into(),
                        );
                    }
                } else {
                    // No tools, just add normal assistant message
                    let filtered_content = llm::strip_reasoning_blocks(&msg.content);
                    messages.push(
                        ChatCompletionRequestAssistantMessage {
                            content: Some(filtered_content.into()),
                            ..Default::default()
                        }
                        .into(),
                    );
                }
            } else {
                // No thinking steps, just add normal assistant message
                let filtered_content = llm::strip_reasoning_blocks(&msg.content);
                messages.push(
                    ChatCompletionRequestAssistantMessage {
                        content: Some(filtered_content.into()),
                        ..Default::default()
                    }
                    .into(),
                );
            }
        }
    }

    messages
}

/// Estimate token usage for a session when provider doesn't report it
async fn estimate_and_send_usage(
    session_manager: &session::SessionManager,
    session_id: &str,
    model_id: &str,
    system_prompt: Option<&str>,
    accumulated_content: &str,
) -> (i64, i64) {
    let mut total_input_tokens = 0i64;
    let mut total_output_tokens = 0i64;

    if let Some(session) = session_manager.get_session(session_id) {
        let mut estimate_messages = Vec::new();

        // Add system message
        let sys_msg_text = if let Some(prompt) = system_prompt {
            prompt.to_string()
        } else {
            llm::combine_prompts(llm::get_ask_prompt())
        };
        let system_msg = ChatCompletionRequestSystemMessage {
            content: sys_msg_text.into(),
            ..Default::default()
        };
        estimate_messages.push(system_msg.into());

        // Add all conversation history
        for msg in &session.messages {
            if msg.role == "user" {
                estimate_messages.push(
                    ChatCompletionRequestUserMessage {
                        content: msg.content.clone().into(),
                        ..Default::default()
                    }
                    .into(),
                );
            } else if msg.role == "assistant" {
                estimate_messages.push(
                    ChatCompletionRequestAssistantMessage {
                        content: Some(msg.content.clone().into()),
                        ..Default::default()
                    }
                    .into(),
                );
            }
        }

        let (estimated_input, _) = tokens::estimate_tokens(model_id, &estimate_messages);
        total_input_tokens = estimated_input;

        // Estimate output tokens from accumulated content
        if !accumulated_content.is_empty() {
            total_output_tokens = tokens::estimate_message_tokens(model_id, accumulated_content);
        }
    }

    (total_input_tokens, total_output_tokens)
}

#[derive(Debug, Deserialize)]
pub struct FileAttachment {
    pub filename: String,
    pub content: String,
}

#[derive(Debug, Deserialize)]
pub struct ChatRequest {
    pub message: String,
    #[serde(default)]
    pub session_id: Option<String>,
    #[serde(default)]
    pub files: Vec<FileAttachment>,
    #[serde(default)]
    pub system_prompt: Option<String>,
    pub agent_id: String,
    #[serde(default)]
    pub use_rag: Option<bool>,
    #[serde(default)]
    pub use_tools: Option<bool>,
}

#[derive(Debug, Serialize, Clone)]
pub struct Source {
    pub title: String,
    pub content: String,
}

#[derive(Debug, Serialize)]
#[serde(tag = "type")]
pub enum StreamEvent {
    #[serde(rename = "session")]
    Session { session_id: String },
    #[serde(rename = "sources")]
    Sources { sources: Vec<Source> },
    #[serde(rename = "content")]
    Content { text: String },
    #[serde(rename = "reasoning")]
    Reasoning { text: String },
    #[serde(rename = "tool_call")]
    ToolCall { name: String, arguments: String },
    #[serde(rename = "tool_result")]
    ToolResult { name: String, result: String },
    #[serde(rename = "usage")]
    Usage {
        input_tokens: i64,
        output_tokens: i64,
        reasoning_tokens: i64,
        cache_tokens: i64,
    },
    #[serde(rename = "tool_approval_request")]
    ToolApprovalRequest {
        approval_id: String,
        tool_name: String,
        tool_args: Value,
        tool_description: String,
    },
    #[serde(rename = "tool_approval_response")]
    ToolApprovalResponse { approval_id: String, approved: bool },
    #[serde(rename = "tool_invocation_completed")]
    ToolInvocationCompleted {
        name: String,
        arguments: Value,
        result: Option<String>,
        error: Option<String>,
    },
    #[serde(rename = "job_status")]
    JobStatus {
        job_id: i64,
        job_name: String,
        status: String,
        result: Option<String>,
        error: Option<String>,
        timestamp: i64,
    },
    #[serde(rename = "error")]
    Error { message: String },
    #[serde(rename = "done")]
    Done,
}

#[derive(Debug, Serialize)]
pub struct SessionMessage {
    pub role: String,
    pub content: String,
    pub sources: Vec<Source>,
    pub timestamp: i64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thinking_steps: Option<Vec<session::ThinkingStep>>,
}

#[derive(Debug, Clone, Serialize)]
pub struct TokenUsageResponse {
    pub total_tokens: i64,
    pub input_tokens: i64,
    pub output_tokens: i64,
    pub reasoning_tokens: i64,
    pub cache_tokens: i64,
    pub context_window: u32,
    pub context_utilization: f64,
}

#[derive(Debug, Serialize)]
pub struct SessionResponse {
    pub session_id: String,
    pub messages: Vec<SessionMessage>,
    pub created_at: i64,
    pub updated_at: i64,
    pub title: Option<String>,
    pub agent_id: Option<String>,
    pub token_usage: TokenUsageResponse,
    pub cost_usd: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct SessionListItem {
    pub session_id: String,
    pub message_count: usize,
    pub created_at: i64,
    pub updated_at: i64,
    pub preview: Option<String>,
    pub title: Option<String>,
    pub agent_id: Option<String>,
    pub token_usage: TokenUsageResponse,
    pub cost_usd: f64,
    pub is_readonly: bool,
}

#[derive(Debug, Serialize)]
pub struct SessionListResponse {
    pub sessions: Vec<SessionListItem>,
    pub total: usize,
}

#[derive(Debug, Deserialize)]
pub struct UpdateSessionRequest {
    pub title: String,
}

#[derive(Debug, Deserialize)]
pub struct LogsQuery {
    #[serde(default = "default_page")]
    pub page: usize,
    #[serde(default = "default_page_size")]
    pub page_size: usize,
    pub level: Option<String>,
    pub session_id: Option<String>,
}

fn default_page() -> usize {
    1
}

fn default_page_size() -> usize {
    50
}

#[derive(Debug, Serialize)]
pub struct LogEntryResponse {
    pub id: i64,
    pub timestamp: i64,
    pub level: String,
    pub target: String,
    pub message: String,
    pub session_id: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct LogsResponse {
    pub logs: Vec<LogEntryResponse>,
    pub total: usize,
    pub page: usize,
    pub page_size: usize,
    pub total_pages: usize,
}

/// Get session history by ID
pub async fn get_session(
    session_id: web::Path<String>,
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    match session_manager.get_session(&session_id) {
        Some(session) => {
            let response = SessionResponse {
                session_id: session.id,
                messages: session
                    .messages
                    .iter()
                    .map(|msg| SessionMessage {
                        role: msg.role.clone(),
                        content: msg.content.clone(),
                        sources: msg
                            .sources
                            .iter()
                            .map(|s| Source {
                                title: s.title.clone(),
                                content: s.content.clone(),
                            })
                            .collect(),
                        timestamp: msg.timestamp,
                        thinking_steps: msg.thinking_steps.clone(),
                    })
                    .collect(),
                created_at: session.created_at,
                updated_at: session.updated_at,
                title: session.title.clone(),
                agent_id: session.agent_id.clone(),
                token_usage: TokenUsageResponse {
                    total_tokens: session.token_usage.total_tokens,
                    input_tokens: session.token_usage.input_tokens,
                    output_tokens: session.token_usage.output_tokens,
                    reasoning_tokens: session.token_usage.reasoning_tokens,
                    cache_tokens: session.token_usage.cache_tokens,
                    context_window: session.token_usage.context_window,
                    context_utilization: session.token_usage.context_utilization,
                },
                cost_usd: session.cost_usd,
            };
            Ok(HttpResponse::Ok().json(response))
        }
        None => Ok(HttpResponse::NotFound().json(serde_json::json!({
            "error": "Session not found"
        }))),
    }
}

/// List all sessions with metadata
pub async fn list_sessions(
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    let session_ids = session_manager.list_sessions();
    let mut sessions = Vec::new();

    for session_id in session_ids {
        if let Some(session) = session_manager.get_session(&session_id) {
            sessions.push(build_session_list_item(&session));
        }
    }

    // Sort by updated_at descending (most recent first)
    sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));

    let total = sessions.len();
    Ok(HttpResponse::Ok().json(SessionListResponse { sessions, total }))
}

/// Delete a session by ID
pub async fn delete_session(
    session_id: web::Path<String>,
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    let deleted = session_manager.delete_session(&session_id);

    if deleted {
        // Broadcast session deletion
        broadcast_session_update(SessionUpdateEvent::Deleted {
            session_id: session_id.to_string(),
        });

        Ok(HttpResponse::Ok().json(serde_json::json!({
            "success": true,
            "message": "Session deleted successfully"
        })))
    } else {
        Ok(HttpResponse::NotFound().json(serde_json::json!({
            "error": "Session not found"
        })))
    }
}

/// Update a session (e.g., rename)
pub async fn update_session(
    session_id: web::Path<String>,
    update_request: web::Json<UpdateSessionRequest>,
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    // Validate title is not empty
    let title = update_request.title.trim();
    if title.is_empty() {
        return Ok(HttpResponse::BadRequest().json(serde_json::json!({
            "error": "Title cannot be empty"
        })));
    }

    match session_manager.update_session_title(&session_id, title.to_string()) {
        Ok(_) => {
            // Broadcast session update
            broadcast_session_update_for_session(&session_manager, &session_id);

            Ok(HttpResponse::Ok().json(serde_json::json!({
                "success": true,
                "message": "Session updated successfully"
            })))
        }
        Err(e) => Ok(HttpResponse::InternalServerError().json(serde_json::json!({
            "error": format!("Failed to update session: {}", e)
        }))),
    }
}

/// Server-Sent Events endpoint for session updates
pub async fn session_events() -> HttpResponse {
    use tokio_stream::StreamExt as TokioSE;

    if let Some(tx) = SESSION_UPDATE_BROADCASTER.get() {
        let rx = tx.subscribe();
        let filtered = TokioSE::filter_map(BroadcastStream::new(rx), |result| match result {
            Ok(event) => Some(event),
            Err(e) => {
                warn!("Session SSE broadcast receive error: {}", e);
                None
            }
        });

        let mapped = TokioSE::map(filtered, |event| {
            let data = serde_json::to_string(&event).unwrap_or_else(|e| {
                log::error!("Failed to serialize session event: {}", e);
                String::from("{}")
            });
            Ok::<_, actix_web::Error>(web::Bytes::from(format!(
                "event: session_update\ndata: {}\n\n",
                data
            )))
        });

        let stream = TokioSE::throttle(mapped, Duration::from_millis(100)); // Limit to 10 updates/sec

        HttpResponse::Ok()
            .insert_header((header::CONTENT_TYPE, "text/event-stream"))
            .insert_header((header::CACHE_CONTROL, "no-cache"))
            .insert_header((header::CONNECTION, "keep-alive"))
            .streaming(Box::pin(stream)
                as Pin<
                    Box<dyn Stream<Item = Result<web::Bytes, actix_web::Error>>>,
                >)
    } else {
        warn!("Session update broadcaster not initialized");
        HttpResponse::ServiceUnavailable().json(serde_json::json!({
            "error": "Session update broadcaster not available"
        }))
    }
}

/// Handles streaming chat requests
pub async fn chat_stream(
    body: web::Json<ChatRequest>,
    app_config: web::Data<Arc<config::Config>>,
    session_manager: web::Data<Arc<session::SessionManager>>,
    approval_map: web::Data<ApprovalStateMap>,
    rag_system: web::Data<Option<Arc<RagSystem>>>,
) -> Result<HttpResponse, Error> {
    let question = body.message.clone();
    let use_rag = body.use_rag.unwrap_or(false);
    let mut use_tools = body.use_tools.unwrap_or(false);

    // Validate file sizes (10MB limit per file)
    const MAX_FILE_SIZE: usize = 10 * 1024 * 1024;
    for file in &body.files {
        if file.content.len() > MAX_FILE_SIZE {
            return Ok(HttpResponse::BadRequest().json(serde_json::json!({
                "error": format!("File '{}' exceeds size limit of 10MB ({} bytes > {} bytes)",
                    file.filename, file.content.len(), MAX_FILE_SIZE)
            })));
        }
    }

    let files: Vec<session::FileAttachment> = body
        .files
        .iter()
        .map(|f| session::FileAttachment {
            filename: f.filename.clone(),
            content: f.content.clone(),
        })
        .collect();
    let system_prompt = body.system_prompt.clone();
    let system_prompt_for_stream = system_prompt.clone(); // Clone for use inside stream
    let app_config_clone = app_config.get_ref().clone();
    let session_manager_clone = session_manager.get_ref().clone();
    let agent_id = body.agent_id.clone();
    let agent_id_for_stream = agent_id.clone(); // Clone for use inside stream

    // Get agent to extract model_id and context_window for token estimation and session storage
    let (model_id, context_window) = match app_config_clone.get_agent(&agent_id) {
        Some(agent) => {
            let ctx_window = agent
                .context_window
                .unwrap_or(app_config_clone.context_window);
            // Enforce agent-level use_tools setting: if the agent disables tools, override the client request
            if !agent.use_tools {
                use_tools = false;
            }
            (agent.model.clone(), ctx_window)
        }
        None => {
            return Ok(HttpResponse::BadRequest().json(serde_json::json!({
                "error": format!("Agent '{}' not found", agent_id)
            })));
        }
    };

    // Get or create session
    let session_id = body
        .session_id
        .clone()
        .unwrap_or_else(|| session_manager_clone.create_session());

    // Create SSE stream
    let stream = async_stream::stream! {
        // Send session ID first
        let session_event = StreamEvent::Session {
            session_id: session_id.clone(),
        };
        let json = serde_json::to_string(&session_event).unwrap_or_default();
        yield Ok::<_, actix_web::Error>(
            web::Bytes::from(format!("data: {}\n\n", json))
        );

        // Query RAG if enabled
        let mut rag_sources = Vec::new();
        if use_rag {
            if let Some(rag_sys) = rag_system.as_ref() {
                match rag_sys.query.execute_structured(&question).await {
                    Ok(results) => {
                        for result in results.iter() {
                            rag_sources.push(Source {
                                title: result.filename.clone(),
                                content: result.chunk_text.clone(),
                            });
                        }

                        // Send RAG sources as sources event
                        if !rag_sources.is_empty() {
                            let sources_event = StreamEvent::Sources {
                                sources: rag_sources.clone(),
                            };
                            let json = serde_json::to_string(&sources_event).unwrap_or_default();
                            yield Ok::<_, actix_web::Error>(
                                web::Bytes::from(format!("data: {}\n\n", json))
                            );
                        }
                    }
                    Err(e) => {
                        warn!("⚠️  RAG query failed: {}", e);
                    }
                }
            } else {
                warn!("⚠️  RAG requested but system not available");
            }
        }

        // Add user message to session and get file sources
        let sources = match session_manager_clone.add_user_message(
            &session_id,
            question.clone(),
            files.clone(),
        ) {
            Ok(sources) => sources,
            Err(e) => {
                let error_event = StreamEvent::Error {
                    message: format!("Failed to add message to session: {}", e),
                };
                let json = serde_json::to_string(&error_event).unwrap_or_default();
                yield Ok::<_, actix_web::Error>(
                    web::Bytes::from(format!("data: {}\n\n", json))
                );
                return;
            }
        };

        // Send file sources if any (these are from attached files, not RAG)
        if !sources.is_empty() {
            let sources_event = StreamEvent::Sources {
                sources: sources.iter().map(|s| Source {
                    title: s.title.clone(),
                    content: s.content.clone(),
                }).collect(),
            };
            let json = serde_json::to_string(&sources_event).unwrap_or_default();
            yield Ok::<_, actix_web::Error>(
                web::Bytes::from(format!("data: {}\n\n", json))
            );
        }

        match create_chat_stream(
            &session_id,
            &question,
            &files,
            system_prompt.as_deref(),
            &agent_id,
            &app_config_clone,
            &session_manager_clone,
            approval_map.get_ref(),
            if use_rag && !rag_sources.is_empty() { Some(rag_sources.clone()) } else { None },
            use_tools,
        ).await {
            Ok(content_stream) => {
                // Accumulate assistant content and token usage as we stream
                let mut accumulated_content = String::new();
                let mut accumulated_reasoning = String::new();
                let mut total_input_tokens = 0i64;
                let mut total_output_tokens = 0i64;
                let mut total_reasoning_tokens = 0i64;
                let mut total_cache_tokens = 0i64;
                let mut received_usage = false; // Track if provider sent usage
                // Track thinking steps in order as they occur during streaming
                let mut thinking_steps_ordered: Vec<session::ThinkingStep> = Vec::new();
                let mut step_order = 0i32;

                // Track reasoning blocks separately - don't merge them
                let mut last_closed_think_pos = 0;

                // Stream each content chunk as it arrives
                let mut pinned_stream = Box::pin(content_stream);
                while let Some(result) = pinned_stream.next().await {
                    match result {
                        Ok(chunk) => {
                            // Accumulate content chunks
                            if let StreamEvent::Content { ref text } = chunk {
                                accumulated_content.push_str(text);

                                // Check if we completed any <think>...</think> blocks
                                // Process each closed block as a separate reasoning step
                                loop {
                                    let search_start = last_closed_think_pos;
                                    if let Some(relative_start) = accumulated_content[search_start..].find("<think>") {
                                        let absolute_start = search_start + relative_start;
                                        if let Some(relative_end) = accumulated_content[absolute_start..].find("</think>") {
                                            let absolute_end = absolute_start + relative_end;
                                            let reasoning_text = accumulated_content[absolute_start + 7..absolute_end].to_string();
                                            if !reasoning_text.trim().is_empty() {
                                                // Add reasoning step immediately as a SEPARATE step (don't merge)
                                                thinking_steps_ordered.push(session::ThinkingStep {
                                                    step_type: "reasoning".to_string(),
                                                    step_order,
                                                    content: Some(reasoning_text),
                                                    tool_name: None,
                                                    tool_arguments: None,
                                                    tool_result: None,
                                                    tool_error: None,
                                                    content_before_tool: None,
                                                });
                                                step_order += 1;
                                            }
                                            last_closed_think_pos = absolute_end + 8;
                                        } else {
                                            break; // Incomplete block, wait for more content
                                        }
                                    } else {
                                        break; // No more think blocks
                                    }
                                }
                            }

                            // Accumulate reasoning chunks
                            if let StreamEvent::Reasoning { ref text } = chunk {
                                accumulated_reasoning.push_str(text);
                            }

                            // Accumulate token usage
                            if let StreamEvent::Usage { input_tokens, output_tokens, reasoning_tokens, cache_tokens } = chunk {
                                total_input_tokens += input_tokens;
                                total_output_tokens += output_tokens;
                                total_reasoning_tokens += reasoning_tokens;
                                total_cache_tokens += cache_tokens;
                                received_usage = true;
                            }

                            // Add tool invocation to thinking steps immediately
                            // This preserves the order: when a tool completes, it gets added right after the last reasoning step
                            if let StreamEvent::ToolInvocationCompleted { name, arguments, result, error } = &chunk {
                                // Capture content accumulated before this tool
                                let content_snapshot = accumulated_content.trim().to_string();

                                // Add tool as thinking step immediately (preserves order)
                                thinking_steps_ordered.push(session::ThinkingStep {
                                    step_type: "tool".to_string(),
                                    step_order,
                                    content: None,
                                    tool_name: Some(name.clone()),
                                    tool_arguments: Some(arguments.clone()),
                                    tool_result: result.clone(),
                                    tool_error: error.clone(),
                                    content_before_tool: if content_snapshot.is_empty() {
                                        None
                                    } else {
                                        Some(content_snapshot)
                                    },
                                });
                                step_order += 1;
                            }

                            let json = serde_json::to_string(&chunk).unwrap_or_default();
                            yield Ok::<_, actix_web::Error>(
                                web::Bytes::from(format!("data: {}\n\n", json))
                            );
                        }
                        Err(e) => {
                            let error_event = StreamEvent::Error {
                                message: e.to_string(),
                            };
                            let json = serde_json::to_string(&error_event).unwrap_or_default();
                            yield Ok::<_, actix_web::Error>(
                                web::Bytes::from(format!("data: {}\n\n", json))
                            );
                            break;
                        }
                    }
                }

                // Add assistant message to session with sources
                // Parse out ALL <think> and <tool_call> tags from accumulated content for final display
                let final_content = sanitize_assistant_content(&accumulated_content);

                // Use the thinking steps we built during streaming
                let thinking_steps_opt = if thinking_steps_ordered.is_empty() {
                    None
                } else {
                    Some(thinking_steps_ordered)
                };

                // Trim whitespace and check if we have actual content
                let final_content_trimmed = final_content.trim();

                // Save assistant message to session with both file sources and RAG sources
                let mut all_sources = sources.clone();
                all_sources.extend(rag_sources.iter().map(|s| session::Source {
                    title: s.title.clone(),
                    content: s.content.clone(),
                }));

                if !final_content_trimmed.is_empty() || thinking_steps_opt.is_some() {
                    match session_manager_clone.add_assistant_message(
                        &session_id,
                        final_content_trimmed.to_string(),
                        all_sources,
                        thinking_steps_opt,
                    ) {
                        Ok(_) => {},
                        Err(e) => debug!("Failed to save assistant message: {}", e),
                    }
                } else {
                    // Skip empty messages
                }

                // If provider didn't send usage stats, estimate them client-side
                if !received_usage {
                    debug!("Provider didn't report token usage, estimating client-side for model: {}", model_id);

                    let (estimated_input, estimated_output) = estimate_and_send_usage(
                        &session_manager_clone,
                        &session_id,
                        &model_id,
                        system_prompt_for_stream.as_deref(),
                        &accumulated_content,
                    ).await;

                    total_input_tokens = estimated_input;
                    total_output_tokens = estimated_output;

                    // Send estimated usage to UI
                    let usage_event = StreamEvent::Usage {
                        input_tokens: total_input_tokens,
                        output_tokens: total_output_tokens,
                        reasoning_tokens: total_reasoning_tokens,
                        cache_tokens: total_cache_tokens,
                    };
                    let json = serde_json::to_string(&usage_event).unwrap_or_default();
                    yield Ok::<_, actix_web::Error>(
                        web::Bytes::from(format!("data: {}\n\n", json))
                    );
                }

                // Update session with token usage and model info
                if (total_input_tokens > 0 || total_output_tokens > 0)
                    && let Err(e) = session_manager_clone.update_token_usage(
                        &session_id,
                        &agent_id_for_stream,
                        crate::session::TokenUsageUpdate {
                            input_tokens: total_input_tokens,
                            output_tokens: total_output_tokens,
                            reasoning_tokens: total_reasoning_tokens,
                            cache_tokens: total_cache_tokens,
                            context_window,
                        },
                    )
                {
                    debug!("Failed to update token usage: {}", e);
                }

                // Send done event
                let done_event = StreamEvent::Done;
                let json = serde_json::to_string(&done_event).unwrap_or_default();
                yield Ok::<_, actix_web::Error>(
                    web::Bytes::from(format!("data: {}\n\n", json))
                );

                // Broadcast session update via SSE
                broadcast_session_update_for_session(&session_manager_clone, &session_id);
            }
            Err(e) => {
                let error_event = StreamEvent::Error {
                    message: e.to_string(),
                };
                let json = serde_json::to_string(&error_event).unwrap_or_default();
                yield Ok::<_, actix_web::Error>(
                    web::Bytes::from(format!("data: {}\n\n", json))
                );
            }
        }
    };

    Ok(HttpResponse::Ok()
        .content_type("text/event-stream")
        .insert_header(("Cache-Control", "no-cache"))
        .insert_header(("X-Accel-Buffering", "no"))
        .streaming(Box::pin(stream)))
}

#[allow(unused_variables)] // approval_map is used inside async_stream::stream! macro
async fn create_chat_stream(
    session_id: &str,
    question: &str,
    files: &[session::FileAttachment],
    system_prompt: Option<&str>,
    agent_id: &str,
    app_config: &config::Config,
    session_manager: &session::SessionManager,
    approval_map: &ApprovalStateMap,
    rag_sources: Option<Vec<Source>>,
    use_tools: bool,
) -> Result<
    impl futures::Stream<Item = Result<StreamEvent, Box<dyn std::error::Error + Send + Sync>>>,
    Box<dyn std::error::Error + Send + Sync>,
> {
    // Get agent config
    let agent = match app_config.get_agent(agent_id) {
        Some(a) => a,
        None => {
            return Err(format!("Agent '{}' not found", agent_id).into());
        }
    };

    let model_id = agent.model.clone();
    debug!("Using API URL: {}", app_config.api_url);
    debug!("Using Agent: {} (model: {})", agent_id, model_id);

    let config = OpenAIConfig::new()
        .with_api_base(&app_config.api_url)
        .with_api_key(app_config.get_api_key());

    let client = Client::with_config(config);

    // Build user message with template rendering support
    let mut user_message = String::new();

    // Add RAG context first if available
    if let Some(ref sources) = rag_sources
        && !sources.is_empty()
    {
        user_message.push_str("# Retrieved Context from Documents\n\n");
        for (idx, source) in sources.iter().enumerate() {
            user_message.push_str(&format!(
                "## Document {}: {}\n\n{}\n\n",
                idx + 1,
                source.title,
                source.content
            ));
        }
        user_message.push_str("---\n\n");
        debug!("✅ Added {} RAG sources to context", sources.len());
    }

    // Add file contents if present
    if !files.is_empty() {
        for file in files {
            user_message.push_str(&format!(
                "Here is the content of '{}':\n\n```\n{}\n```\n\n",
                file.filename, file.content
            ));
        }
    }

    user_message.push_str(&format!("User query: {}", question));

    let default_prompt = llm::combine_prompts(llm::get_ask_prompt());
    // Use agent's prompt if available, then system_prompt parameter, then default
    let final_system_prompt = agent
        .prompt
        .as_deref()
        .or(system_prompt)
        .unwrap_or(&default_prompt);

    // Render template variables in system message
    let renderer = template::TemplateRenderer::new();
    let system_message = renderer
        .render_string(final_system_prompt)
        .unwrap_or_else(|e| {
            warn!("Failed to render system prompt template: {}", e);
            final_system_prompt.to_string()
        });

    // Get conversation history from session
    let session = session_manager
        .get_session(session_id)
        .ok_or("Session not found")?;

    // Capture agent_id for use in stream
    let agent_id_owned = agent_id.to_string();

    // Build conversation messages from session history
    let mut messages = build_messages_from_history(&session, system_message);

    // Add the current user message with full context
    messages.push(
        ChatCompletionRequestUserMessage {
            content: user_message.into(),
            ..Default::default()
        }
        .into(),
    );

    let mut tool_calls: Vec<ChatCompletionMessageToolCall> = Vec::new();

    let output_stream = async_stream::stream! {
        loop {
        let mut request_builder = CreateChatCompletionRequestArgs::default();
        request_builder
            .model(model_id.clone())
            .messages(messages.clone());

        // Only add tools if enabled
        if use_tools {
            request_builder.tools(tools::get_tools());
        }

        let request = request_builder
            .stream_options(ChatCompletionStreamOptions {
                include_usage: Some(true),
                include_obfuscation: None,
            })
            .build();

            let request = match request {
                Ok(req) => req,
                Err(e) => {
                    yield Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>);
                    break;
                }
            };

            debug!("Sending streaming request...");

            let stream_result = client.chat().create_stream(request).await;
            let mut stream = match stream_result {
                Ok(s) => s,
                Err(e) => {
                    yield Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>);
                    break;
                }
            };

            tool_calls.clear();

            while let Some(result) = stream.next().await {
                let response = match result {
                    Ok(r) => r,
                    Err(e) => {
                        yield Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>);
                        break;
                    }
                };

                // Yield token usage statistics from streaming response
                if let Some(usage) = &response.usage {
                    debug!(
                        "Token usage - Prompt: {}, Completion: {}, Total: {}",
                        usage.prompt_tokens, usage.completion_tokens, usage.total_tokens
                    );

                    yield Ok(StreamEvent::Usage {
                        input_tokens: usage.prompt_tokens as i64,
                        output_tokens: usage.completion_tokens as i64,
                        reasoning_tokens: 0, // Not provided by OpenAI streaming API
                        cache_tokens: 0,     // Not provided by OpenAI streaming API
                    });
                }

                for choice in response.choices {
                    // Handle content - yield it immediately
                    if let Some(content) = &choice.delta.content {
                        yield Ok(StreamEvent::Content {
                            text: content.clone(),
                        });
                    }

                // Handle tool calls
                if let Some(tool_calls_delta) = &choice.delta.tool_calls {
                    for tool_call_delta in tool_calls_delta {
                        let index = tool_call_delta.index as usize;

                        // Ensure we have enough space in the vector
                        while tool_calls.len() <= index {
                            tool_calls.push(ChatCompletionMessageToolCall {
                                id: String::new(),
                                function: Default::default(),
                            });
                        }

                        // Update the tool call with delta information
                        if let Some(id) = &tool_call_delta.id {
                            tool_calls[index].id = id.clone();
                        }
                        if let Some(function) = &tool_call_delta.function {
                            if let Some(name) = &function.name {
                                tool_calls[index].function.name = name.clone();
                            }
                            if let Some(args) = &function.arguments {
                                tool_calls[index].function.arguments.push_str(args);
                            }
                        }
                    }
                }

                // Handle finish reason
                if let Some(finish_reason) = &choice.finish_reason {
                    match finish_reason {
                        FinishReason::ToolCalls => {
                            debug!("Executing tool calls...");

                            // Add assistant message with tool calls
                            let assistant_tool_calls: Vec<ChatCompletionMessageToolCalls> =
                                tool_calls.iter().map(|tc| tc.clone().into()).collect();

                            messages.push(
                                ChatCompletionRequestAssistantMessage {
                                    content: None,
                                    tool_calls: Some(assistant_tool_calls),
                                    ..Default::default()
                                }
                                .into(),
                            );

                            // Execute tools with approval handling
                            for tool_call in tool_calls.iter() {
                                let name = &tool_call.function.name;
                                let args_str = &tool_call.function.arguments;
                                let tool_call_id = &tool_call.id;

                                // Parse arguments
                                let args_value: Value = match args_str.parse() {
                                    Ok(v) => v,
                                    Err(e) => {
                                        let error_result = json!({
                                            "error": format!("Failed to parse tool arguments: {}", e)
                                        });
                                        messages.push(
                                            ChatCompletionRequestToolMessage {
                                                content: error_result.to_string().into(),
                                                tool_call_id: tool_call_id.clone(),
                                            }
                                            .into(),
                                        );
                                        continue;
                                    }
                                };

                                // Check permission status
                                let permission_status = tools::check_tool_permission(name, &args_value, &agent_id_owned, app_config);

                                debug!("Tool '{}' permission status: {:?}", name, permission_status);

                                match permission_status {
                                    tools::ToolPermissionStatus::Denied { reason } => {
                                        // Tool is denied, don't execute
                                        let deny_result = json!({
                                            "error": reason,
                                            "skipped": true
                                        });
                                        messages.push(
                                            ChatCompletionRequestToolMessage {
                                                content: deny_result.to_string().into(),
                                                tool_call_id: tool_call_id.clone(),
                                            }
                                            .into(),
                                        );
                                    }
                                    tools::ToolPermissionStatus::Allowed => {
                                        // Tool is auto-allowed, execute directly
                                        let result = tools::execute_tool_direct(name, &args_value, app_config).await;

                                        // Emit tool invocation completed event
                                        yield Ok(StreamEvent::ToolInvocationCompleted {
                                            name: name.clone(),
                                            arguments: args_value.clone(),
                                            result: Some(result.to_string()),
                                            error: None,
                                        });

                                        messages.push(
                                            ChatCompletionRequestToolMessage {
                                                content: result.to_string().into(),
                                                tool_call_id: tool_call_id.clone(),
                                            }
                                            .into(),
                                        );
                                    }
                                    tools::ToolPermissionStatus::NeedsApproval => {
                                        use uuid::Uuid;
                                        use std::time::Duration;

                                        // Generate unique approval ID
                                        let approval_id = Uuid::new_v4().to_string();

                                        // Create oneshot channel for approval response
                                        let (sender, receiver) = tokio::sync::oneshot::channel::<bool>();

                                        // Store approval state in map
                                        {
                                            let mut approvals = approval_map.lock().await;
                                            approvals.insert(approval_id.clone(), ApprovalState {
                                                tool_name: name.clone(),
                                                tool_args: args_value.clone(),
                                                tool_call_id: tool_call_id.clone(),
                                                agent_id: agent_id_owned.clone(),
                                                sender,
                                                created_at: Instant::now(),
                                            });
                                        }

                                        // Yield approval request event
                                        yield Ok(StreamEvent::ToolApprovalRequest {
                                            approval_id: approval_id.clone(),
                                            tool_name: name.clone(),
                                            tool_args: args_value.clone(),
                                            tool_description: get_tool_description(name),
                                        });

                                        // Wait for approval with 5 minute timeout
                                        let approved = match tokio::time::timeout(
                                            Duration::from_secs(300),
                                            receiver
                                        ).await {
                                            Ok(Ok(decision)) => {
                                                decision
                                            }
                                            Ok(Err(_)) => {
                                                warn!("Tool approval channel closed without response");
                                                false
                                            }
                                            Err(_) => {
                                                warn!("Tool approval timed out after 5 minutes");
                                                false
                                            }
                                        };

                                        // Clean up from map
                                        {
                                            let mut approvals = approval_map.lock().await;
                                            approvals.remove(&approval_id);
                                        }

                                        // Yield approval response event
                                        yield Ok(StreamEvent::ToolApprovalResponse {
                                            approval_id: approval_id.clone(),
                                            approved,
                                        });

                                        // Execute based on approval
                                        if approved {
                                            let result = tools::execute_tool_direct(name, &args_value, app_config).await;

                                            // Emit tool invocation completed event
                                            yield Ok(StreamEvent::ToolInvocationCompleted {
                                                name: name.clone(),
                                                arguments: args_value.clone(),
                                                result: Some(result.to_string()),
                                                error: None,
                                            });

                                            messages.push(
                                                ChatCompletionRequestToolMessage {
                                                    content: result.to_string().into(),
                                                    tool_call_id: tool_call_id.clone(),
                                                }
                                                .into(),
                                            );
                                        } else {
                                            let reject_result = json!({
                                                "message": format!("Tool '{}' was not executed because you rejected it.", name),
                                                "skipped": true
                                            });

                                            // Emit tool invocation completed event for rejection to record in thinking steps
                                            yield Ok(StreamEvent::ToolInvocationCompleted {
                                                name: name.clone(),
                                                arguments: args_value.clone(),
                                                result: None,
                                                error: Some("Tool execution rejected by user".to_string()),
                                            });

                                            messages.push(
                                                ChatCompletionRequestToolMessage {
                                                    content: reject_result.to_string().into(),
                                                    tool_call_id: tool_call_id.clone(),
                                                }
                                                .into(),
                                            );
                                        }
                                    }
                                }
                            }

                            // Continue the loop to make another request with tool results
                            break;
                        }
                        FinishReason::Stop => {
                            return;
                        }
                        _ => {}
                    }
                }
            }


        }

            // If we had tool calls, the loop continues to make another request
            // If not, we should have returned by now
            if tool_calls.is_empty() {
                break;
            }
        }
    };

    Ok(output_stream)
}

/// Get logs with pagination
pub async fn get_logs(
    query: web::Query<LogsQuery>,
    app_config: web::Data<Arc<config::Config>>,
) -> Result<HttpResponse, Error> {
    let db_path = &app_config.get_ref().database_path;

    // Calculate offset from page number
    let offset = (query.page.saturating_sub(1)) * query.page_size;

    // Query logs with limit for total count + 1 to check if there are more
    let all_logs = logger::query_logs(
        db_path,
        None, // Get all to calculate total
        query.level.as_deref(),
        query.session_id.as_deref(),
    )
    .map_err(|e| {
        actix_web::error::ErrorInternalServerError(format!("Failed to query logs: {}", e))
    })?;

    let total = all_logs.len();
    let total_pages = total.div_ceil(query.page_size);

    // Get the paginated subset
    let logs: Vec<LogEntryResponse> = all_logs
        .into_iter()
        .skip(offset)
        .take(query.page_size)
        .map(|entry| LogEntryResponse {
            id: entry.id,
            timestamp: entry.timestamp,
            level: entry.level,
            target: entry.target,
            message: entry.message,
            session_id: entry.session_id,
        })
        .collect();

    Ok(HttpResponse::Ok().json(LogsResponse {
        logs,
        total,
        page: query.page,
        page_size: query.page_size,
        total_pages,
    }))
}

#[derive(Debug, Serialize)]
pub struct ConfigResponse {
    pub api_url: String,
    pub context_window: u32,
    pub rag_enabled: bool,
    pub web_sounds: bool,
    pub audio_enabled: bool,
}

/// Get API configuration (default model, etc.)
pub async fn get_config(app_config: web::Data<Arc<config::Config>>) -> Result<HttpResponse, Error> {
    debug!("Fetching API configuration");

    let response = ConfigResponse {
        api_url: app_config.api_url.clone(),
        context_window: app_config.context_window,
        rag_enabled: app_config.rag.enabled,
        web_sounds: app_config.web.sounds,
        audio_enabled: app_config.audio.enabled,
    };

    Ok(HttpResponse::Ok().json(response))
}

#[derive(Debug, Serialize)]
pub struct AgentInfo {
    pub id: String,
    pub name: String,
    pub description: String,
    pub model: String,
    pub enabled: bool,
    pub use_tools: bool,
    pub permissions: crate::agent::AgentPermissions,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pricing_model: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub context_window: Option<u32>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub suggestions: Vec<String>,
}

#[derive(Debug, Serialize)]
pub struct AgentsResponse {
    pub agents: Vec<AgentInfo>,
    pub default_agent: String,
}

pub async fn get_agents(app_config: web::Data<Arc<config::Config>>) -> Result<HttpResponse, Error> {
    debug!("Fetching agents configuration");

    let agents: Vec<AgentInfo> = app_config
        .agents
        .agents
        .iter()
        .filter(|(_, agent)| agent.enabled)
        .map(|(id, agent)| AgentInfo {
            id: id.clone(),
            name: agent.name.clone(),
            description: agent.description.clone(),
            model: agent.model.clone(),
            enabled: agent.enabled,
            use_tools: agent.use_tools,
            permissions: agent.permissions.clone(),
            pricing_model: agent.pricing_model.clone(),
            context_window: agent.context_window,
            suggestions: agent.suggestions.clone(),
        })
        .collect();

    Ok(HttpResponse::Ok().json(AgentsResponse {
        agents,
        default_agent: app_config.agents.default_agent.clone(),
    }))
}

/// Response for agent file content
#[derive(Debug, Serialize)]
pub struct AgentContentResponse {
    pub id: String,
    pub name: String,
    pub content: String,
}

/// Get the raw markdown content for a specific agent
pub async fn get_agent_content(
    agent_id: web::Path<String>,
    app_config: web::Data<Arc<config::Config>>,
) -> Result<HttpResponse, Error> {
    let agent_id = agent_id.into_inner();
    debug!("Fetching content for agent: {}", agent_id);

    // Validate agent_id to prevent path traversal and invalid characters
    // Allow only alphanumerics, underscore, and hyphen.
    let is_valid_agent_id = !agent_id.is_empty()
        && agent_id
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-');
    if !is_valid_agent_id {
        return Ok(HttpResponse::BadRequest().json(serde_json::json!({
            "error": "Invalid agent id"
        })));
    }

    // Check if agent exists in configuration
    let agent = app_config.agents.agents.get(&agent_id);
    if agent.is_none() {
        return Ok(HttpResponse::NotFound().json(serde_json::json!({
            "error": format!("Agent '{}' not found", agent_id)
        })));
    }

    // Read the agent file from the agents directory
    let agents_dir = crate::agent::get_agents_dir(app_config.config_dir.as_deref());
    let agent_file = agents_dir.join(format!("{}.md", agent_id));

    // Resolve symlinks and canonicalize paths to prevent traversal
    let agents_dir = agents_dir.canonicalize().map_err(|e| {
        actix_web::error::ErrorInternalServerError(format!("Agents directory not found: {}", e))
    })?;
    let agent_file = match agent_file.canonicalize() {
        Ok(p) => p,
        Err(_) => {
            return Ok(HttpResponse::NotFound().json(serde_json::json!({
                "error": format!("Agent '{}' not found", agent_id)
            })));
        }
    };

    // Ensure the resolved path is still within the agents directory
    if !agent_file.starts_with(&agents_dir) {
        return Ok(HttpResponse::Forbidden().json(serde_json::json!({
            "error": "Access denied"
        })));
    }

    match std::fs::read_to_string(&agent_file) {
        Ok(content) => Ok(HttpResponse::Ok().json(AgentContentResponse {
            id: agent_id,
            name: agent.unwrap().name.clone(),
            content,
        })),
        Err(e) => Ok(HttpResponse::InternalServerError().json(serde_json::json!({
            "error": format!("Failed to read agent file: {}", e)
        }))),
    }
}

/// Response structure for agent token statistics
#[derive(Debug, Serialize)]
pub struct AgentTokenStatsResponse {
    pub agent_id: String,
    pub total_sessions: i64,
    pub total_tokens: i64,
    pub input_tokens: i64,
    pub output_tokens: i64,
    pub reasoning_tokens: i64,
    pub cache_tokens: i64,
    pub total_cost_usd: f64,
    pub avg_cost_per_session: f64,
    pub first_used_at: i64,
    pub last_used_at: i64,
}

/// Response structure for all agent statistics
#[derive(Debug, Serialize)]
pub struct AllAgentTokenStatsResponse {
    pub agents: Vec<AgentTokenStatsResponse>,
}

/// Get token statistics for all agents
pub async fn get_agent_stats(
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    debug!("Fetching agent token statistics");

    let all_stats = match session_manager.get_all_agent_token_stats() {
        Ok(stats) => stats,
        Err(e) => {
            return Ok(HttpResponse::InternalServerError().json(serde_json::json!({
                "error": format!("Failed to fetch agent token statistics: {}", e)
            })));
        }
    };

    let agents: Vec<AgentTokenStatsResponse> = all_stats
        .into_iter()
        .map(|stat| {
            let avg_cost_per_session = if stat.total_sessions > 0 {
                stat.total_cost_usd / stat.total_sessions as f64
            } else {
                0.0
            };

            AgentTokenStatsResponse {
                agent_id: stat.agent_id,
                total_sessions: stat.total_sessions,
                total_tokens: stat.total_tokens,
                input_tokens: stat.input_tokens,
                output_tokens: stat.output_tokens,
                reasoning_tokens: stat.reasoning_tokens,
                cache_tokens: stat.cache_tokens,
                total_cost_usd: stat.total_cost_usd,
                avg_cost_per_session,
                first_used_at: stat.first_used_at,
                last_used_at: stat.last_used_at,
            }
        })
        .collect();

    Ok(HttpResponse::Ok().json(AllAgentTokenStatsResponse { agents }))
}

/// Get token statistics for a specific agent
pub async fn get_agent_stats_by_id(
    agent_id: web::Path<String>,
    session_manager: web::Data<Arc<session::SessionManager>>,
) -> Result<HttpResponse, Error> {
    debug!("Fetching token statistics for agent: {}", agent_id);

    match session_manager.get_agent_token_stats(&agent_id) {
        Ok(Some(stat)) => {
            let avg_cost_per_session = if stat.total_sessions > 0 {
                stat.total_cost_usd / stat.total_sessions as f64
            } else {
                0.0
            };

            let response = AgentTokenStatsResponse {
                agent_id: stat.agent_id,
                total_sessions: stat.total_sessions,
                total_tokens: stat.total_tokens,
                input_tokens: stat.input_tokens,
                output_tokens: stat.output_tokens,
                reasoning_tokens: stat.reasoning_tokens,
                cache_tokens: stat.cache_tokens,
                total_cost_usd: stat.total_cost_usd,
                avg_cost_per_session,
                first_used_at: stat.first_used_at,
                last_used_at: stat.last_used_at,
            };

            Ok(HttpResponse::Ok().json(response))
        }
        Ok(None) => Ok(HttpResponse::NotFound().json(serde_json::json!({
            "error": "Agent statistics not found"
        }))),
        Err(e) => Ok(HttpResponse::InternalServerError().json(serde_json::json!({
            "error": format!("Failed to fetch agent token statistics: {}", e)
        }))),
    }
}

#[derive(Debug, Deserialize)]
pub struct ToolApprovalRequest {
    pub approval_id: String,
    pub approved: bool,
    #[serde(default)]
    pub save_decision: bool,
    #[serde(default)]
    pub scope: String, // "tool" or "tool:specific" (e.g., "bash:ls")
}

#[derive(Debug, Serialize)]
pub struct ToolApprovalResponse {
    pub success: bool,
    pub message: String,
}

/// Handle tool approval requests from the web UI
pub async fn handle_tool_approval(
    body: web::Json<ToolApprovalRequest>,
    approval_map: web::Data<ApprovalStateMap>,
) -> Result<HttpResponse, Error> {
    // Find the pending approval
    let mut approvals = approval_map.lock().await;

    if let Some(approval_state) = approvals.remove(&body.approval_id) {
        // Send the approval decision through the channel
        if approval_state.sender.send(body.approved).is_err() {
            return Ok(
                HttpResponse::InternalServerError().json(ToolApprovalResponse {
                    success: false,
                    message: "Failed to send approval response".to_string(),
                }),
            );
        }

        // If save_decision is true, update the config file
        if body.save_decision {
            let tool_name = &approval_state.tool_name;
            let agent_id = &approval_state.agent_id;
            let scope = if body.scope.is_empty() {
                tool_name.clone()
            } else {
                body.scope.clone()
            };

            let mut config = config::Config::load();
            let result = if body.approved {
                config.allow_tool_for_agent(agent_id, &scope)
            } else {
                config.deny_tool_for_agent(agent_id, &scope)
            };

            if let Err(e) = result {
                return Ok(
                    HttpResponse::InternalServerError().json(ToolApprovalResponse {
                        success: false,
                        message: format!("Approval processed but failed to save to config: {}", e),
                    }),
                );
            }
        }

        Ok(HttpResponse::Ok().json(ToolApprovalResponse {
            success: true,
            message: "Approval processed successfully".to_string(),
        }))
    } else {
        Ok(HttpResponse::NotFound().json(ToolApprovalResponse {
            success: false,
            message: "Approval request not found or expired".to_string(),
        }))
    }
}

// ========================================
// RAG (Retrieval-Augmented Generation) Endpoints
// ========================================

use crate::rag::RagSystem;

#[derive(Debug, Deserialize)]
pub struct RagQueryRequest {
    pub query: String,
    #[serde(default)]
    pub top_k: Option<usize>,
}

#[derive(Debug, Serialize)]
pub struct RagQueryResponse {
    pub context: String,
    pub sources: Vec<RagSource>,
}

#[derive(Debug, Serialize)]
pub struct RagSource {
    pub filename: String,
    pub text: String,
    pub relevance: f32,
}

#[derive(Debug, Serialize)]
pub struct DocumentListResponse {
    pub documents: Vec<DocumentSummary>,
}

#[derive(Debug, Serialize)]
pub struct DocumentSummary {
    pub id: i64,
    pub filename: String,
    pub file_size: i64,
    pub updated_at: i64,
}

#[derive(Debug, Serialize)]
pub struct RagStatsResponse {
    pub doc_count: i64,
    pub chunk_count: i64,
    pub embedding_count: i64,
    pub avg_chunks_per_doc: f64,
}

#[derive(Debug, Serialize)]
pub struct RagResponse {
    pub success: bool,
    pub message: String,
}

/// Query RAG index for relevant context
pub async fn rag_query(
    body: web::Json<RagQueryRequest>,
    rag_system: web::Data<Option<Arc<RagSystem>>>,
) -> Result<HttpResponse, Error> {
    let Some(rag_system) = rag_system.as_ref() else {
        return Ok(HttpResponse::ServiceUnavailable().json(json!({
            "error": "RAG system is not enabled"
        })));
    };

    match rag_system.query.execute_structured(&body.query).await {
        Ok(results) => {
            let mut context = String::from("# Retrieved Context\n\n");
            let mut sources = Vec::new();

            for (idx, result) in results.iter().enumerate() {
                context.push_str(&format!(
                    "## Source {}: {} (relevance: {:.3})\n\n{}\n\n",
                    idx + 1,
                    result.filename,
                    1.0 - result.distance.min(1.0),
                    result.chunk_text
                ));

                sources.push(RagSource {
                    filename: result.filename.clone(),
                    text: result.chunk_text.clone(),
                    relevance: 1.0 - result.distance.min(1.0),
                });
            }

            Ok(HttpResponse::Ok().json(RagQueryResponse { context, sources }))
        }
        Err(e) => {
            warn!("Failed to execute RAG query: {}", e);
            Ok(HttpResponse::InternalServerError().json(json!({
                "error": format!("Failed to execute query: {}", e)
            })))
        }
    }
}

/// List all indexed documents
pub async fn rag_list_documents(
    rag_system: web::Data<Option<Arc<RagSystem>>>,
) -> Result<HttpResponse, Error> {
    debug!("Listing RAG documents");

    let Some(rag_system) = rag_system.as_ref() else {
        return Ok(HttpResponse::ServiceUnavailable().json(json!({
            "error": "RAG system is not enabled"
        })));
    };

    match rag_system.indexer.list_documents() {
        Ok(docs) => {
            let documents: Vec<DocumentSummary> = docs
                .into_iter()
                .map(|doc| DocumentSummary {
                    id: doc.id,
                    filename: doc.filename,
                    file_size: doc.file_size,
                    updated_at: doc.updated_at,
                })
                .collect();

            Ok(HttpResponse::Ok().json(DocumentListResponse { documents }))
        }
        Err(e) => {
            warn!("Failed to list documents: {}", e);
            Ok(HttpResponse::InternalServerError().json(json!({
                "error": format!("Failed to list documents: {}", e)
            })))
        }
    }
}

/// Delete a document from the RAG index
pub async fn rag_delete_document(
    path: web::Path<String>,
    rag_system: web::Data<Option<Arc<RagSystem>>>,
) -> Result<HttpResponse, Error> {
    let filename = path.into_inner();

    let Some(rag_system) = rag_system.as_ref() else {
        return Ok(HttpResponse::ServiceUnavailable().json(RagResponse {
            success: false,
            message: "RAG system is not enabled".to_string(),
        }));
    };

    match rag_system.indexer.remove_document(&filename) {
        Ok(_) => Ok(HttpResponse::Ok().json(RagResponse {
            success: true,
            message: format!("Document {} deleted successfully", filename),
        })),
        Err(e) => {
            warn!("Failed to delete document: {}", e);
            Ok(HttpResponse::InternalServerError().json(RagResponse {
                success: false,
                message: format!("Failed to delete document: {}", e),
            }))
        }
    }
}

/// Get RAG statistics
pub async fn rag_stats(
    rag_system: web::Data<Option<Arc<RagSystem>>>,
) -> Result<HttpResponse, Error> {
    debug!("Getting RAG statistics");

    let Some(rag_system) = rag_system.as_ref() else {
        return Ok(HttpResponse::ServiceUnavailable().json(json!({
            "error": "RAG system is not enabled"
        })));
    };

    match rag_system.indexer.get_stats() {
        Ok((doc_count, chunk_count, embedding_count)) => {
            let avg_chunks_per_doc = if doc_count > 0 {
                chunk_count as f64 / doc_count as f64
            } else {
                0.0
            };

            Ok(HttpResponse::Ok().json(RagStatsResponse {
                doc_count,
                chunk_count,
                embedding_count,
                avg_chunks_per_doc,
            }))
        }
        Err(e) => {
            warn!("Failed to get RAG stats: {}", e);
            Ok(HttpResponse::InternalServerError().json(json!({
                "error": format!("Failed to get statistics: {}", e)
            })))
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct UploadDocumentRequest {
    pub filename: String,
    pub content: String,
}

/// Upload and index a document
pub async fn rag_upload_document(
    body: web::Json<UploadDocumentRequest>,
    rag_system: web::Data<Option<Arc<RagSystem>>>,
    app_config: web::Data<Arc<config::Config>>,
) -> Result<HttpResponse, Error> {
    if rag_system.as_ref().is_none() {
        return Ok(HttpResponse::ServiceUnavailable().json(RagResponse {
            success: false,
            message: "RAG system is not enabled".to_string(),
        }));
    }

    use std::path::PathBuf;
    use tokio::fs;

    let documents_path = PathBuf::from(&app_config.rag.documents_path);

    if !documents_path.exists()
        && let Err(e) = fs::create_dir_all(&documents_path).await
    {
        return Ok(HttpResponse::InternalServerError().json(RagResponse {
            success: false,
            message: format!("Failed to create documents directory: {}", e),
        }));
    }

    let file_path = documents_path.join(&body.filename);

    if let Err(e) = fs::write(&file_path, &body.content).await {
        return Ok(HttpResponse::InternalServerError().json(RagResponse {
            success: false,
            message: format!("Failed to write file: {}", e),
        }));
    }

    // File will be automatically indexed by the document watcher
    Ok(HttpResponse::Ok().json(RagResponse {
        success: true,
        message: format!(
            "Document {} uploaded successfully. Indexing in progress...",
            body.filename
        ),
    }))
}
