tidaldb/tidal/src/text/writer.rs

use std::collections::HashMap;
use std::sync::MutexGuard;

use tantivy::{IndexWriter, TantivyDocument, Term};

use crate::TidalError;
use crate::schema::EntityId;
use crate::text::index::TantivyFields;

/// Write operations on the Tantivy text index.
///
/// Created by [`TextIndex::writer_guard()`](super::TextIndex::writer_guard).
/// Holds the `MutexGuard` on the `IndexWriter`, preventing concurrent writes
/// (Tantivy enforces single-writer).
///
/// All write operations are batched in memory. Call [`commit()`](Self::commit)
/// to make changes durable and visible to new searchers.
pub struct TextIndexWriter<'a> {
    pub(crate) writer: MutexGuard<'a, IndexWriter>,
    pub(crate) fields: &'a TantivyFields,
}

impl TextIndexWriter<'_> {
    /// Index or re-index an item.
    ///
    /// Performs a delete-then-add in the same batch so updates are atomic
    /// from the reader's perspective (both become visible after `commit()`).
    /// Only metadata keys declared as text fields in the schema are indexed;
    /// other keys are silently ignored.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to add the document.
    #[tracing::instrument(skip(self, metadata), fields(%entity_id))]
    pub fn index_item(
        &mut self,
        entity_id: EntityId,
        metadata: &HashMap<String, String>,
    ) -> crate::Result<()> {
        // Delete any previous version of this document.
        let id_term = Term::from_field_u64(self.fields.entity_id, entity_id.as_u64());
        self.writer.delete_term(id_term);

        // Build new document.
        let mut doc = TantivyDocument::default();
        doc.add_u64(self.fields.entity_id, entity_id.as_u64());

        for (key, tv_field, _field_type) in &self.fields.text_fields {
            if let Some(value) = metadata.get(key) {
                doc.add_text(*tv_field, value);
            }
        }

        self.writer
            .add_document(doc)
            .map_err(|e| TidalError::Internal(format!("tantivy add_document: {e}")))?;

        Ok(())
    }

    /// Schedule deletion of a document by entity ID.
    ///
    /// The delete is applied on the next `commit()`. If the entity is not in
    /// the index, this is a no-op.
    pub fn delete_item(&mut self, entity_id: EntityId) {
        let id_term = Term::from_field_u64(self.fields.entity_id, entity_id.as_u64());
        self.writer.delete_term(id_term);
    }

    /// Commit all pending writes, storing `last_seq` in the Tantivy commit payload.
    ///
    /// After this returns, new `searcher()` instances will see the committed docs.
    /// The `last_seq` is retrievable via [`last_committed_seq()`](Self::last_committed_seq)
    /// for crash recovery.
    ///
    /// Uses Tantivy's two-phase commit: `prepare_commit()` -> `set_payload()` ->
    /// `commit()`. This is the only way to attach a payload in Tantivy 0.22.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to prepare or finalize
    /// the commit.
    #[tracing::instrument(skip(self))]
    pub fn commit(&mut self, last_seq: u64) -> crate::Result<()> {
        let mut prepared = self
            .writer
            .prepare_commit()
            .map_err(|e| TidalError::Internal(format!("tantivy prepare_commit: {e}")))?;
        prepared.set_payload(&last_seq.to_string());
        prepared
            .commit()
            .map_err(|e| TidalError::Internal(format!("tantivy commit: {e}")))?;
        Ok(())
    }

    /// Delete all documents from the index.
    ///
    /// Used by the syncer's `rebuild_from()` to start fresh before re-indexing
    /// from the entity store.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to delete documents.
    pub fn delete_all(&mut self) -> crate::Result<()> {
        self.writer
            .delete_all_documents()
            .map_err(|e| TidalError::Internal(format!("tantivy delete_all: {e}")))?;
        Ok(())
    }

    /// Read the last committed sequence number from the Tantivy index payload.
    ///
    /// Returns `0` if no commit payload exists (fresh index or first run).
    /// This is used on startup for crash recovery to determine the WAL replay
    /// point.
    #[must_use]
    pub fn last_committed_seq(index: &tantivy::Index) -> u64 {
        index
            .load_metas()
            .ok()
            .and_then(|meta| meta.payload)
            .and_then(|p| p.parse::<u64>().ok())
            .unwrap_or(0)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::schema::{TextFieldDef, TextFieldType};
    use crate::text::TextIndex;

    use tantivy::collector::TopDocs;
    use tantivy::query::QueryParser;
    use tantivy::schema::Value;

    fn sample_text_fields() -> Vec<TextFieldDef> {
        vec![
            TextFieldDef {
                key: "title".to_owned(),
                field_type: TextFieldType::Text,
            },
            TextFieldDef {
                key: "description".to_owned(),
                field_type: TextFieldType::Text,
            },
        ]
    }

    fn make_metadata(pairs: &[(&str, &str)]) -> HashMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect()
    }

    #[test]
    fn index_and_search() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "smooth jazz classics")]),
            )
            .unwrap();
            w.index_item(
                EntityId::new(2),
                &make_metadata(&[("title", "rock anthems")]),
            )
            .unwrap();
            w.index_item(
                EntityId::new(3),
                &make_metadata(&[("title", "jazz fusion favorites")]),
            )
            .unwrap();
            w.commit(10).unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();
        let top_docs = searcher.search(&query, &TopDocs::with_limit(10)).unwrap();

        assert_eq!(top_docs.len(), 2, "two docs contain 'jazz'");

        let found_ids: Vec<u64> = top_docs
            .iter()
            .map(|(_score, doc_address)| {
                let doc: TantivyDocument = searcher.doc(*doc_address).unwrap();
                doc.get_first(idx.fields().entity_id)
                    .and_then(|v| v.as_u64())
                    .unwrap()
            })
            .collect();

        assert!(found_ids.contains(&1));
        assert!(found_ids.contains(&3));
    }

    #[test]
    fn delete_removes_document() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(EntityId::new(1), &make_metadata(&[("title", "alpha")]))
                .unwrap();
            w.index_item(EntityId::new(2), &make_metadata(&[("title", "beta")]))
                .unwrap();
            w.commit(1).unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 2);

        {
            let mut w = idx.writer_guard().unwrap();
            w.delete_item(EntityId::new(1));
            w.commit(2).unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 1);
    }

    #[test]
    fn update_replaces_document() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        // Index with title A.
        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "original alpha")]),
            )
            .unwrap();
            w.commit(1).unwrap();
        }

        // Re-index same entity with title B.
        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "replacement beta")]),
            )
            .unwrap();
            w.commit(2).unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();

        // Only one document should exist.
        assert_eq!(searcher.num_docs(), 1);

        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);

        // Search for old title: not found.
        let query_old = qp.parse_query("original").unwrap();
        let results_old = searcher
            .search(&query_old, &TopDocs::with_limit(10))
            .unwrap();
        assert!(results_old.is_empty(), "old title should not be found");

        // Search for new title: found.
        let query_new = qp.parse_query("replacement").unwrap();
        let results_new = searcher
            .search(&query_new, &TopDocs::with_limit(10))
            .unwrap();
        assert_eq!(results_new.len(), 1, "new title should be found");

        let doc: TantivyDocument = searcher.doc(results_new[0].1).unwrap();
        let eid = doc
            .get_first(idx.fields().entity_id)
            .and_then(|v| v.as_u64())
            .unwrap();
        assert_eq!(eid, 1);
    }

    #[test]
    fn commit_stores_sequence() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.commit(42).unwrap();
        }

        let seq = TextIndexWriter::last_committed_seq(&idx.index);
        assert_eq!(seq, 42);
    }

    #[test]
    fn last_committed_seq_returns_zero_fresh() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();
        let seq = TextIndexWriter::last_committed_seq(&idx.index);
        assert_eq!(seq, 0);
    }

    #[test]
    fn unknown_metadata_keys_ignored() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[
                    ("title", "hello world"),
                    ("nonexistent_key", "should be ignored"),
                    ("another_bogus", "also ignored"),
                ]),
            )
            .unwrap();
            w.commit(1).unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();
        assert_eq!(searcher.num_docs(), 1, "document should be indexed");

        // The entity_id should still be readable.
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("hello").unwrap();
        let results = searcher.search(&query, &TopDocs::with_limit(10)).unwrap();
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn delete_nonexistent_is_noop() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            // Delete entity that was never indexed — should not panic or error.
            w.delete_item(EntityId::new(999));
            w.commit(1).unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 0);
    }
}