use crate::agent::{
    assess_evidence, build_query_plan, fallback_evidence_assessment, fallback_query_plan,
    fallback_support_check, verify_answer_support, AgentIteration, EvidenceAssessment,
    EvidenceVerdict, QueryPlan, RetrievalStrategy,
};
use crate::citation::labeled_contexts;
use crate::cli::QueryModeArg;
use crate::commands::query::QueryCommand;
use crate::graph::placeholder_plan;
use crate::models::Generator;
use crate::retrieval::{merge_candidates, prune_candidates, RetrievalCandidate};
use crate::rewrite::{rewrite_query_for_retrieval, QueryRewriteSet};
use crate::store;
use crate::ui::Panel;
use anyhow::Result;
use std::collections::BTreeSet;
use tracing::{field, Instrument};

use super::render::{
    evidence_verdict_label, print_citations, print_contexts, print_query_plan, print_query_trace,
    print_scores,
};
use super::retrieve::retrieve_candidates;
use super::runtime::{mode_label, prepare_runtime};
use super::types::{QueryExecutionPath, QueryJsonReport, QueryResult, QueryRuntime};

pub async fn run(name: Option<&str>, command: QueryCommand) -> Result<()> {
    let runtime = prepare_runtime(name, &command)?;
    let execution = execution_path(command.mode);
    let span = tracing::info_span!(
        "query_pipeline",
        store_name = name.unwrap_or("default"),
        mode = ?command.mode,
        execution_path = execution_path_label(execution),
        question_chars = command.question.chars().count(),
        top_k = command.top_k,
        fetch_k = command.fetch_k,
        max_iterations = command.max_iterations,
        rewrite = command.rewrite,
        rerank = command.rerank,
        has_source_filter = command.source.is_some(),
        has_path_prefix_filter = command.path_prefix.is_some(),
        page = field::debug(command.page),
        has_format_filter = command.format.is_some(),
        hit_count = field::Empty,
        iteration_count = field::Empty,
    );

    let span_inner = span.clone();
    async move {
        let result = match execution {
            QueryExecutionPath::Simple => run_simple_query(&runtime, &command).await?,
            QueryExecutionPath::Agentic => run_agentic_query_command(&runtime, &command).await?,
            QueryExecutionPath::AgenticStub => {
                run_agentic_stub_query_command(&runtime, &command).await?
            }
        };

        span_inner.record("hit_count", result.hits.len());
        span_inner.record("iteration_count", result.iterations.len());

        if result.hits.is_empty() {
            if command.json {
                let report = QueryJsonReport::from_result(
                    &command.question,
                    mode_label(command.mode),
                    &result,
                    None,
                );
                println!("{}", serde_json::to_string_pretty(&report)?);
            } else {
                let mut panel = Panel::new("Query Result");
                panel.kv("status", "No relevant context found in the local store.", 8);
                panel.render();
            }
            return Ok(());
        }

        if command.json {
            let answer = stripped_answer(&runtime, &command, &result).await?;
            let report = QueryJsonReport::from_result(
                &command.question,
                mode_label(command.mode),
                &result,
                Some(answer),
            );
            println!("{}", serde_json::to_string_pretty(&report)?);
            return Ok(());
        }

        print_query_plan(&command, &result);
        print_query_trace(&command, &result);
        print_scores(&command, &result);
        print_citations(&command, &result);
        print_contexts(&command, &result);

        let answer = stripped_answer(&runtime, &command, &result).await?;
        println!();
        let mut panel = Panel::new("Answer");
        panel.prose("", answer.trim(), 0);
        panel.render();
        Ok(())
    }
    .instrument(span)
    .await
}

