// 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, bail};
use kvenginepb::fts::FullTextIndexFieldDef;
use tantivy::{
    directory::{Directory, RamDirectory},
    schema::{OwnedValue, document::TantivyDocument},
};

use crate::{
    PhysicalFieldType, PhysicalIndexSchema, TrackedDirectory,
    physical_index_analysis::{build_physical_index_analysis, legacy_index_field},
};

pub struct TantivyIndexWriter<D: Directory> {
    index_writer: tantivy::SingleSegmentIndexWriter,
    dir: D,
    default_field: tantivy::schema::Field,
    physical_schema: PhysicalIndexSchema,
}

impl<D: Directory> TantivyIndexWriter<D> {
    fn new(tokenizer_name: &str, dir: D, layout: &crate::IndexLayout) -> Result<Self> {
        Self::new_with_field_tokenizers(
            std::iter::repeat_n(tokenizer_name, layout.fields().len()),
            dir,
            layout,
        )
    }

    fn new_with_index_fields(
        index_fields: &[FullTextIndexFieldDef],
        dir: D,
        layout: &crate::IndexLayout,
    ) -> Result<Self> {
        // Build the physical schema once before accepting the first document.
        let (physical_schema, tokenizers) = build_physical_index_analysis(index_fields, layout)?;
        let index_writer = tantivy::IndexBuilder::new()
            .tokenizers(tokenizers)
            .schema(physical_schema.tantivy_schema().clone())
            .single_segment_index_writer(dir.box_clone(), 32_000_000_000)?;
        let default_field = physical_schema.index_layout().first_field().field();
        Ok(Self {
            index_writer,
            dir,
            default_field,
            physical_schema,
        })
    }

    fn new_with_field_tokenizers<I, S>(
        tokenizer_names: I,
        dir: D,
        layout: &crate::IndexLayout,
    ) -> Result<Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        let index_fields = tokenizer_names
            .into_iter()
            .map(|tokenizer_name| legacy_index_field(tokenizer_name.as_ref()))
            .collect::<Result<Vec<_>>>()?;
        Self::new_with_index_fields(&index_fields, dir, layout)
    }

    pub fn add_document(&mut self, body: &str) -> Result<()> {
        self.add_document_fields(std::iter::once((self.default_field, body)))?;
        Ok(())
    }

    pub fn add_document_fields<'a, I>(&mut self, fields: I) -> Result<()>
    where
        I: IntoIterator<Item = (tantivy::schema::Field, &'a str)>,
    {
        self.add_document_values(
            fields
                .into_iter()
                .map(|(field, text)| (field, OwnedValue::Str(text.to_owned()))),
        )
    }

    pub fn add_document_values<I>(&mut self, fields: I) -> Result<()>
    where
        I: IntoIterator<Item = (tantivy::schema::Field, OwnedValue)>,
    {
        let mut doc = TantivyDocument::new();
        let layout = self.physical_schema.index_layout();
        for (field, value) in fields {
            let field_layout = layout
                .fields()
                .get(field.field_id() as usize)
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "field {} is absent from physical FTS schema",
                        field.field_id()
                    )
                })?;
            let expected = field_layout.field_type();
            let actual = physical_field_type_from_value(&value).ok_or_else(|| {
                anyhow::anyhow!(
                    "unsupported Tantivy value for physical FTS field {}",
                    field.field_id()
                )
            })?;
            if actual != expected {
                bail!(
                    "value type {:?} does not match physical field type {:?} for field {}",
                    actual,
                    expected,
                    field.field_id()
                );
            }
            doc.add_field_value(field, value);
        }
        self.index_writer.add_document(doc)?;
        Ok(())
    }

    pub fn add_null(&mut self) -> Result<()> {
        self.index_writer.add_document(tantivy::doc!())?;
        Ok(())
    }

    pub fn finalize(self) -> Result<tantivy::Index> {
        Ok(self.index_writer.finalize()?)
    }

    pub fn finalize_as_dir(self) -> Result<D> {
        self.index_writer.finalize()?;
        Ok(self.dir)
    }
}

