// Copyright 2025 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use anyhow::{Result, anyhow};
use tantivy::{
    TantivyError,
    query::{
        DocumentEvaluation, SingleDocumentEvaluationContext, SingleDocumentEvaluator,
        SingleDocumentPreparer,
    },
    schema::Document,
};

/// Evaluates a parsed FTS query against external documents without an index.
///
/// Documents are prepared with Tantivy's indexing pipeline and are then
/// evaluated by the query's single-document evaluator. This keeps brute-force
/// matching, field norms, and positions aligned with indexed search. BM25 uses
/// prepared index statistics when available and valid empty statistics
/// otherwise.
pub struct BruteScoredSearcher {
    query: crate::Query,
    // `None` until the first fallback document is prepared. The indexed plus
    // fallback reader pipeline runs in this order:
    // 1. `new` clones the shared query while `prepared_bm25` may still be empty.
    // 2. Indexed search installs global statistics with `prepare_bm25_once`.
    // 3. The first fallback document initializes the evaluator with the global statistics. If
    //    unavailable, the query provides empty statistics.
    // Initializing in `new` would cache empty-statistics BM25 weights before
    // indexed search has a chance to install the global statistics.
    evaluator_state: Option<EvaluatorState>,
}

struct EvaluatorState {
    preparer: SingleDocumentPreparer,
    /// Cached after the first document initializes the searcher.
    evaluator: Box<dyn SingleDocumentEvaluator>,
}

impl BruteScoredSearcher {
    /// Creates a lazy brute-force evaluator.
    ///
    /// Query compilation is deferred until a real document is prepared. This
    /// lets indexed-only reads use query types that Tantivy cannot evaluate
    /// without an index, while still returning an explicit error if fallback is
    /// actually required.
    pub fn new(query: &crate::Query) -> Result<Self> {
        match query.info().get_query_type() {
            tipb::FtsQueryType::FtsQueryTypeNoScore | tipb::FtsQueryType::FtsQueryTypeWithScore => {
            }
            query_type => return Err(anyhow!("unsupported FTS query type {query_type:?}")),
        }
        Ok(Self {
            query: query.clone(),
            evaluator_state: None,
        })
    }

    fn compile_evaluator(&self) -> Result<Box<dyn SingleDocumentEvaluator>> {
        let schema = self.query.evaluation_schema();
        let context = match self.query.info().get_query_type() {
            tipb::FtsQueryType::FtsQueryTypeWithScore => {
                SingleDocumentEvaluationContext::with_scoring(schema, &self.query)
            }
            tipb::FtsQueryType::FtsQueryTypeNoScore => {
                SingleDocumentEvaluationContext::without_scoring(schema)
            }
            query_type => return Err(anyhow!("unsupported FTS query type {query_type:?}")),
        };
        self.query
            .tantivy_query()
            .single_document_evaluator(context)
            .map_err(map_evaluator_error)
    }

    fn get_or_init_evaluator_state(&mut self) -> Result<&mut EvaluatorState> {
        if self.evaluator_state.is_some() {
            return self
                .evaluator_state
                .as_mut()
                .ok_or_else(|| anyhow!("evaluator state is not initialized"));
        }

        let evaluator = self.compile_evaluator()?;
        let required_fields = self
            .query
            .document_fields()
            .iter()
            .map(|field| field.field())
            .collect::<Vec<_>>();
        let preparer = SingleDocumentPreparer::for_fields(
            self.query.evaluation_schema(),
            self.query.document_tokenizers(),
            &required_fields,
        )
        .map_err(map_evaluator_error)?;
        Ok(self.evaluator_state.insert(EvaluatorState {
            preparer,
            evaluator,
        }))
    }

    /// Prepares and evaluates one document.
    ///
    /// Scoring uses prepared shard statistics when available and empty
    /// statistics otherwise.
    pub fn evaluate_document<D: Document>(&mut self, document: &D) -> Result<DocumentEvaluation> {
        let evaluator_state = self.get_or_init_evaluator_state()?;
        let prepared_document = evaluator_state
            .preparer
            .prepare(document)
            .map_err(map_evaluator_error)?;
        evaluator_state
            .evaluator
            .evaluate(&prepared_document)
            .map_err(map_evaluator_error)
    }
}

