use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use futures::future::join_all;
use futures::FutureExt;
use serde_json::json;
use tokio::sync::Notify;

use crate::agent_events::{AgentEvent, ToolCall, ToolDefinition, ToolResult};
use crate::agent_tools;
use crate::ai::{self, AiCompletionRequest, AiConfig, AiMessage, AiProvider, AiStreamChunk, AiTaskContract};
use crate::ai_cli_agent::CliAgentCommandSpec;
use crate::connection::AppState;
use crate::models::connection::DatabaseType;
use crate::token_usage::TokenUsage;

/// Default number of agent loop turns to prevent infinite loops.
/// Users can raise the limit in Settings → AI; it is clamped to
/// [`MIN_MAX_AGENT_TURNS`, `MAX_MAX_AGENT_TURNS`] so it can never be unlimited.
///
/// These values are mirrored in apps/desktop/src/components/editor/EditorSettingsDialog.vue
/// (MAX_AGENT_TURNS_DEFAULT/MIN/MAX) for client-side input validation — this module's
/// `clamp_max_agent_turns` remains the actual source of truth, applied on every save/load.
pub const DEFAULT_MAX_AGENT_TURNS: u32 = 30;
pub const MIN_MAX_AGENT_TURNS: u32 = 5;
pub const MAX_MAX_AGENT_TURNS: u32 = 500;

/// Clamp a user-provided agent turn limit into the supported range.
pub fn clamp_max_agent_turns(value: u32) -> u32 {
    value.clamp(MIN_MAX_AGENT_TURNS, MAX_MAX_AGENT_TURNS)
}
const MAX_TOOL_RESULT_CONTEXT_CHARS: usize = 12_000;
const TOOL_RESULT_HEAD_CHARS: usize = 4_000;
const TOOL_RESULT_TAIL_CHARS: usize = 4_000;
const TOOL_RESULT_SAMPLE_ITEMS: usize = 5;
const MAX_CONTRACT_REPAIR_ATTEMPTS: u32 = 2;

fn take_text(m: &std::sync::Mutex<String>) -> String {
    m.lock().unwrap_or_else(|e| e.into_inner()).clone()
}

/// Convert a streaming AI chunk into agent events for the frontend.
/// Pure function — no side effects, easily testable.
fn chunk_to_events(chunk: &AiStreamChunk) -> Vec<AgentEvent> {
    let mut events = Vec::new();
    if !chunk.delta.is_empty() {
        events.push(AgentEvent::TextDelta { delta: chunk.delta.clone() });
    }
    if let Some(ref reasoning) = chunk.reasoning_delta {
        events.push(AgentEvent::ReasoningDelta { delta: reasoning.clone() });
    }
    events
}

enum LoopExit {
    Completed,
    Cancelled,
    Interrupted(String),
    Exhausted,
}

impl LoopExit {
    fn should_break_turns(&self) -> bool {
        matches!(self, LoopExit::Cancelled | LoopExit::Interrupted(_))
    }
}

enum CompactResult {
    Skipped,
    Compacted,
    Cancelled,
}

