tidaldb/tidal/src/text/query.rs

use tantivy::{query::Query, schema::Field};

use crate::{TidalError, schema::TextFieldType, text::index::TantivyFields};

/// Parser for text search queries. Wraps Tantivy's `QueryParser` with
/// tidalDB-specific syntax extensions.
///
/// **Default search fields:** only [`TextFieldType::Text`] fields (tokenized).
/// `Keyword` fields require explicit field scoping (e.g., `category:programming`).
///
/// **Conjunction mode:** AND by default. Multi-word queries like `rust tutorial`
/// match documents containing both terms, not either.
///
/// # Supported syntax
///
/// - Bare terms: `rust tutorial` (conjunction of rust AND tutorial)
/// - Exact phrase: `"exact phrase"` (`PhraseQuery`)
/// - Boolean AND: `jazz AND piano`
/// - Boolean OR: `jazz OR rock`
/// - Boolean NOT / exclusion: `jazz -beginner` or `jazz NOT beginner`
/// - Field-scoped: `title:jazz`
/// - Wildcard prefix: `pian*`
/// - Hashtag: `#jazz` (pre-processed to `jazz`)
pub struct TextQueryParser {
    inner: tantivy::query::QueryParser,
}

impl TextQueryParser {
    /// Create a parser configured for the given index and fields.
    ///
    /// Default fields are the subset of `fields.text_fields` with
    /// `TextFieldType::Text`. Keyword fields are not searched by default --
    /// the caller must use field-scoped syntax (e.g., `category:tech`).
    ///
    /// Sets AND as the default conjunction mode so that multi-word queries
    /// require all terms to be present.
    #[must_use]
    pub fn new(index: &tantivy::Index, fields: &TantivyFields) -> Self {
        let default_fields: Vec<Field> = fields
            .text_fields
            .iter()
            .filter(|(_, _, ft)| *ft == TextFieldType::Text)
            .map(|(_, f, _)| *f)
            .collect();

        let mut inner = tantivy::query::QueryParser::for_index(index, default_fields);
        inner.set_conjunction_by_default(); // "rust tutorial" = rust AND tutorial
        Self { inner }
    }

    /// Parse a query string into a Tantivy [`Query`].
    ///
    /// Applies hashtag pre-processing (`#jazz` -> `jazz`) before delegating
    /// to Tantivy's parser.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` with a descriptive message if the query
    /// string cannot be parsed (e.g., unbalanced quotes, unknown field names).
    pub fn parse(&self, query_str: &str) -> crate::Result<Box<dyn Query>> {
        let preprocessed = preprocess_query(query_str);
        self.inner.parse_query(&preprocessed).map_err(|e| {
            TidalError::internal("text_query_parse", format!("text query parse error: {e}"))
        })
    }
}