fn map_evaluator_error(error: TantivyError) -> anyhow::Error {
    match error {
        TantivyError::UnsupportedQueryForSingleDocumentEvaluation(query_path) => {
            anyhow!("brute-force FTS fallback does not support this query: {query_path}")
        }
        other => other.into(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::test_util::{PlainFtsQueryInfo, make_scored_query, make_unscored_query};

    fn text_document(query: &crate::Query, text: &str) -> tantivy::TantivyDocument {
        let mut document = tantivy::TantivyDocument::new();
        document.add_text(query.document_fields()[0].field(), text);
        document
    }

    #[test]
    fn unscored_documents_return_expected_matches() -> Result<()> {
        let query = make_unscored_query(&PlainFtsQueryInfo {
            query: "popular".to_string(),
            ..Default::default()
        });
        let mut brute = BruteScoredSearcher::new(&query)?;
        let matched = brute.evaluate_document(&text_document(
            &query,
            "Being too popular can be such a hassle",
        ))?;
        let unmatched = brute.evaluate_document(&text_document(&query, "unrelated"))?;

        // Keep brute-force results aligned with indexed no-score search, whose
        // public contract returns a constant score of 1.0.
        assert_eq!(matched, DocumentEvaluation::Match(1.0));
        assert_eq!(unmatched, DocumentEvaluation::NoMatch);
        Ok(())
    }

    #[test]
    fn score_without_prepared_statistics_does_not_depend_on_other_documents() -> Result<()> {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "machine".to_string(),
            ..Default::default()
        });
        let mut brute = BruteScoredSearcher::new(&query)?;
        let first = brute.evaluate_document(&text_document(&query, "machine learning machine"))?;
        assert_eq!(
            brute.evaluate_document(&text_document(&query, "unrelated"))?,
            DocumentEvaluation::NoMatch
        );
        let second = brute.evaluate_document(&text_document(&query, "machine learning machine"))?;

        assert_eq!(first, second);
        Ok(())
    }

    #[test]
    fn scored_single_document_without_prepared_statistics_uses_empty_statistics() -> Result<()> {
        let query_info = PlainFtsQueryInfo {
            query: "machine".to_string(),
            ..Default::default()
        };
        let query = make_scored_query(&query_info);
        let mut brute = BruteScoredSearcher::new(&query)?;
        let actual = brute.evaluate_document(&text_document(&query, "machine learning machine"))?;

        let prepared_query = make_scored_query(&query_info);
        assert!(prepared_query.prepare_bm25_once(crate::Bm25Stats::empty(&prepared_query))?);
        let mut prepared_brute = BruteScoredSearcher::new(&prepared_query)?;
        let expected = prepared_brute
            .evaluate_document(&text_document(&prepared_query, "machine learning machine"))?;

        assert_eq!(actual, expected);
        Ok(())
    }

    #[test]
    fn statistics_prepared_after_searcher_construction_are_used() -> Result<()> {
        let documents = ["machine learning machine", "unrelated"];
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "machine".to_string(),
            ..Default::default()
        });
        let index = crate::index_for_test(&documents)?.finalize()?;
        let reader = crate::IndexReader::from_tantivy_index(index)?;
        let mut brute = BruteScoredSearcher::new(&query)?;

        let mut stats = crate::Bm25Stats::empty(&query);
        reader.accumulate_bm25_stats(&query, &mut stats)?;
        assert!(query.prepare_bm25_once(stats)?);

        let mut indexed_results = Vec::new();
        reader.search(&query, &mut indexed_results)?;
        let evaluation = brute.evaluate_document(&text_document(&query, documents[0]))?;
        assert_eq!(
            evaluation,
            DocumentEvaluation::Match(indexed_results[0].score)
        );
        Ok(())
    }

    fn assert_match_eq(text: &str, query_text: &str) -> Result<()> {
        let query = make_unscored_query(&PlainFtsQueryInfo {
            query: query_text.to_string(),
            ..Default::default()
        });
        let index = crate::index_for_test(&[text])?.finalize()?;
        let reader = crate::IndexReader::from_tantivy_index(index)?;
        let mut indexed_results = Vec::new();
        reader.search(&query, &mut indexed_results)?;

        let mut brute = BruteScoredSearcher::new(&query)?;
        let brute_result = brute.evaluate_document(&text_document(&query, text))?;
        assert_eq!(
            indexed_results.is_empty(),
            brute_result == DocumentEvaluation::NoMatch,
            "indexed and brute-force results differ for query {query_text:?} in {text:?}"
        );
        Ok(())
    }

    #[test]
    fn unscored_queries_match_indexed_search() -> Result<()> {
        let text = "Being too popular can be such a hassle";
        for query in [
            "popular",
            "people",
            "peoples",
            "can",
            "being",
            "beING",
            "CAN",
            "furina",
            "eing",
            "foo bar",
            "foo bar be",
            "ass",
            "ass a",
            "   ass    a   ",
            "elssah",
        ] {
            assert_match_eq(text, query)?;
        }
        Ok(())
    }

    const SCORE_SAMPLE_DOCS: &[&str] = &[
        "machine learning machine learning machine learning",
        "machine learning algorithms. Machine learning models. Machine learning optimization.",
        "Machine learning (Machine learning (Machine learning (Machine learning (Machine learning (is used in AI. Machine learning (Machine learning (Machine learning (Machine learning (Machine learning (is used in AI. Machine learning (Machine learning (Machine learning (Machine learning (Machine learning (is used in AI. Machine learning ends.",
        "Understanding machine learning: basics of learning from machines. Machine learning requires data.",
        "machine learning",
        "This 50-word document briefly mentions machine learning once. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. Generic text about AI. ",
        "Learning machine learning machine. Learning machine applications.",
        "Space Tourism: Risks and Opportunities - Private companies like SpaceX and Blue Origin are making suborbital flights accessible, but safety protocols remain a critical concern.",
        "Sustainable Fashion Trends - Eco-friendly materials like mushroom leather and recycled polyester are reshaping the fashion industry’s environmental impact.",
        "AI-powered tools are transforming medical imaging analysis. Deep learning algorithms now detect early-stage tumors with 95% accuracy.",
    ];

    fn assert_scores_eq(query_text: &str) -> Result<()> {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: query_text.to_string(),
            ..Default::default()
        });
        let index = crate::index_for_test(SCORE_SAMPLE_DOCS)?.finalize()?;
        let reader = crate::IndexReader::from_tantivy_index(index)?;
        let mut stats = crate::Bm25Stats::empty(&query);
        reader.accumulate_bm25_stats(&query, &mut stats)?;
        assert!(query.prepare_bm25_once(stats)?);
        let mut indexed_results = Vec::new();
        reader.search(&query, &mut indexed_results)?;

        let mut brute_results = Vec::new();
        let mut brute = BruteScoredSearcher::new(&query)?;
        for (doc_id, text) in SCORE_SAMPLE_DOCS.iter().enumerate() {
            if let DocumentEvaluation::Match(score) =
                brute.evaluate_document(&text_document(&query, text))?
            {
                brute_results.push(crate::ScoredResult {
                    doc_id: doc_id as u32,
                    score,
                });
            }
        }

        assert_eq!(indexed_results.len(), brute_results.len());
        for (indexed, brute) in indexed_results.iter().zip(&brute_results) {
            assert_eq!(indexed.doc_id, brute.doc_id);
            assert!(
                (indexed.score - brute.score).abs() < 0.001,
                "score mismatch for query {query_text:?}: indexed={indexed:?}, brute={brute:?}"
            );
        }
        Ok(())
    }

    #[test]
    fn scored_queries_match_indexed_search() -> Result<()> {
        for query in [
            "machine learn",
            "machine learning",
            "machine learning machine learning",
            "machine learning machine",
            "machine learning learning",
            "machine learning algorithms",
            "machine learning optimization",
            "machine learning is used in AI",
            "machine learning ends",
            "machine learning requires data",
            "learning machine",
            "learning machine applications",
            "space tourism",
            "tantivy bm25",
        ] {
            assert_scores_eq(query)?;
        }
        Ok(())
    }

    #[test]
    fn duplicate_scored_query_terms_preserve_occurrence_weight() -> Result<()> {
        const DOCUMENT: &str = "machine learning machine";

        let index = crate::index_for_test(&[DOCUMENT, "unrelated"])?.finalize()?;
        let reader = crate::IndexReader::from_tantivy_index(index)?;
        let score = |query_text: &str| -> Result<(f32, f32, usize)> {
            let query = make_scored_query(&PlainFtsQueryInfo {
                query: query_text.to_string(),
                ..Default::default()
            });
            let mut stats = crate::Bm25Stats::empty(&query);
            reader.accumulate_bm25_stats(&query, &mut stats)?;
            assert!(query.prepare_bm25_once(stats)?);

            let mut indexed_results = Vec::new();
            reader.search(&query, &mut indexed_results)?;
            assert_eq!(indexed_results.len(), 1);
            assert_eq!(indexed_results[0].doc_id, 0);

            let mut brute = BruteScoredSearcher::new(&query)?;
            let DocumentEvaluation::Match(brute_score) =
                brute.evaluate_document(&text_document(&query, DOCUMENT))?
            else {
                return Err(anyhow!("expected duplicate-term query to match"));
            };
            Ok((
                indexed_results[0].score,
                brute_score,
                query.bm25_terms().len(),
            ))
        };

        let (single_indexed, single_brute, single_terms) = score("machine")?;
        let (duplicate_indexed, duplicate_brute, duplicate_terms) = score("machine machine")?;

        assert_eq!(single_terms, 1);
        assert_eq!(duplicate_terms, 1);
        assert!((duplicate_indexed - 2.0 * single_indexed).abs() < 0.001);
        assert!((duplicate_brute - 2.0 * single_brute).abs() < 0.001);
        assert!((duplicate_brute - duplicate_indexed).abs() < 0.001);
        Ok(())
    }
}

