terraphim_grep 1.21.2

use super::hybrid_searcher::{HybridResults, RetrievedChunk};

#[derive(Debug, Clone)]
pub struct HeuristicThresholds {
    pub min_coverage: f64,
    pub min_kg_confidence: f64,
    pub min_diversity: usize,
    pub min_results: usize,
}

impl Default for HeuristicThresholds {
    fn default() -> Self {
        Self {
            min_coverage: 0.7,
            min_kg_confidence: 0.5,
            min_diversity: 2,
            min_results: 3,
        }
    }
}

#[derive(Debug, Clone)]
pub enum Sufficiency {
    Sufficient(Vec<RetrievedChunk>),
    NeedsSynthesis(Vec<RetrievedChunk>),
    NeedsExpansion(Vec<RetrievedChunk>),
    Insufficient(Vec<RetrievedChunk>),
}

pub struct SufficiencyJudge {
    thresholds: HeuristicThresholds,
}

impl SufficiencyJudge {
    pub fn new(thresholds: HeuristicThresholds) -> Self {
        Self { thresholds }
    }

    pub fn judge(&self, results: &HybridResults, query: &str) -> Sufficiency {
        let chunks = results.to_chunks();

        if chunks.is_empty() && results.kg_concepts.is_empty() {
            return Sufficiency::Insufficient(vec![]);
        }

        let coverage = self.calculate_coverage(query, &chunks);
        let confidence = self.calculate_kg_confidence(&results.kg_concepts);
        let diversity = self.calculate_diversity(&chunks);

        if chunks.len() < self.thresholds.min_results {
            return Sufficiency::Insufficient(chunks);
        }

        if coverage >= self.thresholds.min_coverage
            && confidence >= self.thresholds.min_kg_confidence
            && diversity >= self.thresholds.min_diversity
        {
            Sufficiency::Sufficient(chunks)
        } else if coverage >= 0.3 && !chunks.is_empty() {
            Sufficiency::NeedsSynthesis(chunks)
        } else if coverage > 0.0 {
            Sufficiency::NeedsExpansion(chunks)
        } else {
            Sufficiency::Insufficient(chunks)
        }
    }

    fn calculate_coverage(&self, query: &str, chunks: &[RetrievedChunk]) -> f64 {
        if chunks.is_empty() {
            return 0.0;
        }

        let query_terms: std::collections::HashSet<String> =
            query.split_whitespace().map(|s| s.to_lowercase()).collect();

        if query_terms.is_empty() {
            return 1.0;
        }

        let mut covered_terms = 0usize;
        for term in &query_terms {
            let term_found = chunks.iter().any(|chunk| {
                chunk.content.to_lowercase().contains(term)
                    || chunk.source.to_lowercase().contains(term)
            });
            if term_found {
                covered_terms += 1;
            }
        }

        covered_terms as f64 / query_terms.len() as f64
    }

    fn calculate_kg_confidence(&self, kg_concepts: &[super::hybrid_searcher::KgConcept]) -> f64 {
        if kg_concepts.is_empty() {
            return 0.0;
        }

        let avg_score: f64 =
            kg_concepts.iter().map(|c| c.score).sum::<f64>() / kg_concepts.len() as f64;
        avg_score.min(1.0)
    }

    fn calculate_diversity(&self, chunks: &[RetrievedChunk]) -> usize {
        let haystacks: std::collections::HashSet<&str> =
            chunks.iter().map(|c| c.haystack).collect();
        haystacks.len()
    }
}

impl Default for SufficiencyJudge {
    fn default() -> Self {
        Self::new(HeuristicThresholds::default())
    }
}

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

    fn make_chunk(content: &str, source: &str, haystack: &'static str) -> RetrievedChunk {
        RetrievedChunk {
            content: content.to_string(),
            source: source.to_string(),
            line_start: Some(1),
            line_end: Some(1),
            relevance_score: 0.8,
            haystack,
        }
    }

    #[test]
    fn test_empty_results_insufficient() {
        let judge = SufficiencyJudge::default();
        let results = HybridResults {
            code_results: vec![],
            doc_results: vec![],
            kg_concepts: vec![],
        };

        let sufficiency = judge.judge(&results, "test query");
        assert!(matches!(sufficiency, Sufficiency::Insufficient(_)));
    }

    #[test]
    fn test_low_results_insufficient() {
        let judge = SufficiencyJudge::default();
        let results = HybridResults {
            code_results: vec![make_chunk("test", "file.rs", "code")],
            doc_results: vec![],
            kg_concepts: vec![],
        };

        let sufficiency = judge.judge(&results, "test query");
        assert!(matches!(sufficiency, Sufficiency::Insufficient(_)));
    }

    #[test]
    fn test_high_coverage_sufficient() {
        let judge = SufficiencyJudge::default();
        let results = HybridResults {
            code_results: vec![
                make_chunk("retry configuration in test file", "retry.rs", "code"),
                make_chunk("backoff settings", "config.rs", "code"),
            ],
            doc_results: vec![make_chunk("retry docs", "docs.md", "docs")],
            kg_concepts: vec![KgConcept {
                id: 1,
                name: "retry".to_string(),
                display_value: None,
                score: 0.9,
            }],
        };

        let sufficiency = judge.judge(&results, "retry configuration");
        assert!(matches!(sufficiency, Sufficiency::Sufficient(_)));
    }

    #[test]
    fn test_coverage_calculation() {
        let judge = SufficiencyJudge::default();
        let chunks = vec![make_chunk("retry configuration", "file.rs", "code")];

        let coverage = judge.calculate_coverage("retry configuration", &chunks);
        assert!(coverage >= 0.9);

        let coverage2 = judge.calculate_coverage("missing term", &chunks);
        assert!(coverage2 < 0.5);
    }

    #[test]
    fn test_diversity_calculation() {
        let judge = SufficiencyJudge::default();
        let chunks = vec![
            make_chunk("code result", "file.rs", "code"),
            make_chunk("code result 2", "file2.rs", "code"),
        ];
        assert_eq!(judge.calculate_diversity(&chunks), 1);

        let chunks2 = vec![
            make_chunk("code result", "file.rs", "code"),
            make_chunk("doc result", "file.md", "docs"),
        ];
        assert_eq!(judge.calculate_diversity(&chunks2), 2);
    }

    #[test]
    fn test_kg_confidence_calculation() {
        let judge = SufficiencyJudge::default();
        let concepts = vec![
            KgConcept {
                id: 1,
                name: "test".to_string(),
                display_value: None,
                score: 0.9,
            },
            KgConcept {
                id: 2,
                name: "test2".to_string(),
                display_value: None,
                score: 0.7,
            },
        ];
        let confidence = judge.calculate_kg_confidence(&concepts);
        assert!((confidence - 0.8).abs() < 0.001);

        let empty_confidence = judge.calculate_kg_confidence(&[]);
        assert_eq!(empty_confidence, 0.0);
    }
}