use std::sync::Arc;

use anyhow::Result;
use tracing::{debug, warn};

use crate::types::{LlmProvider, SearchResult};

/// Uses an LLM to judge whether a set of search results is "sufficient"
/// to answer the user's query, or whether the orchestrator should escalate
/// to a more expensive search tier.
pub struct ResultJudge {
    llm: Arc<dyn LlmProvider>,
    min_relevance_score: f64,
}

impl ResultJudge {
    pub fn new(llm: Arc<dyn LlmProvider>, min_relevance_score: f64) -> Self {
        Self {
            llm,
            min_relevance_score,
        }
    }

    /// Determine whether the given results sufficiently answer the query.
    ///
    /// Uses a two-stage check:
    /// 1. **Heuristic gate**: if no results or all scores are below
    ///    `min_relevance_score`, immediately return `false`.
    /// 2. **LLM judge**: ask the LLM whether the results contain the
    ///    information needed to answer the query.
    pub async fn is_sufficient(&self, query: &str, results: &[SearchResult]) -> Result<bool> {
        // Stage 1: heuristic gate.
        if results.is_empty() {
            debug!(query = %query, "ResultJudge: no results -> insufficient");
            return Ok(false);
        }

        let above_threshold: Vec<&SearchResult> = results
            .iter()
            .filter(|r| r.relevance_score >= self.min_relevance_score)
            .collect();

        if above_threshold.is_empty() {
            debug!(
                query = %query,
                min_score = self.min_relevance_score,
                best_score = results
                    .iter()
                    .map(|r| r.relevance_score)
                    .fold(0.0_f64, f64::max),
                "ResultJudge: all results below threshold -> insufficient"
            );
            return Ok(false);
        }

        // If we have several high-confidence results, skip the LLM call to save cost.
        let high_confidence_count = results
            .iter()
            .filter(|r| r.relevance_score >= 0.8)
            .count();

        if high_confidence_count >= 3 {
            debug!(
                query = %query,
                high_confidence_count,
                "ResultJudge: multiple high-confidence results -> sufficient (fast path)"
            );
            return Ok(true);
        }

        // Stage 2: LLM judge (only if LLM is available).
        if !self.llm.is_available() {
            debug!(
                query = %query,
                above_threshold = above_threshold.len(),
                "ResultJudge: LLM unavailable, using threshold-only heuristic"
            );
            // Without LLM, accept results that passed the threshold gate.
            return Ok(true);
        }
        self.llm_judge(query, &above_threshold).await
    }

    /// Ask the LLM to judge result sufficiency.
    async fn llm_judge(&self, query: &str, results: &[&SearchResult]) -> Result<bool> {
        let system = r#"You are a search result quality judge for a code intelligence system.
Given a query and a set of code search results, determine whether the results contain enough relevant information to answer the query.

Respond with ONLY "YES" or "NO".
- "YES" means the results are sufficient -- they contain the relevant code/information.
- "NO" means the results are insufficient -- deeper search is needed."#;

        let mut result_summaries = String::new();
        for (i, r) in results.iter().enumerate().take(5) {
            let snippet: String = r.content.lines().take(5).collect::<Vec<_>>().join("\n");
            result_summaries.push_str(&format!(
                "\n--- Result {} (score={:.2}, file={}) ---\n{}\n",
                i + 1,
                r.relevance_score,
                r.file_path,
                snippet
            ));
        }

        let prompt = format!(
            "Query: {}\n\nSearch Results:{}\n\nAre these results sufficient to answer the query?",
            query, result_summaries
        );

        let response = tokio::time::timeout(
            std::time::Duration::from_secs(30),
            self.llm.complete_with_system(system, &prompt),
        )
        .await
        .map_err(|_| anyhow::anyhow!("LLM result judgment timed out after 30s"))??;
        let trimmed = response.trim().to_uppercase();

        let sufficient = trimmed.starts_with("YES");

        debug!(
            query = %query,
            llm_response = %trimmed,
            sufficient,
            result_count = results.len(),
            "ResultJudge: LLM judgment"
        );

        if !sufficient {
            warn!(
                query = %query,
                "ResultJudge: LLM determined results are insufficient"
            );
        }

        Ok(sufficient)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::SearchResult;
    use std::sync::Arc;

    struct MockLlm {
        response: String,
    }

    impl MockLlm {
        fn new(response: &str) -> Arc<Self> {
            Arc::new(Self { response: response.to_string() })
        }
    }

    #[async_trait::async_trait]
    impl LlmProvider for MockLlm {
        async fn complete(&self, _prompt: &str) -> anyhow::Result<String> {
            Ok(self.response.clone())
        }
        async fn complete_with_system(&self, _system: &str, _prompt: &str) -> anyhow::Result<String> {
            Ok(self.response.clone())
        }
    }

    fn make_result(score: f64) -> SearchResult {
        SearchResult {
            file_path: "test.rs".into(),
            content: "fn test() {}".into(),
            relevance_score: score,
            line_number: Some(1),
        }
    }

    #[tokio::test]
    async fn test_empty_results_insufficient() {
        let judge = ResultJudge::new(MockLlm::new("YES"), 0.6);
        assert!(!judge.is_sufficient("test query", &[]).await.unwrap());
    }

    #[tokio::test]
    async fn test_all_below_threshold_insufficient() {
        let judge = ResultJudge::new(MockLlm::new("YES"), 0.7);
        let results = vec![make_result(0.3), make_result(0.5)];
        assert!(!judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_high_confidence_fast_path() {
        // 3+ results >= 0.8 should return true without LLM call
        let judge = ResultJudge::new(MockLlm::new("NO"), 0.6);
        let results = vec![make_result(0.9), make_result(0.85), make_result(0.88)];
        // Even though mock LLM says "NO", fast path should return true
        assert!(judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_two_high_confidence_calls_llm() {
        // Only 2 high-confidence results -> LLM is called
        let judge = ResultJudge::new(MockLlm::new("YES"), 0.6);
        let results = vec![make_result(0.9), make_result(0.85), make_result(0.7)];
        assert!(judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_llm_says_no() {
        let judge = ResultJudge::new(MockLlm::new("NO"), 0.6);
        let results = vec![make_result(0.75)];
        assert!(!judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_llm_unexpected_response_treated_as_no() {
        let judge = ResultJudge::new(MockLlm::new("MAYBE"), 0.6);
        let results = vec![make_result(0.75)];
        assert!(!judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_boundary_exactly_three_high_confidence() {
        let judge = ResultJudge::new(MockLlm::new("NO"), 0.6);
        let results = vec![make_result(0.8), make_result(0.8), make_result(0.8)];
        assert!(judge.is_sufficient("test query", &results).await.unwrap());
    }

    #[tokio::test]
    async fn test_custom_threshold() {
        let judge = ResultJudge::new(MockLlm::new("YES"), 0.9);
        let results = vec![make_result(0.85), make_result(0.75)];
        // Both below 0.9 threshold
        assert!(!judge.is_sufficient("test query", &results).await.unwrap());
    }
}