#[cfg(test)]
mod benches {
    use std::{fs, hint::black_box, io};

    use super::*;
    use crate::query::test_util::{PlainFtsQueryInfo, make_scored_query};

    fn prepare_bench_data() -> Vec<String> {
        use io::BufRead;

        let path = std::path::Path::new("../testdata/Arts_Crafts_and_Sewing.jsonl.gz");
        if !path.exists() {
            panic!("Test data file does not exist, see testdata/README.md for details");
        }
        let gz_file = flate2::read::GzDecoder::new(fs::File::open(path).unwrap());
        let lines = io::BufReader::new(gz_file).lines();
        let mut data = Vec::new();
        for line in lines.map_while(Result::ok) {
            let value: serde_json::Value = serde_json::from_str(&line).unwrap();
            data.push(value["text"].as_str().unwrap().to_string());
            if data.len() >= 100 {
                break;
            }
        }
        data
    }

    fn text_document(query: &crate::Query, text: &str) -> tantivy::TantivyDocument {
        let mut document = tantivy::TantivyDocument::new();
        document.add_text(query.document_fields()[0].field(), text);
        document
    }

    #[bench]
    #[ignore]
    fn bench_scored_brute_search(b: &mut test::Bencher) {
        let data = prepare_bench_data();
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "sewing machine".to_string(),
            ..Default::default()
        });
        let documents = data
            .iter()
            .map(|text| text_document(&query, text))
            .collect::<Vec<_>>();

        b.iter(|| {
            let mut searcher = BruteScoredSearcher::new(&query).unwrap();
            let matches = documents
                .iter()
                .filter(|document| {
                    matches!(
                        searcher.evaluate_document(*document).unwrap(),
                        DocumentEvaluation::Match(_)
                    )
                })
                .count();
            black_box(matches);
            assert_eq!(4, matches);
        });
    }

    #[bench]
    #[ignore]
    fn bench_scored_tantivy_searcher_prebuilt(b: &mut test::Bencher) {
        let data = prepare_bench_data();
        let data = data.iter().map(String::as_str).collect::<Vec<_>>();
        let index = crate::index_for_test(&data).unwrap().finalize().unwrap();
        let reader = crate::IndexReader::from_tantivy_index(index).unwrap();
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "sewing machine".to_string(),
            ..Default::default()
        });

        let mut results = Vec::new();
        b.iter(|| {
            reader.search(&query, &mut results).unwrap();
            black_box(results.len());
            assert_eq!(4, results.len());
        });
    }

    #[bench]
    #[ignore]
    fn bench_scored_tantivy_searcher_on_demand(b: &mut test::Bencher) {
        let data = prepare_bench_data();
        let data = data.iter().map(String::as_str).collect::<Vec<_>>();
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "sewing machine".to_string(),
            ..Default::default()
        });

        let mut results = Vec::new();
        b.iter(|| {
            let index = crate::index_for_test(&data).unwrap().finalize().unwrap();
            let reader = crate::IndexReader::from_tantivy_index(index).unwrap();
            reader.search(&query, &mut results).unwrap();
            black_box(results.len());
            assert_eq!(4, results.len());
        });
    }
}