/// Pre-process tidalDB query strings before passing to Tantivy's `QueryParser`.
///
/// Current transformations:
/// - `#jazz` -> `jazz` (strips the `#` prefix from valid hashtags, where the
///   character immediately after `#` is ASCII alphanumeric)
///
/// A lone `#` or `# ` (hash followed by space/EOF) is left unchanged.
fn preprocess_query(query: &str) -> String {
    let mut result = String::with_capacity(query.len());
    let mut chars = query.chars().peekable();
    while let Some(ch) = chars.next() {
        if ch == '#' && chars.peek().is_some_and(char::is_ascii_alphanumeric) {
            // Valid hashtag prefix -- skip the '#', let the word through.
        } else {
            result.push(ch);
        }
    }
    result
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::collections::HashMap;

    use tantivy::{TantivyDocument, collector::TopDocs, schema::Value};

    use super::*;
    use crate::{
        schema::{EntityId, TextFieldDef, TextFieldType},
        text::index::TextIndex,
    };

    /// Helper: create an ephemeral index with "title" (Text) and "category"
    /// (Keyword), populate it with 4 documents, commit, and reload the reader.
    fn setup_index() -> TextIndex {
        let fields = vec![
            TextFieldDef {
                key: "title".to_owned(),
                field_type: TextFieldType::Text,
            },
            TextFieldDef {
                key: "category".to_owned(),
                field_type: TextFieldType::Keyword,
            },
        ];
        let idx = TextIndex::ephemeral(&fields).unwrap();

        let docs: Vec<(u64, &str, &str)> = vec![
            (1, "jazz piano beginner", "music"),
            (2, "rock guitar advanced", "music"),
            (3, "jazz violin intermediate", "music"),
            (4, "rust programming language", "tech"),
        ];

        let mut w = idx.writer_guard().unwrap();
        for (id, title, cat) in docs {
            let mut m = HashMap::new();
            m.insert("title".to_owned(), title.to_owned());
            m.insert("category".to_owned(), cat.to_owned());
            w.index_item(EntityId::new(id), &m).unwrap();
        }
        w.commit().unwrap();
        drop(w);
        idx.reader.reload().unwrap();
        idx
    }

    /// Helper: execute a query and return the matched entity IDs.
    fn search_ids(idx: &TextIndex, query: &dyn Query) -> Vec<u64> {
        let searcher = idx.reader.searcher();
        let top_docs = searcher.search(query, &TopDocs::with_limit(100)).unwrap();
        top_docs
            .iter()
            .map(|(_score, doc_addr)| {
                let doc: TantivyDocument = searcher.doc(*doc_addr).unwrap();
                doc.get_first(idx.fields().entity_id)
                    .and_then(|v| v.as_u64())
                    .unwrap()
            })
            .collect()
    }

    #[test]
    fn preprocess_removes_hashtag() {
        assert_eq!(preprocess_query("#jazz"), "jazz");
        assert_eq!(preprocess_query("#jazz #piano"), "jazz piano");
        assert_eq!(preprocess_query("jazz #piano"), "jazz piano");
        assert_eq!(preprocess_query("no-hashtag"), "no-hashtag");
        // Space after # = not a hashtag.
        assert_eq!(preprocess_query("# notag"), "# notag");
    }

    #[test]
    fn preprocess_preserves_empty_and_plain() {
        assert_eq!(preprocess_query(""), "");
        assert_eq!(preprocess_query("hello world"), "hello world");
        assert_eq!(preprocess_query("#"), "#");
    }

    #[test]
    fn parse_bare_terms_conjunction() {
        // "jazz piano" with set_conjunction_by_default() -> AND.
        // Should find entity 1 (has both jazz AND piano) but not entity 3
        // (jazz but not piano).
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("jazz piano").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(ids.contains(&1), "entity 1 should match jazz AND piano");
        assert!(!ids.contains(&2), "entity 2 should not match");
        assert!(!ids.contains(&3), "entity 3 has jazz but not piano");
        idx.close().unwrap();
    }

    #[test]
    fn parse_exact_phrase() {
        // "\"jazz piano\"" -> PhraseQuery, only entity 1 has the contiguous
        // sequence "jazz piano".
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("\"jazz piano\"").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(ids.contains(&1), "entity 1 has 'jazz piano' phrase");
        assert!(
            !ids.contains(&3),
            "entity 3 has 'jazz violin', not 'jazz piano'"
        );
        idx.close().unwrap();
    }

    #[test]
    fn parse_boolean_or() {
        // "jazz OR rock" -> entities 1, 2, 3 (all contain jazz or rock).
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("jazz OR rock").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(ids.contains(&1), "entity 1 has jazz");
        assert!(ids.contains(&2), "entity 2 has rock");
        assert!(ids.contains(&3), "entity 3 has jazz");
        assert!(!ids.contains(&4), "entity 4 has neither jazz nor rock");
        idx.close().unwrap();
    }

    #[test]
    fn parse_exclusion_minus() {
        // "jazz -beginner" -> jazz items excluding beginners.
        // Entity 3 (jazz, no beginner) should match. Entity 1 (jazz beginner) excluded.
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("jazz -beginner").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(
            ids.contains(&3),
            "entity 3 (jazz, no beginner) should match"
        );
        assert!(
            !ids.contains(&1),
            "entity 1 (jazz beginner) should be excluded"
        );
        idx.close().unwrap();
    }

    #[test]
    fn parse_field_scoped_keyword() {
        // "category:tech" -> only entity 4 (rust programming).
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("category:tech").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(ids.contains(&4), "entity 4 should match category:tech");
        assert!(!ids.contains(&1), "entity 1 has category:music, not tech");
        idx.close().unwrap();
    }

    #[test]
    fn parse_field_scoped_text() {
        // "title:guitar" -> only entity 2.
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("title:guitar").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert_eq!(ids, vec![2], "only entity 2 has 'guitar' in title");
        idx.close().unwrap();
    }

    #[test]
    fn parse_wildcard_prefix() {
        // "jaz*" -> verify parsing does not panic or produce a hard error.
        // Tantivy 0.22's QueryParser may or may not support wildcard expansion
        // depending on the index settings. We verify the parse itself succeeds
        // and, if results are returned, they are plausible.
        let idx = setup_index();
        let parser = idx.query_parser();
        let result = parser.parse("jaz*");
        // Accept either success or a specific parse error -- do not panic.
        if let Ok(q) = result {
            // Wildcard support is best-effort. If Tantivy resolves it, great;
            // if not, an empty result set is acceptable.
            let _ids = search_ids(&idx, q.as_ref());
        }
        idx.close().unwrap();
    }

    #[test]
    fn parse_hashtag() {
        // "#jazz" should produce same results as "jazz".
        let idx = setup_index();
        let parser = idx.query_parser();
        let q_hash = parser.parse("#jazz").unwrap();
        let q_bare = parser.parse("jazz").unwrap();
        let ids_hash: std::collections::HashSet<u64> =
            search_ids(&idx, q_hash.as_ref()).into_iter().collect();
        let ids_bare: std::collections::HashSet<u64> =
            search_ids(&idx, q_bare.as_ref()).into_iter().collect();
        assert_eq!(
            ids_hash, ids_bare,
            "#jazz and jazz should return same entity IDs"
        );
        idx.close().unwrap();
    }

    #[test]
    fn parse_invalid_query_returns_error() {
        let idx = TextIndex::ephemeral(&[TextFieldDef {
            key: "title".to_owned(),
            field_type: TextFieldType::Text,
        }])
        .unwrap();
        let parser = idx.query_parser();
        // An unbalanced quote is a parse error in Tantivy.
        let result = parser.parse("\"unclosed phrase");
        // Tantivy may either error or be lenient. Verify no panic either way.
        // If it does error, verify it is TidalError::Internal.
        if let Err(e) = result {
            assert!(
                e.to_string().contains("text query parse error")
                    || e.to_string().contains("internal error"),
                "error should describe parse failure, got: {e}"
            );
        }
        idx.close().unwrap();
    }

    #[test]
    fn default_fields_exclude_keyword() {
        // A bare search for "music" should not match because "category" is a
        // Keyword field and not in the default search fields.
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("music").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        // "music" does not appear in any title, and category (Keyword) is not
        // a default field, so no results.
        assert!(
            ids.is_empty(),
            "bare 'music' should not search keyword-only field; got: {ids:?}"
        );
        idx.close().unwrap();
    }

    #[test]
    fn query_parser_method_on_text_index() {
        // Verify the convenience method exists and produces a working parser.
        let idx = setup_index();
        let parser = idx.query_parser();
        let q = parser.parse("rust").unwrap();
        let ids = search_ids(&idx, q.as_ref());
        assert!(ids.contains(&4), "entity 4 has 'rust' in title");
        idx.close().unwrap();
    }
}