fn physical_field_type_from_value(value: &OwnedValue) -> Option<PhysicalFieldType> {
    match value {
        OwnedValue::Str(_) => Some(PhysicalFieldType::Text),
        OwnedValue::I64(_) => Some(PhysicalFieldType::I64),
        OwnedValue::U64(_) => Some(PhysicalFieldType::U64),
        OwnedValue::F64(_) => Some(PhysicalFieldType::F64),
        OwnedValue::Bytes(_) => Some(PhysicalFieldType::Bytes),
        _ => None,
    }
}

impl TantivyIndexWriter<TrackedDirectory<RamDirectory>> {
    pub fn new_in_memory(tokenizer_name: &str) -> Result<Self> {
        Self::new_in_memory_with_layout(tokenizer_name, &crate::IndexLayout::legacy_body())
    }

    pub fn new_in_memory_with_layout(
        tokenizer_name: &str,
        layout: &crate::IndexLayout,
    ) -> Result<Self> {
        let dir = TrackedDirectory::wrap(RamDirectory::default());
        TantivyIndexWriter::new(tokenizer_name, dir, layout)
    }

    pub fn new_in_memory_with_index_fields(
        index_fields: &[FullTextIndexFieldDef],
        layout: &crate::IndexLayout,
    ) -> Result<Self> {
        let dir = TrackedDirectory::wrap(RamDirectory::default());
        TantivyIndexWriter::new_with_index_fields(index_fields, dir, layout)
    }

    #[inline]
    pub fn physical_schema(&self) -> &PhysicalIndexSchema {
        &self.physical_schema
    }
}

/// Helper function to build a Tantivy index in memory for testing purposes.
/// The index will contain the provided documents as the body field.
/// Empty documents will be represented as null documents.
pub fn index_for_test(docs: &[&str]) -> Result<TantivyIndexWriter<TrackedDirectory<RamDirectory>>> {
    let mut writer = TantivyIndexWriter::new_in_memory("STANDARD_V1")?;
    if docs.is_empty() {
        writer.add_null()?;
    } else {
        for doc in docs {
            if doc.is_empty() {
                writer.add_null()?;
            } else {
                writer.add_document(doc)?;
            }
        }
    }
    Ok(writer)
}

// Unit tests are placed in index_reader.rs

#[cfg(test)]
mod tests {
    use tantivy::{
        DocSet,
        postings::Postings,
        schema::{Field, FieldType, IndexRecordOption, Schema},
    };

    use super::*;
    use crate::test_util::ngram_index_field;

    fn build_schema_with_index_fields(
        index_fields: &[FullTextIndexFieldDef],
        layout: &crate::IndexLayout,
    ) -> Result<Schema> {
        let physical_schema =
            crate::physical_index_analysis::build_physical_index_analysis(index_fields, layout)?.0;
        Ok(physical_schema.tantivy_schema().clone())
    }

    fn build_schema(tokenizer_name: &str, layout: &crate::IndexLayout) -> Result<Schema> {
        let index_fields = std::iter::repeat_n(tokenizer_name, layout.fields().len())
            .map(crate::physical_index_analysis::legacy_index_field)
            .collect::<Result<Vec<_>>>()?;
        build_schema_with_index_fields(&index_fields, layout)
    }

    fn standard_index_field(enable_bm25: bool) -> FullTextIndexFieldDef {
        let mut index_field = FullTextIndexFieldDef::new();
        let mut index_field_option = kvenginepb::fts::FullTextIndexFieldOption::new();
        index_field_option
            .set_analyzer_type(kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerStandardV1);
        index_field_option.set_enable_bm25(enable_bm25);
        index_field.set_index_field_option(index_field_option);
        index_field
    }

    fn index_field(
        analyzer_type: kvenginepb::fts::FullTextAnalyzerType,
        enable_bm25: bool,
        delimiter: Option<&str>,
    ) -> FullTextIndexFieldDef {
        let mut index_field = FullTextIndexFieldDef::new();
        let option = index_field.mut_index_field_option();
        option.set_analyzer_type(analyzer_type);
        option.set_enable_bm25(enable_bm25);
        if let Some(delimiter) = delimiter {
            option
                .mut_path_hierarchy()
                .set_delimiter(delimiter.to_string());
        }
        index_field
    }

