komp_ac/common/src/search.rs

use std::path::{Path, PathBuf};

use tantivy::schema::{
    Field, IndexRecordOption, JsonObjectOptions, Schema, Term, TextFieldIndexing, TextOptions,
    FAST, INDEXED, STORED, STRING,
};
use tantivy::tokenizer::{
    AsciiFoldingFilter, LowerCaser, NgramTokenizer, RawTokenizer, RemoveLongFilter,
    SimpleTokenizer, TextAnalyzer, TokenStream,
};
use tantivy::Index;

pub const F_PG_ID: &str = "pg_id";
pub const F_ROW_REVISION: &str = "row_revision";
pub const F_IS_ARCHIVED: &str = "is_archived";
pub const F_TABLE_DEFINITION_ID: &str = "table_definition_id";
pub const F_VERSION_NUMBER: &str = "version_number";
pub const F_TABLE_NAME: &str = "table_name";
pub const F_ROW_KEY: &str = "row_key";
pub const F_ALL_TEXT: &str = "all_text";
pub const F_DATA_WORD: &str = "data_word";
pub const F_DATA_NGRAM: &str = "data_ngram";
pub const F_DATA_EXACT: &str = "data_exact";
pub const JOURNAL_TABLE_NAME: &str = "general_ledger";
pub const ARCHIVED_ROW_KEY_PREFIX: &str = "__archive__:";
pub const SEARCH_INDEX_FORMAT_DIRECTORY: &str = "v3";