/// Context for an agent loop run.
pub struct AgentLoopContext {
    pub state: Arc<AppState>,
    pub connection_id: String,
    pub database: String,
    /// Selected schema that scopes Agent metadata and SQL execution.
    pub schema: Option<String>,
    pub db_type: DatabaseType,
    pub cli_mcp_server_command: Option<CliAgentCommandSpec>,
    pub sql_permissions: agent_tools::AgentSqlPermissions,
    /// Turn limit for this run, already clamped by the settings layer.
    /// Callers that have no user setting should pass [`DEFAULT_MAX_AGENT_TURNS`].
    pub max_agent_turns: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FunctionCallingSupport {
    Supported,
    Unsupported,
    Unknown,
}

/// Resolve function calling support for the selected provider/model.
///
/// Native Ollama capabilities are model-specific. Missing metadata is kept as
/// unknown so older servers and compatible proxies can still attempt tools.
async fn provider_function_calling_support(config: &AiConfig) -> FunctionCallingSupport {
    if !matches!(config.provider, AiProvider::Ollama) {
        return FunctionCallingSupport::Supported;
    }

    match ai::ollama_selected_model_tool_support(config).await {
        Ok(Some(true)) => FunctionCallingSupport::Supported,
        Ok(Some(false)) => FunctionCallingSupport::Unsupported,
        Ok(None) => FunctionCallingSupport::Unknown,
        Err(error) => {
            log::debug!(
                "[agent][ollama] tool capability unavailable for model {}: {error}; attempting tool call",
                config.model
            );
            FunctionCallingSupport::Unknown
        }
    }
}

/// Run the agent loop: call LLM with tools, execute tool calls, feed results back, repeat.
///
/// The `on_event` callback receives streaming events for the frontend.
/// Returns the final accumulated assistant text.
///
/// If the selected model explicitly does not support function calling,
/// automatically degrades to a text-only completion with schema context
/// injected into the system prompt.
#[allow(clippy::too_many_arguments)]
pub async fn run_agent_loop(
    config: &AiConfig,
    system_prompt: &str,
    messages: &[AiMessage],
    agent_ctx: &AgentLoopContext,
    on_event: impl Fn(AgentEvent) + Send + Sync + Clone + 'static,
    cancelled: &Notify,
    max_tokens: Option<u32>,
    task_contract: Option<&AiTaskContract>,
    is_agent_mode: bool,
) -> Result<String, String> {
    let contract_system_prompt = augment_system_prompt_with_task_contract(system_prompt, task_contract, is_agent_mode);
    let system_prompt = contract_system_prompt.as_str();

    if crate::ai::is_cli_provider(&config.provider) {
        let connection_name = {
            let configs = agent_ctx.state.configs.read().await;
            configs
                .get(&agent_ctx.connection_id)
                .map(|config| config.name.clone())
                .unwrap_or_else(|| agent_ctx.connection_id.clone())
        };
        let options = crate::ai_cli_agent::CliAgentRunOptions {
            connection_id: agent_ctx.connection_id.clone(),
            connection_name,
            database: agent_ctx.database.clone(),
            schema: agent_ctx.schema.clone(),
            agent_mode: is_agent_mode,
            allow_writes: agent_ctx.sql_permissions.allow_writes,
            allow_dangerous: agent_ctx.sql_permissions.allow_dangerous,
            confirmed_write_sql: agent_ctx.sql_permissions.confirmed_write_sql.clone(),
            mcp_server_command: agent_ctx.cli_mcp_server_command.clone(),
        };
        if matches!(config.provider, AiProvider::ClaudeCodeCli) {
            let prompt = crate::ai_claude_code_cli::build_claude_code_prompt(
                system_prompt,
                messages,
                agent_ctx.sql_permissions.allow_writes,
            );
            return crate::ai_claude_code_cli::run_claude_code_agent(config, &prompt, options, cancelled, on_event)
                .await;
        }
        if matches!(config.provider, AiProvider::PiAgentCli) {
            let prompt = crate::ai_pi_agent_cli::build_pi_agent_prompt(
                system_prompt,
                messages,
                agent_ctx.sql_permissions.allow_writes,
            );
            return crate::ai_pi_agent_cli::run_pi_agent(config, &prompt, options, cancelled, on_event).await;
        }
        if matches!(config.provider, AiProvider::OpenCodeCli) {
            let prompt = crate::ai_opencode_cli::build_opencode_prompt(
                system_prompt,
                messages,
                agent_ctx.sql_permissions.allow_writes,
            );
            return crate::ai_opencode_cli::run_opencode_agent(config, &prompt, options, cancelled, on_event).await;
        }
        if matches!(config.provider, AiProvider::CursorCli) {
            let prompt = crate::ai_cursor_cli::build_cursor_prompt(
                system_prompt,
                messages,
                agent_ctx.sql_permissions.allow_writes,
            );
            return crate::ai_cursor_cli::run_cursor_agent(config, &prompt, options, cancelled, on_event).await;
        }
        if matches!(config.provider, AiProvider::GrokCli) {
            let prompt =
                crate::ai_grok_cli::build_grok_prompt(system_prompt, messages, agent_ctx.sql_permissions.allow_writes);
            return crate::ai_grok_cli::run_grok_agent(config, &prompt, options, cancelled, on_event).await;
        }
        let prompt =
            crate::ai_codex_cli::build_codex_prompt(system_prompt, messages, agent_ctx.sql_permissions.allow_writes);
        return crate::ai_codex_cli::run_codex_agent(config, &prompt, options, cancelled, on_event).await;
    }

    // Auto-degrade only models that explicitly advertise no tool support.
    // Unknown Ollama capabilities still get a chance to use tools so older
    // servers and compatible proxies are not disabled provider-wide.
    let function_calling_support = tokio::select! {
        support = provider_function_calling_support(config) => support,
        _ = cancelled.notified() => {
            let message = "Agent run was cancelled before producing output.".to_string();
            on_event(AgentEvent::TextDelta { delta: message.clone() });
            on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
            return Ok(message);
        }
    };
    if function_calling_support == FunctionCallingSupport::Unsupported {
        return run_agent_loop_text_only(
            config,
            system_prompt,
            messages,
            agent_ctx,
            on_event,
            cancelled,
            max_tokens,
            task_contract,
        )
        .await;
    }
    let tools = if is_agent_mode {
        agent_tools::all_tools(agent_ctx.db_type, agent_ctx.sql_permissions.clone())
    } else {
        agent_tools::read_only_tools(agent_ctx.db_type)
    };
    let task_contract = task_contract.cloned();
    let mut conversation_messages: Vec<AiMessage> = messages.to_vec();
    let mut final_text = String::new();
    let mut loop_exit = LoopExit::Exhausted;
    let mut total_usage = TokenUsage::default();
    let mut contract_repair_attempts = 0;

    let max_agent_turns = clamp_max_agent_turns(agent_ctx.max_agent_turns);

    for turn in 0..max_agent_turns {
        // Check for cancellation before each turn
        if cancelled.notified().now_or_never().is_some() {
            loop_exit = LoopExit::Cancelled;
            break;
        }

        // Check and maybe compact context
        if matches!(
            maybe_compact(
                config,
                system_prompt,
                &tools,
                &mut conversation_messages,
                max_tokens,
                &on_event,
                cancelled,
                false,
            )
            .await,
            CompactResult::Cancelled
        ) {
            loop_exit = LoopExit::Cancelled;
            break;
        }

        on_event(AgentEvent::TurnStart { turn });

        let mut stream_result: Option<(Vec<ToolCall>, Option<TokenUsage>, String)> = None;
        let mut last_stream_error: Option<String> = None;

        for attempt in 0..2 {
            // Build the LLM request with tools. Rebuild after retry compaction so the request
            // reflects the latest conversation_messages.
            let request = build_tool_request(
                config,
                system_prompt,
                &conversation_messages,
                &tools,
                max_tokens,
                task_contract.clone(),
            );

            // Stream the LLM response, collecting text and tool_calls.
            let accumulated_text = Arc::new(Mutex::new(String::new()));
            let emitted_any_chunk = Arc::new(AtomicBool::new(false));
            let session_id =
                if attempt == 0 { format!("agent-turn-{turn}") } else { format!("agent-turn-{turn}-retry") };

            let acc = accumulated_text.clone();
            let emitted = emitted_any_chunk.clone();
            let on_event2 = on_event.clone();
            let on_chunk = move |chunk: AiStreamChunk| {
                if !chunk.delta.is_empty() {
                    emitted.store(true, Ordering::Relaxed);
                    acc.lock().unwrap_or_else(|e| e.into_inner()).push_str(&chunk.delta);
                }
                if chunk.reasoning_delta.is_some() {
                    emitted.store(true, Ordering::Relaxed);
                }
                for event in chunk_to_events(&chunk) {
                    on_event2(event);
                }
            };

            match stream_with_tools(config, &request, &session_id, &tools, cancelled, on_chunk).await {
                Ok((tool_calls, usage)) => {
                    let accumulated_text = take_text(&accumulated_text);
                    stream_result = Some((tool_calls, usage, accumulated_text));
                    break;
                }
                Err(err)
                    if turn == 0
                        && matches!(config.provider, AiProvider::Ollama)
                        && is_tool_unsupported_error(&err)
                        && !emitted_any_chunk.load(Ordering::Relaxed) =>
                {
                    log::debug!(
                        "[agent][ollama] model {} rejected tools before producing output; falling back to text-only mode",
                        config.model
                    );
                    on_event(AgentEvent::TurnEnd { turn });
                    return run_agent_loop_text_only(
                        config,
                        system_prompt,
                        messages,
                        agent_ctx,
                        on_event,
                        cancelled,
                        max_tokens,
                        task_contract.as_ref(),
                    )
                    .await;
                }
                Err(err)
                    if attempt == 0 && is_context_length_error(&err) && !emitted_any_chunk.load(Ordering::Relaxed) =>
                {
                    last_stream_error = Some(err);
                    let compacted = maybe_compact(
                        config,
                        system_prompt,
                        &tools,
                        &mut conversation_messages,
                        max_tokens,
                        &on_event,
                        cancelled,
                        true,
                    )
                    .await;
                    match compacted {
                        CompactResult::Compacted => continue,
                        CompactResult::Cancelled => {
                            loop_exit = LoopExit::Cancelled;
                            break;
                        }
                        CompactResult::Skipped => {
                            final_text = take_text(&accumulated_text);
                            loop_exit =
                                LoopExit::Interrupted(last_stream_error.take().unwrap_or_else(|| {
                                    "LLM request failed after context compaction retry".to_string()
                                }));
                        }
                    }
                    break;
                }
                Err(err) if err == ai::AGENT_CANCELLED_ERROR => {
                    final_text = take_text(&accumulated_text);
                    loop_exit = LoopExit::Cancelled;
                    break;
                }
                Err(err) => {
                    final_text = take_text(&accumulated_text);
                    loop_exit = LoopExit::Interrupted(err);
                    break;
                }
            }
        }

        if loop_exit.should_break_turns() {
            break;
        }

        let Some((collected_tool_calls, turn_usage, accumulated_text)) = stream_result else {
            return Err(
                last_stream_error.unwrap_or_else(|| "LLM request failed after context compaction retry".to_string())
            );
        };

        if let Some(usage) = turn_usage {
            total_usage.add(&usage);
        }

        on_event(AgentEvent::TurnEnd { turn });

        // Add assistant message to conversation (including tool_use blocks)
        conversation_messages.push(AiMessage {
            role: "assistant".to_string(),
            content: accumulated_text.clone(),
            tool_call_id: None,
            tool_calls: collected_tool_calls
                .iter()
                .map(|tc| ai::ToolCallRef {
                    id: tc.id.clone(),
                    name: tc.name.clone(),
                    arguments: tc.arguments.clone(),
                    provider_payload: tc.provider_payload.clone(),
                })
                .collect(),
        });

        if collected_tool_calls.is_empty() {
            match validate_final_answer(task_contract.as_ref(), &accumulated_text) {
                FinalAnswerCheck::Satisfied => {
                    final_text = accumulated_text;
                    loop_exit = LoopExit::Completed;
                    break;
                }
                FinalAnswerCheck::NeedsRepair(reason) if contract_repair_attempts < MAX_CONTRACT_REPAIR_ATTEMPTS => {
                    contract_repair_attempts += 1;
                    conversation_messages.push(AiMessage {
                        role: "user".to_string(),
                        content: build_contract_repair_prompt(task_contract.as_ref(), is_agent_mode, &reason),
                        tool_call_id: None,
                        tool_calls: Vec::new(),
                    });
                    continue;
                }
                FinalAnswerCheck::NeedsRepair(reason) => {
                    let message = append_contract_failure_note(accumulated_text, &reason);
                    on_event(AgentEvent::TextDelta { delta: message.clone() });
                    final_text = message;
                    loop_exit = LoopExit::Completed;
                    break;
                }
            }
        }

        // Honor a cancellation that arrived after the stream finished but before we
        // run the requested tools; otherwise a long execute_query would keep running.
        if cancelled.notified().now_or_never().is_some() {
            final_text = accumulated_text;
            loop_exit = LoopExit::Cancelled;
            break;
        }

        // Execute each tool call
        // Emit all ToolCallStart events first
        for tc in &collected_tool_calls {
            on_event(AgentEvent::ToolCallStart {
                tool_call_id: tc.id.clone(),
                tool_name: tc.name.clone(),
                args: tc.arguments.clone(),
            });
        }

        // Execute tool calls: parallel for read tools, sequential for execute_query
        let state2 = Arc::clone(&agent_ctx.state);
        let conn2 = agent_ctx.connection_id.clone();
        let db2 = agent_ctx.database.clone();
        let schema2 = agent_ctx.schema.clone();
        let db_type = agent_ctx.db_type;
        let sql_permissions = agent_ctx.sql_permissions.clone();

        // Split by index into parallel and sequential groups using tool metadata
        let tool_parallel_map: std::collections::HashMap<&str, bool> =
            tools.iter().map(|t| (t.name, t.parallel_ok)).collect();
        let (parallel_indices, sequential_indices): (Vec<usize>, Vec<usize>) = (0..collected_tool_calls.len())
            .partition(|&i| *tool_parallel_map.get(collected_tool_calls[i].name.as_str()).unwrap_or(&false));

        let make_tc = |tc: &ToolCall| ToolCall {
            id: tc.id.clone(),
            name: tc.name.clone(),
            arguments: tc.arguments.clone(),
            provider_payload: tc.provider_payload.clone(),
        };

        // Run parallel group
        let parallel_futures: Vec<_> =
            parallel_indices
                .iter()
                .map(|&i| {
                    let tc = make_tc(&collected_tool_calls[i]);
                    let state = Arc::clone(&state2);
                    let conn = conn2.clone();
                    let db = db2.clone();
                    let schema = schema2.clone();
                    let perms = sql_permissions.clone();
                    async move {
                        agent_tools::execute_tool(&tc, &state, &conn, &db, schema.as_deref(), &db_type, perms).await
                    }
                })
                .collect();
        let parallel_results = join_all(parallel_futures).await;

        // Run sequential group one-by-one
        let mut sequential_results = Vec::with_capacity(sequential_indices.len());
        for &i in &sequential_indices {
            let tc = make_tc(&collected_tool_calls[i]);
            sequential_results.push(
                agent_tools::execute_tool(
                    &tc,
                    &state2,
                    &conn2,
                    &db2,
                    schema2.as_deref(),
                    &db_type,
                    sql_permissions.clone(),
                )
                .await,
            );
        }

        // Merge results back into original order
        let mut results: Vec<Option<ToolResult>> = vec![None; collected_tool_calls.len()];
        for (pos, &i) in parallel_indices.iter().enumerate() {
            results[i] = Some(parallel_results[pos].clone());
        }
        for (pos, &i) in sequential_indices.iter().enumerate() {
            results[i] = Some(sequential_results[pos].clone());
        }
        let results: Vec<ToolResult> = results.into_iter().map(|r| r.unwrap()).collect();

        // Process results in order, emitting ToolCallEnd events
        for (tc, result) in collected_tool_calls.iter().zip(results) {
            on_event(AgentEvent::ToolCallEnd {
                tool_call_id: tc.id.clone(),
                tool_name: tc.name.clone(),
                result: match &result.explain_data {
                    Some(ed) => json!({ "content": result.content, "explain_data": ed }),
                    None => json!({ "content": result.content }),
                },
                is_error: result.is_error,
            });
            conversation_messages.push(AiMessage {
                role: "tool".to_string(),
                content: tool_result_for_followup_context(&tc.name, &result.content),
                tool_call_id: Some(tc.id.clone()),
                tool_calls: Vec::new(),
            });
        }
    }

    match loop_exit {
        LoopExit::Completed => {}
        LoopExit::Cancelled => {
            let message = if final_text.trim().is_empty() {
                "Agent run was cancelled before producing output.".to_string()
            } else {
                "\n\nAgent run was cancelled. Partial output above was preserved.".to_string()
            };
            on_event(AgentEvent::TextDelta { delta: message.clone() });
            final_text.push_str(&message);
        }
        LoopExit::Interrupted(error) => {
            let message = if final_text.trim().is_empty() {
                format!("Agent stream stopped before completion: {error}.")
            } else {
                format!("\n\nAgent stream stopped before completion: {error}. Partial output above was preserved.")
            };
            on_event(AgentEvent::TextDelta { delta: message.clone() });
            final_text.push_str(&message);
        }
        LoopExit::Exhausted => {
            let message = if final_text.trim().is_empty() {
                format!("Agent reached the {max_agent_turns}-turn safety limit before producing output. Send Continue to let the agent keep working, or raise the limit in Settings → AI.")
            } else {
                format!(
                    "\n\nAgent reached the {max_agent_turns}-turn safety limit before a final answer. The partial output above was preserved; send Continue to let the agent keep working, or raise the limit in Settings → AI."
                )
            };
            on_event(AgentEvent::TextDelta { delta: message.clone() });
            final_text.push_str(&message);
        }
    }

    on_event(AgentEvent::AgentEnd {
        input_tokens: if total_usage.input_tokens > 0 { Some(total_usage.input_tokens) } else { None },
        output_tokens: if total_usage.output_tokens > 0 { Some(total_usage.output_tokens) } else { None },
    });
    Ok(final_text)
}

/// Build an LLM request that includes tool definitions.
fn build_tool_request(
    config: &AiConfig,
    system_prompt: &str,
    messages: &[AiMessage],
    _tools: &[ToolDefinition], // Tools are injected in ai::stream_with_tools, not via AiCompletionRequest.
    max_tokens: Option<u32>,
    task_contract: Option<AiTaskContract>,
) -> AiCompletionRequest {
    // Note: tools are passed via the body, not via AiCompletionRequest.
    // The actual injection happens in stream_with_tools.
    AiCompletionRequest {
        config: config.clone(),
        system_prompt: system_prompt.to_string(),
        messages: messages.to_vec(),
        task_contract,
        max_tokens: max_tokens.or(Some(4096)),
    }
}

fn augment_system_prompt_with_task_contract(
    system_prompt: &str,
    task_contract: Option<&AiTaskContract>,
    is_agent_mode: bool,
) -> String {
    let Some(contract) = task_contract else {
        return system_prompt.to_string();
    };

    let action = contract.action.as_deref().unwrap_or("unknown");
    let mode = contract.mode.as_deref().unwrap_or(if is_agent_mode { "agent" } else { "ask" });
    let user_request = contract.user_request.as_deref().unwrap_or("(not provided)");
    let mode_rule = if action_requires_sql_deliverable(action) {
        "This is a SQL-producing action: produce the final SQL in a fenced ```sql code block. Use tools only as intermediate evidence for schema/dialect; do not stop at a tool-result summary. In Agent mode, execute a query only when the original request explicitly asks for real data/results, not when it merely asks to generate SQL."
    } else {
        match action.to_ascii_lowercase().as_str() {
            "general" => "This is a general Q&A mode. Answer the user's question directly and naturally using your knowledge and any available database context. Adapt to the user's intent.",
            "query" => "This is a data-query task: call execute_query to obtain real results, then answer based on the actual data. Do not stop after merely outputting SQL text.",
            "exploreschema" => "This is a schema-inspection task: use list_tables/get_columns to obtain authoritative structure, then summarize. Do not execute data queries unless the user explicitly asks for data.",
            "executeandexplain" => "This is an execute-and-explain task: call execute_query to run the current SQL, then explain the real results.",
            _ if is_agent_mode => "For data-query intents, obtain real results with execute_query when safe; otherwise state the blocker.",
            _ => "In Ask mode, produce SQL/explanation only and do not claim execution.",
        }
    };

    format!(
        "{system_prompt}\n\n[TASK CONTRACT]\n\
Original user request: {user_request}\n\
Action: {action}\n\
Mode: {mode}\n\
Tool results are intermediate evidence. Continue the original task after every tool call; never treat a tool-result summary as the final answer unless the user explicitly requested that summary.\n\
{mode_rule}\n\
If the final deliverable cannot be produced safely, state the exact missing information and ask one concise clarification question."
    )
}

#[derive(Debug, PartialEq, Eq)]
enum FinalAnswerCheck {
    Satisfied,
    NeedsRepair(String),
}

fn validate_final_answer(task_contract: Option<&AiTaskContract>, text: &str) -> FinalAnswerCheck {
    let Some(contract) = task_contract else {
        return FinalAnswerCheck::Satisfied;
    };

    let action = contract.action.as_deref().unwrap_or_default();
    if action_requires_sql_deliverable(action)
        && !contains_sql_deliverable(text)
        && !looks_like_blocker_or_clarification(text)
    {
        return FinalAnswerCheck::NeedsRepair(
            "SQL-producing actions require a final SQL code block, or a concise blocker/clarification when SQL cannot be produced safely.".to_string(),
        );
    }

    FinalAnswerCheck::Satisfied
}

fn build_contract_repair_prompt(task_contract: Option<&AiTaskContract>, is_agent_mode: bool, reason: &str) -> String {
    let action = task_contract.and_then(|c| c.action.as_deref()).unwrap_or("unknown");
    let mode = task_contract.and_then(|c| c.mode.as_deref()).unwrap_or(if is_agent_mode { "agent" } else { "ask" });
    let user_request = task_contract.and_then(|c| c.user_request.as_deref()).unwrap_or("(not provided)");
    let mode_rule = if action_requires_sql_deliverable(action) {
        "For this SQL-producing action, produce SQL in a fenced ```sql code block. Tool results are evidence only; do not answer by summarizing schema/tool output. Execute a query only when the original request explicitly asks for real data/results."
    } else {
        match action.to_ascii_lowercase().as_str() {
            "general" => "For this general Q&A, answer the user's question directly and naturally.",
            "query" => "For this data-query task, call execute_query and answer based on real data; do not stop at SQL text or a schema summary.",
            "exploreschema" => "For this schema-inspection task, summarize real structure from list_tables/get_columns; do not invent columns.",
            "executeandexplain" => "For this execute-and-explain task, run the current SQL via execute_query and explain the real results.",
            _ if is_agent_mode => "If the original request asks for real data and it can be answered safely, call execute_query before the final answer.",
            _ => "In Ask mode, generate SQL and concise explanation only; do not claim the SQL was executed.",
        }
    };

    format!(
        "[SYSTEM-GENERATED TASK CONTRACT CHECK]\n\
Your previous response did not satisfy the current task contract.\n\
Issue: {reason}\n\
Original user request: {user_request}\n\
Action: {action}\n\
Mode: {mode}\n\n\
Tool results in the conversation are intermediate evidence only. Continue the original user task; do not summarize tool results unless the user explicitly requested a summary.\n\
{mode_rule}\n\
Produce a final answer that satisfies the action contract now. If required tables/columns are missing or ambiguous, state exactly what is missing and ask one concise clarification question."
    )
}

fn action_requires_sql_deliverable(action: &str) -> bool {
    matches!(action.to_ascii_lowercase().as_str(), "generate" | "optimize" | "fix" | "convert" | "sampledata")
}

fn append_contract_failure_note(text: String, reason: &str) -> String {
    if text.trim().is_empty() {
        return format!("Unable to produce a contract-compliant final answer: {reason}");
    }

    format!("{text}\n\nTask contract warning: {reason}")
}

fn contains_sql_deliverable(text: &str) -> bool {
    let lower = text.to_ascii_lowercase();
    let mut rest = lower.as_str();

    while let Some(fence_start) = rest.find("```") {
        let after_open = &rest[fence_start + 3..];
        let Some(info_end) = after_open.find('\n') else {
            return false;
        };
        let info = after_open[..info_end].trim();
        let after_info = &after_open[info_end + 1..];
        let Some(fence_end) = after_info.find("```") else {
            return false;
        };
        let body = &after_info[..fence_end];

        if (info.is_empty() || info.starts_with("sql")) && contains_sql_keyword(body) {
            return true;
        }

        rest = &after_info[fence_end + 3..];
    }

    false
}

fn contains_sql_keyword(lower_text: &str) -> bool {
    lower_text.split(|ch: char| !ch.is_ascii_alphanumeric() && ch != '_').any(|token| {
        matches!(
            token,
            "select" | "with" | "show" | "describe" | "explain" | "insert" | "update" | "delete" | "create" | "alter"
        )
    })
}

fn looks_like_blocker_or_clarification(text: &str) -> bool {
    let lower = text.to_ascii_lowercase();
    let markers = [
        "missing",
        "not enough",
        "cannot determine",
        "can't determine",
        "unable to determine",
        "which column",
        "please clarify",
        "need to know",
        "缺少",
        "不足",
        "无法确定",
        "不能确定",
        "没有找到",
        "未找到",
        "请确认",
        "请提供",
        "需要明确",
        "哪个字段",
    ];
    markers.iter().any(|marker| lower.contains(marker))
}

/// Stream an LLM response with tool support, parsing tool_calls from SSE deltas.
///
/// Reasoning and tool call arguments are emitted incrementally as they arrive.
/// Assistant text is buffered until it satisfies the task contract so an
/// intermediate tool-result summary is not shown as the final answer.
async fn stream_with_tools(
    config: &AiConfig,
    request: &AiCompletionRequest,
    session_id: &str,
    tools: &[ToolDefinition],
    cancelled: &Notify,
    on_chunk: impl Fn(AiStreamChunk) + Send + Sync + 'static,
) -> Result<(Vec<ToolCall>, Option<TokenUsage>), String> {
    // Return early if the user cancelled before the LLM call started.
    if cancelled.notified().now_or_never().is_some() {
        return Err(ai::AGENT_CANCELLED_ERROR.to_string());
    }

    ai::stream_with_tools(config, request, session_id, tools, cancelled, on_chunk).await
}

fn is_context_length_error(error: &str) -> bool {
    let lower = error.to_lowercase();
    [
        "context length",
        "context_length",
        "maximum context",
        "max context",
        "token limit",
        "too many tokens",
        "prompt is too long",
        "input is too long",
        "reduce the length",
    ]
    .iter()
    .any(|marker| lower.contains(marker))
}

fn is_tool_unsupported_error(error: &str) -> bool {
    let lower = error.to_lowercase();
    let mentions_tool_use = lower.contains("tool") || lower.contains("function call");
    let rejects_capability = [
        "does not support",
        "doesn't support",
        "not supported",
        "unsupported",
        "unknown field",
        "unknown parameter",
        "unrecognized field",
        "unrecognized parameter",
    ]
    .iter()
    .any(|marker| lower.contains(marker));
    mentions_tool_use && rejects_capability
}

/// Text-only fallback for providers that don't support function calling.
///
/// Injects database schema context into the system prompt so the LLM can still
/// give informed answers, then performs a single non-streaming completion.
#[allow(clippy::too_many_arguments)]
async fn run_agent_loop_text_only(
    config: &AiConfig,
    system_prompt: &str,
    messages: &[AiMessage],
    agent_ctx: &AgentLoopContext,
    on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    cancelled: &Notify,
    max_tokens: Option<u32>,
    task_contract: Option<&AiTaskContract>,
) -> Result<String, String> {
    // Build a schema-enriched system prompt so the LLM can answer schema questions
    // even without tool access.
    let enriched_prompt = build_schema_prompt(agent_ctx, system_prompt).await;

    // Honor a cancellation requested while loading schema context.
    if cancelled.notified().now_or_never().is_some() {
        on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
        return Ok("Agent run was cancelled before producing output.".to_string());
    }

    let mut request = AiCompletionRequest {
        config: config.clone(),
        system_prompt: enriched_prompt,
        messages: messages.to_vec(),
        task_contract: task_contract.cloned(),
        max_tokens: max_tokens.or(Some(4096)),
    };

    for attempt in 0..=MAX_CONTRACT_REPAIR_ATTEMPTS {
        // Use non-streaming completions so contract repair can suppress incomplete drafts.
        // Race the (non-cancellable) HTTP call against cancellation so Stop still works.
        let result = tokio::select! {
            result = ai::complete(&request) => result?,
            _ = cancelled.notified() => {
                on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
                return Ok("Agent run was cancelled before producing output.".to_string());
            }
        };
        match validate_final_answer(task_contract, &result) {
            FinalAnswerCheck::Satisfied => {
                on_event(AgentEvent::TextDelta { delta: result.clone() });
                on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
                return Ok(result);
            }
            FinalAnswerCheck::NeedsRepair(reason) if attempt < MAX_CONTRACT_REPAIR_ATTEMPTS => {
                request.messages.push(AiMessage {
                    role: "assistant".to_string(),
                    content: result,
                    tool_call_id: None,
                    tool_calls: Vec::new(),
                });
                request.messages.push(AiMessage {
                    role: "user".to_string(),
                    content: build_contract_repair_prompt(task_contract, false, &reason),
                    tool_call_id: None,
                    tool_calls: Vec::new(),
                });
            }
            FinalAnswerCheck::NeedsRepair(reason) => {
                let message = append_contract_failure_note(result, &reason);
                on_event(AgentEvent::TextDelta { delta: message.clone() });
                on_event(AgentEvent::AgentEnd { input_tokens: None, output_tokens: None });
                return Ok(message);
            }
        }
    }

    Err("Text-only agent fallback failed to produce a final answer".to_string())
}

/// Build a system prompt enriched with database schema information
/// for text-only mode where the LLM cannot use tools.
async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str) -> String {
    let mut enriched = system_prompt.to_string();

    // Fetch real schema data using the same core functions the tools would use
    let tables_result = crate::schema::list_tables_core(
        &agent_ctx.state,
        &agent_ctx.connection_id,
        &agent_ctx.database,
        agent_ctx.schema.as_deref().unwrap_or(""),
        None,
        Some(50), // smaller limit for prompt injection
        None,
        None,
        None,
    )
    .await;

    match tables_result {
        Ok(tables) if !tables.is_empty() => {
            enriched.push_str("\n\n## Database Schema (for context — no tools available)\n");
            enriched.push_str(&format!("Database: {}\n", agent_ctx.database));
            if let Some(schema) = agent_ctx.schema.as_deref() {
                enriched.push_str(&format!("Schema: {schema}\n"));
            }
            enriched.push_str("Tables:\n");
            for t in &tables {
                enriched.push_str(&format!("  - {} ({})", t.name, t.table_type));
                if let Some(ref comment) = t.comment {
                    if !comment.trim().is_empty() {
                        enriched.push_str(&format!(" — {}", comment.trim()));
                    }
                }
                enriched.push('\n');
            }
        }
        _ => {
            enriched.push_str("\n\n(Note: Unable to load database schema for this request.)\n");
        }
    }

    enriched
}