    #[test]
    fn test_ngram_tokenizer_names_and_bm25_index_options() -> Result<()> {
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(0), "0", true),
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(1), "1", false),
        ]);
        let schema = build_schema_with_index_fields(
            &[
                ngram_index_field(10, 3, 3, "WORD", false, true),
                ngram_index_field(11, 2, 3, "CHAR", true, false),
            ],
            &layout,
        )?;

        for (field_id, expected_name, expected_bm25) in [
            (0, "NGRAM_V1_FIELD_0", true),
            (1, "NGRAM_V1_FIELD_1", false),
        ] {
            let FieldType::Str(options) = schema
                .get_field_entry(Field::from_field_id(field_id))
                .field_type()
            else {
                panic!("expected text field {field_id}");
            };
            let indexing = options.get_indexing_options().unwrap();
            assert_eq!(indexing.tokenizer(), expected_name);
            assert_eq!(indexing.fieldnorms(), expected_bm25);
            assert_eq!(
                indexing.index_option(),
                if expected_bm25 {
                    IndexRecordOption::WithFreqsAndPositions
                } else {
                    IndexRecordOption::Basic
                }
            );
        }
        Ok(())
    }

    #[test]
    fn test_ngram_fields_use_independent_specs_in_one_index() -> Result<()> {
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field(Field::from_field_id(0), "0"),
            crate::SingleFieldLayout::new_with_field(Field::from_field_id(1), "1"),
        ]);
        let mut writer = TantivyIndexWriter::new_in_memory_with_index_fields(
            &[
                ngram_index_field(10, 3, 3, "WORD", false, true),
                ngram_index_field(11, 2, 2, "CHAR", true, true),
            ],
            &layout,
        )?;
        writer.add_document_fields([
            (Field::from_field_id(0), "Hello world"),
            (Field::from_field_id(1), "AB C"),
        ])?;
        let index = writer.finalize()?;
        let searcher = index.reader()?.searcher();

        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(
                Field::from_field_id(0),
                "Hel"
            ))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(
                Field::from_field_id(0),
                "hel"
            ))?,
            0
        );
        for term in ["ab", "b ", " c"] {
            assert_eq!(
                searcher.doc_freq(&tantivy::Term::from_field_text(
                    Field::from_field_id(1),
                    term
                ))?,
                1,
                "term: {term:?}"
            );
        }
        Ok(())
    }

    #[test]
    fn test_ngram_writer_indexes_frequency_and_positions() -> Result<()> {
        let field = Field::from_field_id(0);
        let layout =
            crate::IndexLayout::new(vec![crate::SingleFieldLayout::new_with_field_and_bm25(
                field, "0", true,
            )]);
        let mut writer = TantivyIndexWriter::new_in_memory_with_index_fields(
            &[ngram_index_field(10, 3, 3, "WORD", false, true)],
            &layout,
        )?;
        writer.add_document("hello hello")?;
        let index = writer.finalize()?;
        let reader = index.reader()?;
        let searcher = reader.searcher();
        let inverted_index = searcher.segment_readers()[0].inverted_index(field)?;
        let term = tantivy::Term::from_field_text(field, "hel");
        let mut postings = inverted_index
            .read_postings(&term, IndexRecordOption::WithFreqsAndPositions)?
            .unwrap();

        assert_eq!(postings.doc(), 0);
        assert_eq!(postings.term_freq(), 2);
        let mut positions = Vec::new();
        postings.positions(&mut positions);
        assert_eq!(positions, [0, 6]);
        Ok(())
    }

    #[test]
    fn test_ngram_writer_rejects_missing_and_invalid_options() {
        let layout = crate::IndexLayout::from_field_names(["0"]);

        let missing = index_field(
            kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerNgramV1,
            true,
            None,
        );
        let error = TantivyIndexWriter::new_in_memory_with_index_fields(&[missing], &layout)
            .err()
            .unwrap()
            .to_string();
        assert!(error.contains("requires Ngram options"), "{error}");

        let invalid = ngram_index_field(42, 1, 3, "CHAR", false, true);
        let error = TantivyIndexWriter::new_in_memory_with_index_fields(&[invalid], &layout)
            .err()
            .unwrap()
            .to_string();
        assert!(error.contains("field 0 with col_id 42"), "{error}");
        assert!(error.contains("min_gram=1"), "{error}");

        let invalid = ngram_index_field(42, 2, 3, "TOKEN", false, true);
        let error = TantivyIndexWriter::new_in_memory_with_index_fields(&[invalid], &layout)
            .err()
            .unwrap()
            .to_string();
        assert!(error.contains("granularity=\"TOKEN\""), "{error}");
    }

    #[test]
    fn test_legacy_ngram_writer_requires_field_options() {
        let error = TantivyIndexWriter::new_in_memory("NGRAM_V1")
            .err()
            .unwrap()
            .to_string();
        assert!(error.contains("requires Ngram options"), "{error}");
    }

    #[test]
    fn test_build_schema_with_index_fields_disable_bm25() -> Result<()> {
        let layout =
            crate::IndexLayout::new(vec![crate::SingleFieldLayout::new_with_field_and_bm25(
                Field::from_field_id(0),
                "0",
                false,
            )]);
        let schema = build_schema_with_index_fields(&[standard_index_field(false)], &layout)?;
        let field_entry = schema.get_field_entry(layout.first_field().field());
        let FieldType::Str(text_options) = field_entry.field_type() else {
            panic!("expected text field");
        };
        let indexing = text_options.get_indexing_options().unwrap();
        assert_eq!(indexing.tokenizer(), "STANDARD_V1");
        assert!(!indexing.fieldnorms());
        assert_eq!(indexing.index_option(), IndexRecordOption::Basic);
        Ok(())
    }

    #[test]
    fn test_write_index_with_all_bm25_disabled() -> Result<()> {
        let field = Field::from_field_id(0);
        let layout =
            crate::IndexLayout::new(vec![crate::SingleFieldLayout::new_with_field_and_bm25(
                field, "0", false,
            )]);
        let mut writer = TantivyIndexWriter::new_in_memory_with_index_fields(
            &[standard_index_field(false)],
            &layout,
        )?;

        writer.add_document("Hello world")?;
        let index = writer.finalize()?;
        layout.validate_schema(&index.schema())?;
        assert_eq!(
            index
                .reader()?
                .searcher()
                .doc_freq(&tantivy::Term::from_field_text(field, "hello"))?,
            1
        );
        Ok(())
    }

    #[test]
    fn test_build_text_schema_supports_exact_field() -> Result<()> {
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field(Field::from_field_id(0), "0"),
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(1), "1", false),
        ]);
        let exact_field = index_field(
            kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerExactValueV1,
            false,
            None,
        );

        let schema =
            build_schema_with_index_fields(&[standard_index_field(true), exact_field], &layout)?;
        layout.validate_schema(&schema)?;
        let FieldType::Str(options) = schema.get_field_entry(Field::from_field_id(1)).field_type()
        else {
            panic!("expected exact text field");
        };
        let indexing = options.get_indexing_options().unwrap();
        assert_eq!(indexing.tokenizer(), "EXACT_VALUE_V1");
        assert!(!indexing.fieldnorms());
        Ok(())
    }

    #[test]
    fn test_build_schema_with_path_text_field() -> Result<()> {
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field(Field::from_field_id(0), "0"),
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(1), "1", false),
        ]);
        let path_field = index_field(
            kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1,
            false,
            Some("/"),
        );

        let schema =
            build_schema_with_index_fields(&[standard_index_field(true), path_field], &layout)?;

        layout.validate_schema(&schema)?;
        for field_id in 0..2 {
            assert!(matches!(
                schema
                    .get_field_entry(Field::from_field_id(field_id))
                    .field_type(),
                FieldType::Str(_)
            ));
        }
        let FieldType::Str(path_options) =
            schema.get_field_entry(Field::from_field_id(1)).field_type()
        else {
            unreachable!()
        };
        assert_eq!(
            path_options.get_indexing_options().unwrap().tokenizer(),
            "PATH_HIERARCHY_V1_FIELD_1"
        );
        Ok(())
    }

    #[test]
    fn test_write_path_hierarchy_terms() -> Result<()> {
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field(Field::from_field_id(0), "0"),
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(1), "1", false),
            crate::SingleFieldLayout::new_with_field_and_bm25(Field::from_field_id(2), "2", true),
        ]);
        let index_fields = [
            standard_index_field(true),
            index_field(
                kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1,
                false,
                Some("/"),
            ),
            index_field(
                kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1,
                true,
                Some("::"),
            ),
        ];
        let mut writer =
            TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_fields([
            (Field::from_field_id(0), "target"),
            (Field::from_field_id(1), "/src/pkg/main.go"),
            (Field::from_field_id(2), "database::table::column"),
        ])?;
        let index = writer.finalize()?;
        let schema = index.schema();
        let FieldType::Str(scored_path_options) =
            schema.get_field_entry(Field::from_field_id(2)).field_type()
        else {
            unreachable!()
        };
        let scored_path_indexing = scored_path_options.get_indexing_options().unwrap();
        assert!(scored_path_indexing.fieldnorms());
        assert_eq!(
            scored_path_indexing.index_option(),
            IndexRecordOption::WithFreqsAndPositions
        );
        let searcher = index.reader()?.searcher();

        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(
                Field::from_field_id(1),
                "/src",
            ))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(
                Field::from_field_id(1),
                "/src/pkg/main.go",
            ))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(
                Field::from_field_id(2),
                "database::table",
            ))?,
            1
        );
        Ok(())
    }

    #[test]
    fn test_write_native_physical_values() -> Result<()> {
        let field_types = [
            crate::PhysicalFieldType::I64,
            crate::PhysicalFieldType::U64,
            crate::PhysicalFieldType::F64,
            crate::PhysicalFieldType::Bytes,
        ];
        let layout = crate::IndexLayout::new(
            field_types
                .iter()
                .enumerate()
                .map(|(field_idx, field_type)| {
                    crate::SingleFieldLayout::new_with_field_type_and_bm25(
                        Field::from_field_id(field_idx as u32),
                        field_idx.to_string(),
                        *field_type,
                        false,
                    )
                })
                .collect(),
        );
        let index_fields = field_types
            .iter()
            .map(|_| {
                index_field(
                    kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerExactValueV1,
                    false,
                    None,
                )
            })
            .collect::<Vec<_>>();
        let mut writer =
            TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_values([
            (Field::from_field_id(0), OwnedValue::I64(-7)),
            (Field::from_field_id(1), OwnedValue::U64(7)),
            (Field::from_field_id(2), OwnedValue::F64(7.5)),
            (
                Field::from_field_id(3),
                OwnedValue::Bytes(b"seven".to_vec()),
            ),
        ])?;

        let index = writer.finalize()?;
        layout.validate_schema(&index.schema())?;
        let searcher = index.reader()?.searcher();
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_i64(Field::from_field_id(0), -7))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_u64(Field::from_field_id(1), 7))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_f64(Field::from_field_id(2), 7.5,))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_bytes(
                Field::from_field_id(3),
                b"seven",
            ))?,
            1
        );
        Ok(())
    }

    #[test]
    fn test_type_mismatch_error_does_not_expose_value() -> Result<()> {
        let field = Field::from_field_id(0);
        let layout = crate::IndexLayout::new(vec![
            crate::SingleFieldLayout::new_with_field_type_and_bm25(
                field,
                "0",
                crate::PhysicalFieldType::I64,
                false,
            ),
        ]);
        let index_fields = vec![index_field(
            kvenginepb::fts::FullTextAnalyzerType::FtsAnalyzerExactValueV1,
            false,
            None,
        )];
        let mut writer =
            TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;

        for (value, payload) in [
            (
                OwnedValue::Str("sensitive-text".to_owned()),
                "sensitive-text",
            ),
            (
                OwnedValue::Bytes(b"sensitive-bytes".to_vec()),
                "sensitive-bytes",
            ),
        ] {
            let error = writer
                .add_document_values([(field, value)])
                .unwrap_err()
                .to_string();
            assert!(error.contains("does not match physical field type"));
            assert!(!error.contains(payload));
        }

        let error = writer
            .add_document_values([(field, OwnedValue::Bool(true))])
            .unwrap_err()
            .to_string();
        assert!(error.contains("unsupported Tantivy value"));
        assert!(!error.contains("true"));
        Ok(())
    }

    #[test]
    fn test_build_schema_legacy_tokenizer_keeps_bm25_enabled() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let schema = build_schema("STANDARD_V1", &layout)?;
        let field_entry = schema.get_field_entry(layout.first_field().field());
        let FieldType::Str(text_options) = field_entry.field_type() else {
            panic!("expected text field");
        };
        let indexing = text_options.get_indexing_options().unwrap();
        assert_eq!(indexing.tokenizer(), "STANDARD_V1");
        assert!(indexing.fieldnorms());
        assert_eq!(
            indexing.index_option(),
            IndexRecordOption::WithFreqsAndPositions
        );
        Ok(())
    }
}