async fn run_simple_query(runtime: &QueryRuntime, command: &QueryCommand) -> Result<QueryResult> {
    let span = tracing::info_span!("run_simple_query", mode = ?command.mode);
    async move {
        let mut trace = vec![format!(
            "mode {} uses the current hybrid retrieval pipeline",
            mode_label(command.mode)
        )];
        let rewrite_set = build_rewrite_set(runtime, command, &mut trace).await;
        let hits = retrieve_candidates(runtime, command, &rewrite_set.query_variants(), &mut trace)
            .await?;

        Ok(QueryResult {
            requested_mode: command.mode,
            execution_label: mode_label(command.mode),
            rewrite_set,
            plan: None,
            iterations: Vec::new(),
            support_check: None,
            answer: None,
            hits,
            trace,
        })
    }
    .instrument(span)
    .await
}

async fn run_agentic_query_command(
    runtime: &QueryRuntime,
    command: &QueryCommand,
) -> Result<QueryResult> {
    let span = tracing::info_span!(
        "run_agentic_query",
        mode = ?command.mode,
        max_iterations = command.max_iterations,
        rewrite = command.rewrite,
        rerank = command.rerank,
    );
    async move {
        let generator = Generator::new(
            runtime.cfg.ollama.base_url.clone(),
            runtime.gen_model_name.clone(),
        );
        let mut trace = vec![format!(
            "mode {} uses the agentic retrieval loop",
            mode_label(command.mode)
        )];
        let plan = match build_query_plan(&generator, &command.question).await {
            Ok(plan) => {
                trace.push("planner produced a structured query plan".to_string());
                plan
            }
            Err(err) => {
                trace.push(format!(
                    "planner failed; falling back to heuristic plan ({})",
                    err.root_cause()
                ));
                fallback_query_plan(&command.question)
            }
        };

        let rewrite_set = build_rewrite_set(runtime, command, &mut trace).await;
        let mut iterations = Vec::new();
        let mut previous_keys = Vec::new();
        let mut active_queries = initial_agent_queries(&command.question, &plan, &rewrite_set);
        let mut accumulated_hits = Vec::new();

        for iteration in 1..=command.max_iterations.max(1) {
            trace.push(format!(
                "iteration {iteration}: querying {} variant(s)",
                active_queries.len()
            ));
            let iteration_hits =
                retrieve_candidates(runtime, command, &active_queries, &mut trace).await?;
            let retrieved_count = iteration_hits.len();
            accumulated_hits = merge_candidates([accumulated_hits, iteration_hits]);
            let kept_hits = prune_candidates(accumulated_hits.clone(), command.top_k);
            let assessment = match assess_evidence(&generator, &command.question, &kept_hits).await
            {
                Ok(assessment) => assessment,
                Err(err) => {
                    trace.push(format!(
                        "evidence assessment failed; using heuristic fallback ({})",
                        err.root_cause()
                    ));
                    fallback_evidence_assessment(&kept_hits)
                }
            };
            let notes = assessment_notes(&assessment);
            iterations.push(AgentIteration {
                iteration,
                query_variants: active_queries.clone(),
                retrieved_count,
                kept_count: kept_hits.len(),
                sufficiency: assessment.verdict.clone(),
                notes: notes.clone(),
            });
            for note in notes {
                trace.push(format!("iteration {iteration}: {note}"));
            }

            if matches!(assessment.verdict, EvidenceVerdict::Sufficient) {
                trace.push(format!("iteration {iteration}: evidence marked sufficient"));
                break;
            }
            if iteration >= command.max_iterations.max(1) {
                trace.push("iteration budget exhausted".to_string());
                break;
            }

            let current_keys = kept_hits
                .iter()
                .map(|hit| hit.dedupe_key())
                .collect::<Vec<_>>();
            if !previous_keys.is_empty() && current_keys == previous_keys {
                trace.push("retrieval produced no material evidence change; halting".to_string());
                break;
            }
            previous_keys = current_keys;

            let next_queries =
                refine_agent_queries(&command.question, &plan, &rewrite_set, &assessment);
            if next_queries == active_queries {
                trace.push("no better follow-up queries available; halting".to_string());
                break;
            }
            active_queries = next_queries;
        }

        let final_hits = prune_candidates(accumulated_hits, command.top_k);
        if final_hits.is_empty() {
            trace.push("no evidence retained; skipping answer generation".to_string());
            return Ok(QueryResult {
                requested_mode: command.mode,
                execution_label: "agentic",
                rewrite_set,
                plan: Some(plan),
                iterations,
                support_check: None,
                answer: None,
                hits: final_hits,
                trace,
            });
        }

        let answer = generate_answer_from_hits(runtime, command, &final_hits).await?;
        let support_check = match verify_answer_support(
            &generator,
            &command.question,
            &answer,
            &final_hits,
        )
        .await
        {
            Ok(check) => check,
            Err(err) => {
                trace.push(format!(
                    "support verification failed; using heuristic fallback ({})",
                    err.root_cause()
                ));
                fallback_support_check(&answer, &final_hits)
            }
        };
        if support_check.supported {
            trace.push("answer support verification passed".to_string());
        } else {
            trace.push(format!(
                "answer support verification found {} unsupported claim(s)",
                support_check.unsupported_claims.len()
            ));
        }

        Ok(QueryResult {
            requested_mode: command.mode,
            execution_label: "agentic",
            rewrite_set,
            plan: Some(plan),
            iterations,
            support_check: Some(support_check),
            answer: Some(answer),
            hits: final_hits,
            trace,
        })
    }
    .instrument(span)
    .await
}