/// Estimate text tokens conservatively for mixed English, Chinese, SQL, and JSON content.
fn estimate_text_tokens(text: &str) -> u32 {
    let chars = text.chars().count() as u32;
    if chars == 0 {
        return 0;
    }

    let non_ascii = text.chars().filter(|c| !c.is_ascii()).count() as u32;
    let alpha = non_ascii as f32 / chars as f32;
    let ascii_est = (text.len() as f32 / 3.5).ceil();
    let nonascii_est = (chars as f32 * 1.2).ceil();
    let estimated = (ascii_est * (1.0 - alpha) + nonascii_est * alpha).ceil() as u32;
    estimated.max(1)
}

fn estimate_message_tokens(message: &AiMessage) -> u32 {
    let mut tokens = estimate_text_tokens(&message.content) + 4;

    if let Some(tool_call_id) = &message.tool_call_id {
        tokens += estimate_text_tokens(tool_call_id) + 2;
    }

    for tool_call in &message.tool_calls {
        tokens += estimate_text_tokens(&tool_call.id) + estimate_text_tokens(&tool_call.name) + 4;
        if let Ok(args) = serde_json::to_string(&tool_call.arguments) {
            tokens += estimate_text_tokens(&args);
        }
        if let Some(provider_payload) = &tool_call.provider_payload {
            if let Ok(payload) = serde_json::to_string(provider_payload) {
                tokens += estimate_text_tokens(&payload);
            }
        }
    }

    tokens
}