/// Root for the current on-disk index format. The format component prevents
/// an executable with new projection semantics from opening old documents.
pub fn search_index_root() -> PathBuf {
    std::env::var_os("TANTIVY_INDEX_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("./tantivy_indexes"))
        .join(SEARCH_INDEX_FORMAT_DIRECTORY)
}

pub const TOK_WORD: &str = "kw_word";
pub const TOK_NGRAM: &str = "kw_ngram";
pub const TOK_EXACT: &str = "kw_exact";

/// Returns the on-disk path for a profile search index.
pub fn search_index_path(root: &Path, profile_name: &str) -> PathBuf {
    root.join(profile_name)
}

/// Returns the unique index key for one table row inside a profile index.
pub fn search_row_key(table_name: &str, row_id: i64) -> String {
    format!("{}:{}", table_name, row_id)
}

/// Returns the unique index key for an immutable archived row version.
pub fn archived_search_row_key(
    table_definition_id: i64,
    row_id: i64,
    version_number: i64,
) -> String {
    format!(
        "{ARCHIVED_ROW_KEY_PREFIX}{table_definition_id}:{row_id}:{version_number}"
    )
}

/// Decodes an archived index key into (table definition, row, version).
pub fn parse_archived_search_row_key(row_key: &str) -> Option<(i64, i64, i64)> {
    let encoded = row_key.strip_prefix(ARCHIVED_ROW_KEY_PREFIX)?;
    let mut parts = encoded.split(':');
    let table_definition_id = parts.next()?.parse().ok()?;
    let row_id = parts.next()?.parse().ok()?;
    let version_number = parts.next()?.parse().ok()?;
    if parts.next().is_some() {
        return None;
    }
    Some((table_definition_id, row_id, version_number))
}

/// Normalizes user-entered values for exact-mode terms.
pub fn normalize_exact(input: &str) -> String {
    let trimmed = input.trim();
    if trimmed.is_empty() {
        return String::new();
    }

    let mut analyzer = exact_analyzer();
    let mut stream = analyzer.token_stream(trimmed);
    let mut out = String::with_capacity(trimmed.len());
    while let Some(token) = stream.next() {
        out.push_str(&token.text);
    }
    out
}

/// Canonicalizes an exact constraint with the same scalar spelling used by
/// the indexer. Catalog field types are public type names, not PostgreSQL type
/// names, so this deliberately matches the catalog vocabulary.
pub fn canonical_exact_search_value(input: &str, field_type: &str) -> Result<String, String> {
    let normalized_type = field_type.trim().to_ascii_lowercase();
    if matches!(normalized_type.as_str(), "numeric" | "money")
        || normalized_type.starts_with("decimal(")
    {
        return input
            .parse::<rust_decimal::Decimal>()
            .map(|value| value.normalize().to_string())
            .map_err(|error| format!("Exact numeric search value is invalid: {error}"));
    }
    if matches!(normalized_type.as_str(), "int" | "bigint")
        || normalized_type.starts_with("link(")
    {
        return input
            .parse::<i64>()
            .map(|value| value.to_string())
            .map_err(|error| format!("Exact integer search value is invalid: {error}"));
    }
    Ok(input.to_string())
}

/// Canonical spelling for a JSON number stored in a Tantivy text term.
pub fn canonical_search_number(input: &str) -> String {
    input
        .parse::<rust_decimal::Decimal>()
        .map(|value| value.normalize().to_string())
        .unwrap_or_else(|_| input.to_string())
}

/// Normalizes a column name to the JSON-key form used at index time.
pub fn normalize_column_name(column: &str) -> String {
    column.to_ascii_lowercase()
}

/// Creates the column-aware search schema.
pub fn create_search_schema() -> Schema {
    let mut schema_builder = Schema::builder();

    schema_builder.add_u64_field(F_PG_ID, INDEXED | STORED | FAST);
    schema_builder.add_u64_field(F_ROW_REVISION, STORED | FAST);
    schema_builder.add_u64_field(F_IS_ARCHIVED, INDEXED | STORED | FAST);
    schema_builder.add_u64_field(F_TABLE_DEFINITION_ID, STORED | FAST);
    schema_builder.add_u64_field(F_VERSION_NUMBER, STORED | FAST);
    schema_builder.add_text_field(F_TABLE_NAME, STRING | STORED);
    schema_builder.add_text_field(F_ROW_KEY, STRING | STORED);
    schema_builder.add_text_field(F_ALL_TEXT, text_options(TOK_WORD));

    schema_builder.add_json_field(F_DATA_WORD, json_options(TOK_WORD, true, false));
    schema_builder.add_json_field(F_DATA_NGRAM, json_options(TOK_NGRAM, true, false));
    schema_builder.add_json_field(F_DATA_EXACT, json_options(TOK_EXACT, false, false));

    schema_builder.build()
}

fn text_options(tokenizer_name: &str) -> TextOptions {
    let indexing = TextFieldIndexing::default()
        .set_tokenizer(tokenizer_name)
        .set_index_option(IndexRecordOption::WithFreqsAndPositions);

    TextOptions::default().set_indexing_options(indexing)
}

fn json_options(tokenizer_name: &str, with_positions: bool, stored: bool) -> JsonObjectOptions {
    let index_option = if with_positions {
        IndexRecordOption::WithFreqsAndPositions
    } else {
        IndexRecordOption::Basic
    };

    let indexing = TextFieldIndexing::default()
        .set_tokenizer(tokenizer_name)
        .set_index_option(index_option);

    let mut options = JsonObjectOptions::default().set_indexing_options(indexing);
    if stored {
        options = options.set_stored();
    }
    options
}

/// Registers all required tokenizers with the index.
pub fn register_tokenizers(index: &Index) -> tantivy::Result<()> {
    let tokenizer_manager = index.tokenizers();

    tokenizer_manager.register(TOK_WORD, word_analyzer());
    tokenizer_manager.register(TOK_NGRAM, ngram_analyzer()?);
    tokenizer_manager.register(TOK_EXACT, exact_analyzer());

    Ok(())
}

fn word_analyzer() -> TextAnalyzer {
    TextAnalyzer::builder(SimpleTokenizer::default())
        .filter(RemoveLongFilter::limit(80))
        .filter(LowerCaser)
        .filter(AsciiFoldingFilter)
        .build()
}

fn ngram_analyzer() -> tantivy::Result<TextAnalyzer> {
    Ok(TextAnalyzer::builder(NgramTokenizer::new(3, 3, false)?)
        .filter(RemoveLongFilter::limit(80))
        .filter(LowerCaser)
        .filter(AsciiFoldingFilter)
        .build())
}

fn exact_analyzer() -> TextAnalyzer {
    TextAnalyzer::builder(RawTokenizer::default())
        .filter(LowerCaser)
        .filter(AsciiFoldingFilter)
        .build()
}

/// Tokenizes text the same way `data_word` is indexed.
pub fn tokenize_word(text: &str) -> Vec<String> {
    tokenize_with(word_analyzer(), text)
}

/// Tokenizes text the same way `data_ngram` is indexed.
pub fn tokenize_ngram(text: &str) -> Vec<String> {
    match ngram_analyzer() {
        Ok(analyzer) => tokenize_with(analyzer, text),
        Err(_) => Vec::new(),
    }
}

fn tokenize_with(mut analyzer: TextAnalyzer, text: &str) -> Vec<String> {
    let mut stream = analyzer.token_stream(text);
    let mut out = Vec::new();
    while let Some(token) = stream.next() {
        out.push(token.text.clone());
    }
    out
}

/// Builds a term scoped to a specific JSON path within a JSON field.
pub fn json_path_term(field: Field, column: &str, text: &str) -> Term {
    let mut term = Term::from_field_json_path(field, column, false);
    term.append_type_and_str(text);
    term
}

/// Returns all required schema fields or fails loudly on mismatch.
pub struct SchemaFields {
    pub pg_id: Field,
    pub row_revision: Field,
    pub is_archived: Field,
    pub table_definition_id: Field,
    pub version_number: Field,
    pub table_name: Field,
    pub row_key: Field,
    pub all_text: Field,
    pub data_word: Field,
    pub data_ngram: Field,
    pub data_exact: Field,
}

impl SchemaFields {
    pub fn from(schema: &Schema) -> tantivy::Result<Self> {
        Ok(Self {
            pg_id: get_field(schema, F_PG_ID)?,
            row_revision: get_field(schema, F_ROW_REVISION)?,
            is_archived: get_field(schema, F_IS_ARCHIVED)?,
            table_definition_id: get_field(schema, F_TABLE_DEFINITION_ID)?,
            version_number: get_field(schema, F_VERSION_NUMBER)?,
            table_name: get_field(schema, F_TABLE_NAME)?,
            row_key: get_field(schema, F_ROW_KEY)?,
            all_text: get_field(schema, F_ALL_TEXT)?,
            data_word: get_field(schema, F_DATA_WORD)?,
            data_ngram: get_field(schema, F_DATA_NGRAM)?,
            data_exact: get_field(schema, F_DATA_EXACT)?,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{
        archived_search_row_key, canonical_exact_search_value, parse_archived_search_row_key,
    };

    #[test]
    fn archived_row_key_round_trips_stable_identity() {
        let key = archived_search_row_key(42, 7, 3);
        assert_eq!(parse_archived_search_row_key(&key), Some((42, 7, 3)));
        assert_eq!(parse_archived_search_row_key("customers:7"), None);
        assert_eq!(parse_archived_search_row_key("__archive__:42:7"), None);
    }

    #[test]
    fn exact_scalar_values_follow_catalog_field_types() {
        assert_eq!(canonical_exact_search_value("001", "int").unwrap(), "1");
        assert_eq!(canonical_exact_search_value("001", "bigint").unwrap(), "1");
        assert_eq!(canonical_exact_search_value("001", "link(adresar)").unwrap(), "1");
        assert_eq!(canonical_exact_search_value("10.50", "numeric").unwrap(), "10.5");
        assert_eq!(canonical_exact_search_value("10.50", "money").unwrap(), "10.5");
        assert_eq!(
            canonical_exact_search_value("10.50", "decimal(12, 2)").unwrap(),
            "10.5"
        );
    }
}

fn get_field(schema: &Schema, name: &str) -> tantivy::Result<Field> {
    schema.get_field(name).map_err(|e| {
        tantivy::TantivyError::SchemaError(format!("schema is missing field '{name}': {e}"))
    })
}