async fn run_agentic_stub_query_command(
    runtime: &QueryRuntime,
    command: &QueryCommand,
) -> Result<QueryResult> {
    let span = tracing::info_span!("run_agentic_stub_query", mode = ?command.mode);
    async move {
        let mut trace = vec![format!(
            "requested {} mode routed through the graph-mode placeholder path",
            mode_label(command.mode)
        )];
        let rewrite_set = build_rewrite_set(runtime, command, &mut trace).await;
        let stub_plan = placeholder_plan(command.mode, &rewrite_set);
        trace.extend(stub_plan.notes.iter().cloned());
        let hits =
            retrieve_candidates(runtime, command, &stub_plan.query_variants, &mut trace).await?;

        Ok(QueryResult {
            requested_mode: command.mode,
            execution_label: stub_plan.execution_label,
            rewrite_set,
            plan: None,
            iterations: Vec::new(),
            support_check: None,
            answer: None,
            hits,
            trace,
        })
    }
    .instrument(span)
    .await
}

async fn build_rewrite_set(
    runtime: &QueryRuntime,
    command: &QueryCommand,
    trace: &mut Vec<String>,
) -> QueryRewriteSet {
    let span = tracing::info_span!(
        "build_rewrite_set",
        enabled = command.rewrite,
        question_chars = command.question.chars().count(),
    );
    async move {
        if !command.rewrite {
            trace.push("rewrite disabled; using original query only".to_string());
            return QueryRewriteSet::fallback(&command.question);
        }

        let generator = Generator::new(
            runtime.cfg.ollama.base_url.clone(),
            runtime.gen_model_name.clone(),
        );
        match rewrite_query_for_retrieval(&generator, &command.question).await {
            Ok(rewrite_set) => {
                let variants = rewrite_set.query_variants();
                trace.push(format!(
                    "rewrite enabled; {} retrieval query variant(s) prepared",
                    variants.len()
                ));
                for variant in variants.iter().skip(1) {
                    trace.push(format!("rewrite variant: {variant}"));
                }
                rewrite_set
            }
            Err(err) => {
                trace.push(format!(
                    "rewrite failed; falling back to original query ({})",
                    err.root_cause()
                ));
                QueryRewriteSet::fallback(&command.question)
            }
        }
    }
    .instrument(span)
    .await
}

async fn stripped_answer(
    runtime: &QueryRuntime,
    command: &QueryCommand,
    result: &QueryResult,
) -> Result<String> {
    let answer = match &result.answer {
        Some(answer) => answer.clone(),
        None => generate_answer_from_hits(runtime, command, &result.hits).await?,
    };
    Ok(store::strip_thinking(&answer))
}