/// Estimate tokens for a slice of messages.
fn estimate_tokens(messages: &[AiMessage]) -> u32 {
    messages.iter().map(estimate_message_tokens).sum()
}

fn estimate_tool_schema_tokens(tools: &[ToolDefinition]) -> u32 {
    tools
        .iter()
        .map(|tool| {
            let schema_tokens =
                serde_json::to_string(&tool.parameters).map(|schema| estimate_text_tokens(&schema)).unwrap_or_default();
            estimate_text_tokens(tool.name) + estimate_text_tokens(tool.description) + schema_tokens + 16
        })
        .sum()
}

fn estimate_current_prompt_tokens(system_prompt: &str, tools: &[ToolDefinition], messages: &[AiMessage]) -> u32 {
    estimate_text_tokens(system_prompt) + estimate_tool_schema_tokens(tools) + estimate_tokens(messages) + 16
}

/// Returns the context window size for a given model name.
fn context_window_for_model(model: &str) -> u32 {
    let m = model.to_lowercase();
    // GPT-4.1 family: 1M context
    if m.contains("gpt-4.1") {
        return 1_000_000;
    }
    if m.contains("claude") || m.contains("o1") || m.starts_with("o3") || m.starts_with("o4") {
        200_000
    } else if m.contains("gpt-4") {
        128_000
    } else if m.contains("gemini") {
        1_000_000
    } else {
        128_000
    }
}