async fn generate_answer_from_hits(
    runtime: &QueryRuntime,
    command: &QueryCommand,
    hits: &[RetrievalCandidate],
) -> Result<String> {
    let span = tracing::info_span!(
        "generate_answer_from_hits",
        hit_count = hits.len(),
        question_chars = command.question.chars().count(),
        max_tokens = command.max_tokens,
    );
    async move {
        let generator = Generator::new(
            runtime.cfg.ollama.base_url.clone(),
            runtime.gen_model_name.clone(),
        );
        let contexts = labeled_contexts(hits);
        generator
            .generate_answer(&contexts, &command.question, command.max_tokens)
            .await
    }
    .instrument(span)
    .await
}

fn initial_agent_queries(
    question: &str,
    plan: &QueryPlan,
    rewrite_set: &QueryRewriteSet,
) -> Vec<String> {
    match plan.strategy {
        RetrievalStrategy::Decompose if !plan.subqueries.is_empty() => plan.subqueries.clone(),
        _ => rewrite_set.query_variants(),
    }
    .into_iter()
    .chain(std::iter::once(question.to_string()))
    .collect::<BTreeSet<_>>()
    .into_iter()
    .collect()
}

fn refine_agent_queries(
    question: &str,
    plan: &QueryPlan,
    rewrite_set: &QueryRewriteSet,
    assessment: &EvidenceAssessment,
) -> Vec<String> {
    if !assessment.missing_aspects.is_empty() {
        return assessment
            .missing_aspects
            .iter()
            .cloned()
            .chain(std::iter::once(question.to_string()))
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect();
    }

    if matches!(plan.strategy, RetrievalStrategy::Decompose) && !plan.subqueries.is_empty() {
        return plan
            .subqueries
            .iter()
            .cloned()
            .chain(rewrite_set.query_variants())
            .collect::<BTreeSet<_>>()
            .into_iter()
            .collect();
    }

    rewrite_set.query_variants()
}

fn assessment_notes(assessment: &EvidenceAssessment) -> Vec<String> {
    let mut notes = vec![format!(
        "evidence verdict={}",
        evidence_verdict_label(&assessment.verdict)
    )];
    if !assessment.missing_aspects.is_empty() {
        notes.push(format!(
            "missing_aspects={}",
            assessment.missing_aspects.join(" | ")
        ));
    }
    notes
}

fn execution_path(mode: QueryModeArg) -> QueryExecutionPath {
    match mode {
        QueryModeArg::Naive | QueryModeArg::Hybrid => QueryExecutionPath::Simple,
        QueryModeArg::Agentic => QueryExecutionPath::Agentic,
        QueryModeArg::Local | QueryModeArg::Global | QueryModeArg::Mix => {
            QueryExecutionPath::AgenticStub
        }
    }
}