fn prompt_budget(window: u32, max_tokens: Option<u32>) -> u32 {
    let output_reserve = max_tokens.unwrap_or(4096).min(window / 2);
    let safety_reserve = (window / 10).clamp(2048, 16_384).min(window / 2);
    window.saturating_sub(output_reserve).saturating_sub(safety_reserve)
}

fn keep_recent_budget(prompt_budget: u32) -> u32 {
    if prompt_budget <= 4096 {
        prompt_budget / 2
    } else {
        (prompt_budget * 6 / 10).clamp(4096, 50_000)
    }
}

const COMPACT_SYSTEM_PROMPT: &str = "\
You are a conversation summarizer. Produce a concise structured summary of the conversation \
provided. Format:\n\
## Progress\n## Key Decisions\n## Critical Context\n## Next Steps\n\
Be factual. No commentary.";

async fn maybe_compact(
    config: &AiConfig,
    system_prompt: &str,
    tools: &[ToolDefinition],
    messages: &mut Vec<AiMessage>,
    max_tokens: Option<u32>,
    on_event: &(impl Fn(AgentEvent) + Send + Sync),
    cancelled: &Notify,
    force: bool,
) -> CompactResult {
    let window = config.context_window.unwrap_or_else(|| context_window_for_model(&config.model));
    let budget = prompt_budget(window, max_tokens);
    let estimated_before = estimate_current_prompt_tokens(system_prompt, tools, messages);

    if !force && estimated_before <= budget {
        return CompactResult::Skipped;
    }

    if messages.len() <= 2 {
        return CompactResult::Skipped;
    }

    // Find cut point: keep a dynamic budget of recent messages and summarize older context.
    let keep_recent_tokens = if force { keep_recent_budget(budget) / 2 } else { keep_recent_budget(budget) };
    let mut recent_tokens = 0u32;
    let mut cut = messages.len();
    for i in (0..messages.len()).rev() {
        let t = estimate_message_tokens(&messages[i]);
        if recent_tokens + t > keep_recent_tokens && i > 0 {
            cut = i + 1;
            break;
        }
        recent_tokens += t;
        cut = i;
    }

    if cut >= messages.len() {
        cut = messages.len().saturating_sub(1);
    }

    let cut = adjust_cut_for_tool_pair_integrity(messages, cut);

    // Always keep messages[0] (the original user question) verbatim outside the summary.
    // Only summarize messages[1..cut].
    if cut <= 1 {
        return CompactResult::Skipped;
    }
    let summary_start = 1usize;

    let compacted_messages = cut - summary_start;

    let convo_text: String = messages[summary_start..cut].iter().map(format_message_for_summary).collect();

    let summary_request = AiCompletionRequest {
        config: config.clone(),
        system_prompt: COMPACT_SYSTEM_PROMPT.to_string(),
        messages: vec![AiMessage {
            role: "user".to_string(),
            content: format!("<conversation>\n{convo_text}</conversation>\n\nSummarize the above."),
            tool_call_id: None,
            tool_calls: Vec::new(),
        }],
        task_contract: None,
        max_tokens: Some(1024),
    };

    let summary = match cancelled.notified().now_or_never() {
        Some(_) => return CompactResult::Cancelled,
        None => match tokio::select! {
            result = ai::complete(&summary_request) => result,
            _ = cancelled.notified() => return CompactResult::Cancelled,
        } {
            Ok(s) => s,
            Err(_) => fallback_summary(messages, cut),
        },
    };

    let summary = if validate_summary(&summary) { summary } else { fallback_summary(messages, cut) };

    let summary_tokens = estimate_text_tokens(&summary) + 4;

    let summary_content = format!(
        "[SYSTEM-GENERATED CONTEXT SUMMARY - earlier conversation compressed; background only, not a new user request]\n\n{summary}"
    );

    messages.drain(summary_start..cut);
    messages.insert(
        summary_start,
        AiMessage {
            role: "user".to_string(),
            content: summary_content.clone(),
            tool_call_id: None,
            tool_calls: Vec::new(),
        },
    );

    let estimated_after = estimate_current_prompt_tokens(system_prompt, tools, messages);

    on_event(AgentEvent::ContextCompacted {
        summary: summary_content,
        summary_tokens,
        compacted_messages,
        estimated_before,
        estimated_after,
    });
    CompactResult::Compacted
}