fn execution_path_label(path: QueryExecutionPath) -> &'static str {
    match path {
        QueryExecutionPath::Simple => "simple",
        QueryExecutionPath::Agentic => "agentic",
        QueryExecutionPath::AgenticStub => "agentic_stub",
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::commands::index;
    use crate::test_support::{sequential_json_server, with_test_env};

    use super::super::runtime::retrieval_limit;

    #[tokio::test(flavor = "current_thread")]
    async fn test_run_answers_from_indexed_store_with_mock_ollama() {
        let dir = tempfile::tempdir().unwrap();
        let docs = dir.path().join("docs");
        std::fs::create_dir_all(&docs).unwrap();
        let input = docs.join("note.txt");
        std::fs::write(&input, "The project is a local RAG CLI.").unwrap();

        let index_server = sequential_json_server(vec![r#"{"embeddings":[[0.1,0.2]]}"#]);
        with_test_env(dir.path(), Some(&index_server), || async {
            index::run(
                Some("e2e"),
                input.clone(),
                Some(200),
                Some(0),
                None,
                None,
                Vec::new(),
                false,
                false,
            )
            .await
            .unwrap();
        })
        .await;

        let query_server = sequential_json_server(vec![
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"ragcli is a local RAG CLI."}}"#,
        ]);
        with_test_env(dir.path(), Some(&query_server), || async {
            run(
                Some("e2e"),
                QueryCommand {
                    question: "What is this project?".to_string(),
                    mode: QueryModeArg::Hybrid,
                    top_k: 5,
                    fetch_k: 20,
                    max_iterations: 2,
                    rewrite: false,
                    rerank: false,
                    show_context: true,
                    show_plan: false,
                    show_scores: false,
                    show_citations: false,
                    show_trace: false,
                    source: None,
                    path_prefix: None,
                    page: None,
                    format: None,
                    gen_model: None,
                    max_tokens: 64,
                    json: false,
                },
            )
            .await
            .unwrap();
        })
        .await;
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_rewrite_falls_back_to_original_query_when_json_is_invalid() {
        let dir = tempfile::tempdir().unwrap();
        let docs = dir.path().join("docs");
        std::fs::create_dir_all(&docs).unwrap();
        let input = docs.join("note.txt");
        std::fs::write(
            &input,
            "Configuration resolution happens during query execution.",
        )
        .unwrap();

        let index_server = sequential_json_server(vec![r#"{"embeddings":[[0.1,0.2]]}"#]);
        with_test_env(dir.path(), Some(&index_server), || async {
            index::run(
                Some("rewrite-fallback"),
                input.clone(),
                Some(200),
                Some(0),
                None,
                None,
                Vec::new(),
                false,
                false,
            )
            .await
            .unwrap();
        })
        .await;

        let query_server = sequential_json_server(vec![
            r#"{"message":{"content":"not json"}}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"Configuration resolution happens during query execution."}}"#,
        ]);
        with_test_env(dir.path(), Some(&query_server), || async {
            run(
                Some("rewrite-fallback"),
                QueryCommand {
                    question: "How does config resolution work?".to_string(),
                    mode: QueryModeArg::Hybrid,
                    top_k: 5,
                    fetch_k: 20,
                    max_iterations: 2,
                    rewrite: true,
                    rerank: false,
                    show_context: false,
                    show_plan: false,
                    show_scores: false,
                    show_citations: false,
                    show_trace: false,
                    source: None,
                    path_prefix: None,
                    page: None,
                    format: None,
                    gen_model: None,
                    max_tokens: 64,
                    json: false,
                },
            )
            .await
            .unwrap();
        })
        .await;
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_agentic_mode_retries_when_evidence_is_partial() {
        let dir = tempfile::tempdir().unwrap();
        let docs = dir.path().join("docs");
        std::fs::create_dir_all(&docs).unwrap();
        let input = docs.join("note.txt");
        std::fs::write(
            &input,
            "Config resolution interacts with metadata validation during query execution.",
        )
        .unwrap();

        let index_server = sequential_json_server(vec![r#"{"embeddings":[[0.1,0.2]]}"#]);
        with_test_env(dir.path(), Some(&index_server), || async {
            index::run(
                Some("agentic-retry"),
                input.clone(),
                Some(200),
                Some(0),
                None,
                None,
                Vec::new(),
                false,
                false,
            )
            .await
            .unwrap();
        })
        .await;

        let query_server = sequential_json_server(vec![
            r#"{"message":{"content":"{\"question_type\":\"MultiHop\",\"strategy\":\"Decompose\",\"reasoning\":\"needs two facts\",\"subqueries\":[\"config resolution\",\"metadata validation\"]}"}}"#,
            r#"{"message":{"content":"{\"semantic_variant\":\"How does config resolution interact with metadata validation?\",\"keyword_variant\":\"config resolution metadata validation\",\"subqueries\":[\"config resolution\",\"metadata validation\"]}"}}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"{\"verdict\":\"partial\",\"missing_aspects\":[\"interaction between config resolution and metadata validation\"]}"}}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"{\"verdict\":\"sufficient\",\"missing_aspects\":[]}"}}"#,
            r#"{"message":{"content":"Config resolution interacts with metadata validation during query execution. [1]"}}"#,
            r#"{"message":{"content":"{\"supported\":true,\"unsupported_claims\":[],\"notes\":[\"evidence covers the answer\"]}"}}"#,
        ]);
        with_test_env(dir.path(), Some(&query_server), || async {
            run(
                Some("agentic-retry"),
                QueryCommand {
                    question: "How do config resolution and metadata validation interact?"
                        .to_string(),
                    mode: QueryModeArg::Agentic,
                    top_k: 5,
                    fetch_k: 20,
                    max_iterations: 2,
                    rewrite: true,
                    rerank: false,
                    show_context: false,
                    show_plan: false,
                    show_scores: false,
                    show_citations: false,
                    show_trace: false,
                    source: None,
                    path_prefix: None,
                    page: None,
                    format: None,
                    gen_model: None,
                    max_tokens: 64,
                    json: false,
                },
            )
            .await
            .unwrap();
        })
        .await;
    }

    #[test]
    fn test_retrieval_limit_uses_fetch_k_when_larger_than_top_k() {
        let command = QueryCommand {
            question: "What is ragcli?".to_string(),
            mode: QueryModeArg::Hybrid,
            top_k: 5,
            fetch_k: 20,
            max_iterations: 2,
            rewrite: false,
            rerank: false,
            show_context: false,
            show_plan: false,
            show_scores: false,
            show_citations: false,
            show_trace: false,
            source: None,
            path_prefix: None,
            page: None,
            format: None,
            gen_model: None,
            max_tokens: 256,
            json: false,
        };

        assert_eq!(retrieval_limit(&command), 20);
    }

    #[test]
    fn test_execution_path_routes_agentic_modes_to_stub() {
        assert_eq!(
            execution_path(QueryModeArg::Hybrid),
            QueryExecutionPath::Simple
        );
        assert_eq!(
            execution_path(QueryModeArg::Naive),
            QueryExecutionPath::Simple
        );
        assert_eq!(
            execution_path(QueryModeArg::Agentic),
            QueryExecutionPath::Agentic
        );
        assert_eq!(
            execution_path(QueryModeArg::Local),
            QueryExecutionPath::AgenticStub
        );
        assert_eq!(
            execution_path(QueryModeArg::Global),
            QueryExecutionPath::AgenticStub
        );
        assert_eq!(
            execution_path(QueryModeArg::Mix),
            QueryExecutionPath::AgenticStub
        );
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_local_global_and_mix_modes_use_distinct_placeholder_paths() {
        let dir = tempfile::tempdir().unwrap();
        let docs = dir.path().join("docs");
        std::fs::create_dir_all(&docs).unwrap();
        let input = docs.join("note.txt");
        std::fs::write(&input, "Config checks use metadata validation.").unwrap();

        let index_server = sequential_json_server(vec![r#"{"embeddings":[[0.1,0.2]]}"#]);
        with_test_env(dir.path(), Some(&index_server), || async {
            index::run(
                Some("graph-modes"),
                input.clone(),
                Some(200),
                Some(0),
                None,
                None,
                Vec::new(),
                false,
                false,
            )
            .await
            .unwrap();
        })
        .await;

        let query_server = sequential_json_server(vec![
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"local answer [1]"}}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"global answer [1]"}}"#,
            r#"{"embeddings":[[0.1,0.2]]}"#,
            r#"{"message":{"content":"mix answer [1]"}}"#,
        ]);
        with_test_env(dir.path(), Some(&query_server), || async {
            for mode in [QueryModeArg::Local, QueryModeArg::Global, QueryModeArg::Mix] {
                run(
                    Some("graph-modes"),
                    QueryCommand {
                        question: "How do config checks work?".to_string(),
                        mode,
                        top_k: 5,
                        fetch_k: 20,
                        max_iterations: 2,
                        rewrite: false,
                        rerank: false,
                        show_context: false,
                        show_plan: false,
                        show_scores: false,
                        show_citations: false,
                        show_trace: false,
                        source: None,
                        path_prefix: None,
                        page: None,
                        format: None,
                        gen_model: None,
                        max_tokens: 64,
                        json: false,
                    },
                )
                .await
                .unwrap();
            }
        })
        .await;
    }
}