fn tool_result_for_followup_context(tool_name: &str, content: &str) -> String {
    let result = compact_tool_result_for_context(tool_name, content);
    format!(
        "[TOOL RESULT - INTERMEDIATE EVIDENCE]\n\
Tool: {tool_name}\n\
Use this result to continue the original user task. Do not summarize this tool result as the final answer unless the user explicitly asked for a tool-result or schema summary.\n\n\
{result}"
    )
}

fn compact_tool_result_for_context(tool_name: &str, content: &str) -> String {
    if content.chars().count() <= MAX_TOOL_RESULT_CONTEXT_CHARS {
        return content.to_string();
    }

    if let Ok(value) = serde_json::from_str::<serde_json::Value>(content) {
        return compact_json_tool_result(tool_name, content, &value);
    }

    compact_text_tool_result(tool_name, content)
}

fn compact_json_tool_result(tool_name: &str, original: &str, value: &serde_json::Value) -> String {
    let compacted = match value {
        serde_json::Value::Array(items) => json!({
            "type": "array",
            "totalItems": items.len(),
            "head": items.iter().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>(),
            "tail": items.iter().rev().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>().into_iter().rev().collect::<Vec<_>>(),
        }),
        serde_json::Value::Object(map) => {
            let mut object = serde_json::Map::new();
            object.insert("type".to_string(), json!("object"));
            object.insert("keys".to_string(), json!(map.keys().cloned().collect::<Vec<_>>()));
            for (key, field_value) in map {
                match field_value {
                    serde_json::Value::Array(items) if items.len() > TOOL_RESULT_SAMPLE_ITEMS * 2 => {
                        object.insert(
                            key.clone(),
                            json!({
                                "totalItems": items.len(),
                                "head": items.iter().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>(),
                                "tail": items.iter().rev().take(TOOL_RESULT_SAMPLE_ITEMS).collect::<Vec<_>>().into_iter().rev().collect::<Vec<_>>(),
                            }),
                        );
                    }
                    _ => {
                        object.insert(key.clone(), field_value.clone());
                    }
                }
            }
            serde_json::Value::Object(object)
        }
        _ => value.clone(),
    };

    let compacted_text = serde_json::to_string_pretty(&compacted).unwrap_or_else(|_| compacted.to_string());
    if compacted_text.chars().count() <= MAX_TOOL_RESULT_CONTEXT_CHARS {
        format!(
            "[TOOL RESULT COMPACTED FOR CONTEXT]\nTool: {tool_name}\nOriginal chars: {}\nCompaction: parsed JSON with sampled arrays/fields. UI events preserve the full result.\n\n{}",
            original.chars().count(),
            compacted_text
        )
    } else {
        compact_text_tool_result(tool_name, original)
    }
}

fn compact_text_tool_result(tool_name: &str, content: &str) -> String {
    let original_chars = content.chars().count();
    let head = content.chars().take(TOOL_RESULT_HEAD_CHARS).collect::<String>();
    let tail_chars = content.chars().rev().take(TOOL_RESULT_TAIL_CHARS).collect::<Vec<_>>();
    let tail = tail_chars.into_iter().rev().collect::<String>();
    format!(
        "[TOOL RESULT COMPACTED FOR CONTEXT]\nTool: {tool_name}\nOriginal chars: {original_chars}\nCompaction: kept the head and tail; middle omitted. UI events preserve the full result.\n\n{head}\n\n...[middle omitted from tool result context]...\n\n{tail}"
    )
}

fn validate_summary(summary: &str) -> bool {
    let trimmed = summary.trim();
    trimmed.len() >= 50 && trimmed.len() <= 6000
}

fn fallback_summary(messages: &[AiMessage], cut: usize) -> String {
    // summary_start=1: messages[0] is kept verbatim, only messages[1..cut] are compacted.
    let compacted_messages = &messages[1..cut];
    let tool_calls = compacted_messages.iter().filter(|m| !m.tool_calls.is_empty()).count();
    let tool_results = compacted_messages.iter().filter(|m| m.role == "tool").count();
    let user_messages = compacted_messages.iter().filter(|m| m.role == "user").count();
    let assistant_messages = compacted_messages.iter().filter(|m| m.role == "assistant").count();

    let recent_roles = compacted_messages
        .iter()
        .rev()
        .take(8)
        .map(|message| format!("{}{}", message.role, if message.content.is_empty() { "" } else { ": content" }))
        .collect::<Vec<_>>()
        .join(", ");

    [
        "## Progress".to_string(),
        "- Context summarized by fallback generator because the LLM summary was unavailable or low quality.".to_string(),
        "## Key Decisions".to_string(),
        format!(
            "- Compacted {} messages: {} user, {} assistant, {} tool results, {} assistant tool-call messages.",
            compacted_messages.len(), user_messages, assistant_messages, tool_results, tool_calls
        ),
        "## Critical Context".to_string(),
        format!("- Recent compacted roles: {recent_roles}"),
        "## Next Steps".to_string(),
        "- Continue from the remaining recent conversation and recover any missing detail from tool handles or source paths if needed.".to_string(),
    ]
    .join("\n")
}

fn adjust_cut_for_tool_pair_integrity(messages: &[AiMessage], mut cut: usize) -> usize {
    if cut >= messages.len() {
        return cut;
    }

    while cut < messages.len() && messages[cut].role == "tool" {
        let Some(origin) = find_originating_assistant(messages, cut) else {
            break;
        };
        if origin >= cut {
            break;
        }
        cut = origin;
    }

    cut
}

fn find_originating_assistant(messages: &[AiMessage], tool_index: usize) -> Option<usize> {
    let tool_call_id = messages.get(tool_index)?.tool_call_id.as_deref();

    for i in (0..tool_index).rev() {
        let message = &messages[i];
        if message.role != "assistant" {
            continue;
        }

        if let Some(tool_call_id) = tool_call_id {
            if message.tool_calls.iter().any(|tool_call| tool_call.id == tool_call_id) {
                return Some(i);
            }
        } else if !message.tool_calls.is_empty() {
            return Some(i);
        }
    }

    None
}

fn format_message_for_summary(message: &AiMessage) -> String {
    let mut header = format!("[{}", message.role);
    if let Some(tool_call_id) = &message.tool_call_id {
        header.push_str(&format!(" tool_call_id={tool_call_id}"));
    }
    if !message.tool_calls.is_empty() {
        let tool_names =
            message.tool_calls.iter().map(|tool_call| tool_call.name.as_str()).collect::<Vec<_>>().join(", ");
        header.push_str(&format!(" tool_calls={tool_names}"));
    }
    header.push(']');

    format!("{header}: {}\n", summarize_message_content(&message.content))
}

fn summarize_message_content(content: &str) -> String {
    let char_count = content.chars().count();
    if char_count <= 4000 {
        return content.to_string();
    }

    let head = content.chars().take(1500).collect::<String>();
    let tail_chars = content.chars().rev().take(1500).collect::<Vec<_>>();
    let tail = tail_chars.into_iter().rev().collect::<String>();
    format!("{head}\n\n...[middle omitted for summary input]...\n\n{tail}")
}

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

    #[test]
    fn clamp_max_agent_turns_enforces_bounds() {
        assert_eq!(clamp_max_agent_turns(0), MIN_MAX_AGENT_TURNS);
        assert_eq!(clamp_max_agent_turns(MIN_MAX_AGENT_TURNS), MIN_MAX_AGENT_TURNS);
        assert_eq!(clamp_max_agent_turns(DEFAULT_MAX_AGENT_TURNS), DEFAULT_MAX_AGENT_TURNS);
        assert_eq!(clamp_max_agent_turns(200), 200);
        assert_eq!(clamp_max_agent_turns(u32::MAX), MAX_MAX_AGENT_TURNS);
    }

    #[test]
    fn identifies_explicit_tool_rejection_errors() {
        for error in [
            "model does not support tools",
            "tool use is not supported by this model",
            "unsupported parameter: tools",
            "unknown field `tools`",
            "function calling is not supported",
        ] {
            assert!(is_tool_unsupported_error(error), "{error}");
        }
    }

    #[test]
    fn does_not_hide_unrelated_ollama_errors_as_tool_rejections() {
        for error in [
            "connection refused",
            "model not found",
            "context length exceeded",
            "invalid tool arguments returned by model",
            "request timed out",
        ] {
            assert!(!is_tool_unsupported_error(error), "{error}");
        }
    }

    fn generate_contract(user_request: &str, mode: &str) -> AiTaskContract {
        AiTaskContract {
            action: Some("generate".to_string()),
            mode: Some(mode.to_string()),
            user_request: Some(user_request.to_string()),
        }
    }

    #[test]
    fn generate_contract_rejects_schema_summary_without_sql() {
        let contract = generate_contract("帮我生成统计 2026年1月2日新注册会员数量的 sql", "ask");
        let answer = "The tb_customer table contains comprehensive customer information with key columns:\n\nCore Identity\n- c_no: customer id\nContact Information\n- c_tele: mobile";

        let check = validate_final_answer(Some(&contract), answer);

        assert!(matches!(check, FinalAnswerCheck::NeedsRepair(_)));
    }

    #[test]
    fn generate_contract_accepts_sql_code_block() {
        let contract = generate_contract("帮我生成统计新注册会员数量的 sql", "ask");
        let answer = "```sql\nSELECT COUNT(*) AS member_count FROM tb_customer WHERE created_at >= '2026-01-02' AND created_at < '2026-01-03';\n```";

        let check = validate_final_answer(Some(&contract), answer);

        assert_eq!(check, FinalAnswerCheck::Satisfied);
    }

    #[test]
    fn generate_contract_rejects_unfenced_sql_mention() {
        let contract = generate_contract("帮我生成统计新注册会员数量的 sql", "ask");
        let answer = "You can use SQL to query the table, for example SELECT COUNT(*) FROM tb_customer.";

        let check = validate_final_answer(Some(&contract), answer);

        assert!(matches!(check, FinalAnswerCheck::NeedsRepair(_)));
    }

    #[test]
    fn generate_contract_accepts_missing_column_blocker() {
        let contract = generate_contract("帮我生成统计新注册会员数量的 sql", "ask");
        let answer = "没有找到明确表示会员注册时间的字段，请确认应该使用哪个字段作为注册时间。";

        let check = validate_final_answer(Some(&contract), answer);

        assert_eq!(check, FinalAnswerCheck::Satisfied);
    }

    #[test]
    fn agent_generate_sql_accepts_sql_without_execute_query() {
        let contract = generate_contract("统计 2026年1月2日新注册会员数量", "agent");
        let answer = "```sql\nSELECT COUNT(*) FROM tb_customer;\n```";

        let check = validate_final_answer(Some(&contract), answer);

        assert_eq!(check, FinalAnswerCheck::Satisfied);
    }

    #[test]
    fn agent_generate_sql_prompt_does_not_force_execution() {
        let contract = generate_contract("帮我生成统计 2026年1月2日新注册会员数量的 sql", "agent");

        let prompt = augment_system_prompt_with_task_contract("base", Some(&contract), true);

        assert!(prompt.contains("SQL-producing action"));
        assert!(prompt.contains("execute a query only when the original request explicitly asks for real data/results"));
    }

    #[test]
    fn wraps_tool_results_as_intermediate_evidence() {
        let wrapped = tool_result_for_followup_context("get_columns", "Columns of tb_customer:\n  - c_no: VARCHAR");

        assert!(wrapped.contains("INTERMEDIATE EVIDENCE"));
        assert!(wrapped.contains("continue the original user task"));
        assert!(wrapped.contains("Columns of tb_customer"));
    }

    // --- chunk_to_events tests ---

    #[test]
    fn chunk_to_events_emits_text_delta_for_text() {
        let chunk = AiStreamChunk {
            session_id: "test".to_string(),
            delta: "hello".to_string(),
            reasoning_delta: None,
            done: false,
        };
        let events = chunk_to_events(&chunk);
        assert_eq!(events.len(), 1);
        assert!(matches!(&events[0], AgentEvent::TextDelta { delta } if delta == "hello"));
    }

    #[test]
    fn chunk_to_events_emits_reasoning_delta_for_reasoning() {
        let chunk = AiStreamChunk {
            session_id: "test".to_string(),
            delta: String::new(),
            reasoning_delta: Some("thinking...".to_string()),
            done: false,
        };
        let events = chunk_to_events(&chunk);
        assert_eq!(events.len(), 1);
        assert!(matches!(&events[0], AgentEvent::ReasoningDelta { delta } if delta == "thinking..."));
    }

    #[test]
    fn chunk_to_events_emits_both_for_mixed_chunk() {
        let chunk = AiStreamChunk {
            session_id: "test".to_string(),
            delta: "answer".to_string(),
            reasoning_delta: Some("thinking...".to_string()),
            done: false,
        };
        let events = chunk_to_events(&chunk);
        assert_eq!(events.len(), 2);
        assert!(matches!(&events[0], AgentEvent::TextDelta { delta } if delta == "answer"));
        assert!(matches!(&events[1], AgentEvent::ReasoningDelta { delta } if delta == "thinking..."));
    }

    #[test]
    fn chunk_to_events_returns_empty_for_empty_chunk() {
        let chunk =
            AiStreamChunk { session_id: "test".to_string(), delta: String::new(), reasoning_delta: None, done: false };
        let events = chunk_to_events(&chunk);
        assert!(events.is_empty());
    }

    #[test]
    fn chunk_to_events_reasoning_only_no_text() {
        let chunk = AiStreamChunk {
            session_id: "test".to_string(),
            delta: String::new(),
            reasoning_delta: Some("reasoning".to_string()),
            done: false,
        };
        let events = chunk_to_events(&chunk);
        assert_eq!(events.len(), 1);
        assert!(matches!(&events[0], AgentEvent::ReasoningDelta { .. }));
    }

    #[test]
    fn chunk_to_events_text_only_no_reasoning() {
        let chunk = AiStreamChunk {
            session_id: "test".to_string(),
            delta: "text only".to_string(),
            reasoning_delta: None,
            done: false,
        };
        let events = chunk_to_events(&chunk);
        assert_eq!(events.len(), 1);
        assert!(matches!(&events[0], AgentEvent::TextDelta { .. }));
    }

    // --- Agent task-action contract tests (query / exploreSchema / executeAndExplain) ---

    fn contract_for(action: &str, user_request: &str, mode: &str) -> AiTaskContract {
        AiTaskContract {
            action: Some(action.to_string()),
            mode: Some(mode.to_string()),
            user_request: Some(user_request.to_string()),
        }
    }

    #[test]
    fn query_action_contract_requires_execute_query() {
        let contract = contract_for("query", "统计今天订单数", "agent");
        let prompt = augment_system_prompt_with_task_contract("base", Some(&contract), true);

        assert!(prompt.contains("data-query task"), "prompt should mark this as a data-query task");
        assert!(prompt.contains("call execute_query"), "prompt should instruct the LLM to call execute_query");
        assert!(!prompt.contains("SQL-producing action"), "query must not be treated as a SQL-producing action");
    }

    #[test]
    fn explore_schema_contract_uses_metadata_tools_not_execute_query() {
        let contract = contract_for("exploreSchema", "看一下 orders 表的结构", "agent");
        let prompt = augment_system_prompt_with_task_contract("base", Some(&contract), true);

        assert!(prompt.contains("schema-inspection task"));
        assert!(prompt.contains("list_tables/get_columns"));
        assert!(prompt.contains("Do not execute data queries"));
    }

    #[test]
    fn execute_and_explain_contract_runs_current_sql() {
        let contract = contract_for("executeAndExplain", "执行并解释当前 SQL", "agent");
        let prompt = augment_system_prompt_with_task_contract("base", Some(&contract), true);

        assert!(prompt.contains("execute-and-explain task"));
        assert!(prompt.contains("run the current SQL"));
    }

    #[test]
    fn task_actions_do_not_require_sql_deliverable() {
        // query / exploreSchema / executeAndExplain are task-oriented, not SQL-producing:
        // a final answer without a fenced SQL block must still satisfy the contract.
        let answer_without_sql = "今天共有 42 笔订单。";
        for action in ["query", "exploreSchema", "executeAndExplain"] {
            let contract = contract_for(action, "统计今天订单数", "agent");
            assert_eq!(
                validate_final_answer(Some(&contract), answer_without_sql),
                FinalAnswerCheck::Satisfied,
                "action {action} should not require a SQL deliverable",
            );
        }
    }

    #[test]
    fn query_action_repair_prompt_targets_execute_query() {
        let contract = contract_for("query", "统计今天订单数", "agent");
        let repair = build_contract_repair_prompt(Some(&contract), true, "previous answer did not execute");

        assert!(repair.contains("data-query task"));
        assert!(repair.contains("call execute_query"));
    }

    #[test]
    fn general_action_skips_sql_validation() {
        let contract = AiTaskContract {
            action: Some("general".to_string()),
            mode: Some("ask".to_string()),
            user_request: Some("你好".to_string()),
        };
        let answer = "你好！我是 DBX 的数据库助手。有什么可以帮你的吗？";
        assert_eq!(validate_final_answer(Some(&contract), answer), FinalAnswerCheck::Satisfied);
    }
}