xz-memory-engine 0.2.0

//! FTS5-backed [`IndexSearcher`] implementation.
//!
//! Provides full-text keyword search using SQLite's FTS5 virtual table
//! with BM25 relevance ranking.  The FTS5 index stores entries directly
//! — entries must be inserted into the virtual table separately from the
//! [`EntryStore`] (e.g. by a higher-level synchronisation layer).

use async_trait::async_trait;
use sqlx::{Row, SqlitePool};
use xz_memory_core::{Entry, IndexSearcher, ScoredEntry, SearchOptions, StoreError};

/// FTS5-backed index searcher using SQLite's full-text search engine.
///
/// Stores entries directly in an FTS5 virtual table with BM25 ranking.
/// The table is created on construction with columns `entry_id`, `partition`,
/// `body`, and `recorded_at` — only `body` is tokenised for search.
///
/// # Examples
///
/// ```
/// use xz_memory_engine::backends::fts5::Fts5IndexSearcher;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let pool = sqlx::SqlitePool::connect("sqlite::memory:").await?;
/// let searcher = Fts5IndexSearcher::new(pool, "fts_entries").await?;
/// # Ok(())
/// # }
/// ```
pub struct Fts5IndexSearcher {
    pool: SqlitePool,
    table_name: String,
}

impl Fts5IndexSearcher {
    /// Create a new [`Fts5IndexSearcher`] backed by an FTS5 virtual table
    /// named `table_name`.
    ///
    /// The table is created automatically with `tokenize='unicode61'` on the
    /// `body` column.  Other columns (`entry_id`, `partition`, `recorded_at`)
    /// are stored but not tokenised, so they are invisible to `MATCH` queries.
    pub async fn new(pool: SqlitePool, table_name: &str) -> Result<Self, StoreError> {
        let searcher = Self { pool, table_name: table_name.to_string() };
        searcher.create_table().await?;
        Ok(searcher)
    }

    /// Create the FTS5 virtual table if it doesn't already exist.
    async fn create_table(&self) -> Result<(), StoreError> {
        let sql = format!(
            "CREATE VIRTUAL TABLE IF NOT EXISTS {table} USING fts5(\
                entry_id, \
                partition, \
                body, \
                recorded_at, \
                tokenize='unicode61'\
            )",
            table = self.table_name
        );
        sqlx::query(&sql)
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }

    /// Index (or replace) an entry so it becomes searchable via [`IndexSearcher::search`].
    ///
    /// Call this after writing durable storage (`EntryStore::append`) so the
    /// FTS virtual table stays in sync. Uses delete-then-insert on `entry_id`.
    pub async fn index_entry(&self, entry: &Entry) -> Result<(), StoreError> {
        self.remove_entry(&entry.id).await?;
        let sql = format!(
            "INSERT INTO {table} (entry_id, partition, body, recorded_at) VALUES (?, ?, ?, ?)",
            table = self.table_name
        );
        sqlx::query(&sql)
            .bind(&entry.id)
            .bind(&entry.partition)
            .bind(&entry.body)
            .bind(entry.recorded_at as i64)
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }

    /// Remove an entry from the FTS index by id (no-op if missing).
    pub async fn remove_entry(&self, id: &str) -> Result<(), StoreError> {
        let sql = format!(
            "DELETE FROM {table} WHERE entry_id = ?",
            table = self.table_name
        );
        sqlx::query(&sql)
            .bind(id)
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }
}

#[async_trait]
impl IndexSearcher for Fts5IndexSearcher {
    async fn search(
        &self,
        partitions: &[String],
        query: &str,
        opts: &SearchOptions,
    ) -> Result<Vec<ScoredEntry>, StoreError> {
        // Empty query → empty results (not an error)
        if query.trim().is_empty() {
            return Ok(vec![]);
        }

        // Build SQL with dynamic partition filter.
        // Partition names originate from internal code, not user input,
        // but we parameterise them anyway to follow sqlx best-practice.
        let partition_filter = if partitions.is_empty() {
            String::new()
        } else {
            let placeholders: Vec<&str> = vec!["?"; partitions.len()];
            format!(" AND partition IN ({})", placeholders.join(", "))
        };

        let sql = format!(
            "SELECT entry_id, partition, body, recorded_at, \
                    bm25({table}) as rank \
             FROM {table} \
             WHERE {table} MATCH ?{partition_filter} \
             ORDER BY rank \
             LIMIT ?",
            table = self.table_name,
            partition_filter = partition_filter,
        );

        let mut q = sqlx::query(&sql).bind(query);
        for p in partitions {
            q = q.bind(p);
        }
        q = q.bind(opts.limit as i64);

        let rows = q.fetch_all(&self.pool).await.map_err(|e| StoreError::Backend(e.to_string()))?;

        // Empty results → empty vec (not an error)
        if rows.is_empty() {
            return Ok(vec![]);
        }

        // Collect raw BM25 scores for normalization
        let mut raw_scores: Vec<f64> = Vec::with_capacity(rows.len());
        let mut raw_entries: Vec<Entry> = Vec::with_capacity(rows.len());

        for row in &rows {
            raw_entries.push(Entry {
                id: row.get("entry_id"),
                partition: row.get("partition"),
                body: row.get("body"),
                recorded_at: row.get::<i64, _>("recorded_at") as u64,
            });
            raw_scores.push(row.get("rank"));
        }

        // Normalize BM25 scores to [0.0, 1.0] range.
        // BM25 returns negative values where higher is more relevant.
        // We use min-max normalization so the best score maps to 1.0.
        let normalized = normalize_scores(&raw_scores);

        let mut results: Vec<ScoredEntry> = raw_entries
            .into_iter()
            .zip(normalized.into_iter())
            .map(|(entry, relevance)| ScoredEntry { entry, relevance })
            .collect();

        // Filter by minimum relevance threshold
        if let Some(min_rel) = opts.min_relevance {
            results.retain(|se| se.relevance >= min_rel);
        }

        // Truncate to requested limit (already limited in SQL, but filter step may
        // have reduced count below limit — re-apply for safety)
        results.truncate(opts.limit);

        Ok(results)
    }
}

/// Normalize a slice of BM25 scores into the [0.0, 1.0] range using
/// min-max scaling.
///
/// Single-element or uniform-score slices map every element to 1.0.
fn normalize_scores(scores: &[f64]) -> Vec<f32> {
    if scores.is_empty() {
        return vec![];
    }
    if scores.len() == 1 {
        return vec![1.0_f32];
    }

    let min_score = scores.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

    let range = max_score - min_score;
    if range.abs() < f64::EPSILON {
        return vec![1.0_f32; scores.len()];
    }

    scores.iter().map(|&s| ((s - min_score) / range) as f32).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use xz_memory_core::SearchOptions;

    /// Helper: create an in-memory pool and FTS5 searcher.
    async fn new_searcher(table_name: &str) -> (SqlitePool, Fts5IndexSearcher) {
        let pool = SqlitePool::connect("sqlite::memory:").await.expect("failed to create pool");
        let searcher = Fts5IndexSearcher::new(pool.clone(), table_name)
            .await
            .expect("failed to create searcher");
        (pool, searcher)
    }

    /// Helper: insert an entry directly into the FTS5 virtual table.
    async fn insert_entry(
        pool: &SqlitePool,
        table_name: &str,
        id: &str,
        partition: &str,
        body: &str,
        recorded_at: u64,
    ) {
        let sql = format!(
            "INSERT INTO {table} (entry_id, partition, body, recorded_at) VALUES (?, ?, ?, ?)",
            table = table_name
        );
        sqlx::query(&sql)
            .bind(id)
            .bind(partition)
            .bind(body)
            .bind(recorded_at as i64)
            .execute(pool)
            .await
            .expect("failed to insert test entry");
    }

    // ── Basic search tests ──────────────────────────────────────────

    #[tokio::test]
    async fn search_with_entries_returns_results() {
        let (pool, searcher) = new_searcher("test_fts_1").await;

        insert_entry(&pool, "test_fts_1", "e1", "facts", "Tokyo is the capital of Japan", 1000)
            .await;
        insert_entry(&pool, "test_fts_1", "e2", "facts", "Paris is the capital of France", 2000)
            .await;
        insert_entry(&pool, "test_fts_1", "e3", "other", "random data not matching", 3000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results =
            searcher.search(&["facts".to_string()], "Tokyo", &opts).await.expect("search failed");

        assert!(!results.is_empty(), "expected non-empty results");
        assert_eq!(results.len(), 1, "should match exactly one entry");
        assert_eq!(results[0].entry.id, "e1");
        assert!(
            results[0].relevance >= 0.0 && results[0].relevance <= 1.0,
            "relevance {} must be in [0.0, 1.0]",
            results[0].relevance
        );
    }

    #[tokio::test]
    async fn search_multiple_matches_returns_ranked() {
        let (pool, searcher) = new_searcher("test_fts_2").await;

        insert_entry(&pool, "test_fts_2", "e1", "facts", "machine learning is powerful", 1000)
            .await;
        insert_entry(
            &pool,
            "test_fts_2",
            "e2",
            "facts",
            "deep learning with machine techniques",
            2000,
        )
        .await;
        insert_entry(&pool, "test_fts_2", "e3", "facts", "learning to program in Rust", 3000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string()], "machine learning", &opts)
            .await
            .expect("search failed");

        assert_eq!(results.len(), 2, "should match 2 entries");
        // All scores should be in [0.0, 1.0]
        for se in &results {
            assert!(
                se.relevance >= 0.0 && se.relevance <= 1.0,
                "relevance {} out of range",
                se.relevance
            );
        }
    }

    #[tokio::test]
    async fn empty_database_returns_empty_vec() {
        let (_pool, searcher) = new_searcher("test_fts_empty").await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results =
            searcher.search(&["facts".to_string()], "Tokyo", &opts).await.expect("search failed");

        assert!(results.is_empty(), "expected empty vec for empty database");
    }

    #[tokio::test]
    async fn empty_query_returns_empty_vec() {
        let (pool, searcher) = new_searcher("test_fts_empty_query").await;

        insert_entry(&pool, "test_fts_empty_query", "e1", "facts", "some data", 1000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string()], "", &opts)
            .await
            .expect("search on empty query should not error");

        assert!(results.is_empty(), "expected empty vec for empty query");
    }

    #[tokio::test]
    async fn whitespace_only_query_returns_empty_vec() {
        let (pool, searcher) = new_searcher("test_fts_ws_query").await;

        insert_entry(&pool, "test_fts_ws_query", "e1", "facts", "some data", 1000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string()], "   \t  ", &opts)
            .await
            .expect("search on whitespace query should not error");

        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn no_match_returns_empty_vec() {
        let (pool, searcher) = new_searcher("test_fts_no_match").await;

        insert_entry(&pool, "test_fts_no_match", "e1", "facts", "apples and oranges", 1000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string()], "zzzzzzzzzz", &opts)
            .await
            .expect("search failed");

        assert!(results.is_empty(), "no-match query should return empty vec");
    }

    #[tokio::test]
    async fn unknown_table_name_errors_gracefully() {
        // Create a pool but DON'T create the FTS table — simulate a table
        // that was created by something else with a different schema.
        let pool = SqlitePool::connect("sqlite::memory:").await.expect("pool creation");

        // Create a regular (non-FTS5) table with the same name
        sqlx::query("CREATE TABLE bad_table (id INTEGER PRIMARY KEY, data TEXT)")
            .execute(&pool)
            .await
            .expect("regular table creation");

        // Now try to create the FTS5 searcher on the same name.
        // The IF NOT EXISTS will silently skip creation, but the
        // virtual-table query during search will fail.
        let searcher = Fts5IndexSearcher::new(pool, "bad_table").await;
        assert!(searcher.is_ok(), "IF NOT EXISTS should not error on existing non-FTS table");

        let searcher = searcher.unwrap();
        let opts = SearchOptions { limit: 10, min_relevance: None };
        let result = searcher.search(&["facts".to_string()], "hello", &opts).await;

        assert!(result.is_err(), "search on non-FTS table should error");
    }

    // ── Partition filtering tests ───────────────────────────────────

    #[tokio::test]
    async fn partition_filter_excludes_other_partition() {
        let (pool, searcher) = new_searcher("test_fts_partition").await;

        insert_entry(&pool, "test_fts_partition", "e1", "facts", "machine learning rocks", 1000)
            .await;
        insert_entry(
            &pool,
            "test_fts_partition",
            "e2",
            "projects",
            "machine learning projects",
            2000,
        )
        .await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results =
            searcher.search(&["facts".to_string()], "machine", &opts).await.expect("search failed");

        assert_eq!(results.len(), 1, "should only return facts partition entry");
        assert_eq!(results[0].entry.id, "e1");
    }

    #[tokio::test]
    async fn multiple_partitions_returns_all_matching() {
        let (pool, searcher) = new_searcher("test_fts_multi_part").await;

        insert_entry(&pool, "test_fts_multi_part", "e1", "facts", "machine learning", 1000).await;
        insert_entry(&pool, "test_fts_multi_part", "e2", "projects", "deep learning", 2000).await;
        insert_entry(&pool, "test_fts_multi_part", "e3", "other", "learning rust", 3000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string(), "projects".to_string()], "learning", &opts)
            .await
            .expect("search failed");

        assert_eq!(results.len(), 2, "should match facts + projects");
    }

    // ── Limit and relevance threshold tests ─────────────────────────

    #[tokio::test]
    async fn respect_search_limit() {
        let (pool, searcher) = new_searcher("test_fts_limit").await;

        for i in 0..5 {
            insert_entry(
                &pool,
                "test_fts_limit",
                &format!("e{i}"),
                "facts",
                &format!("learning rust tip {i}"),
                i * 100,
            )
            .await;
        }

        let opts = SearchOptions { limit: 2, min_relevance: None };
        let results = searcher
            .search(&["facts".to_string()], "learning", &opts)
            .await
            .expect("search failed");

        assert_eq!(results.len(), 2, "should respect limit=2");
    }

    #[tokio::test]
    async fn min_relevance_filters_low_scores() {
        let (pool, searcher) = new_searcher("test_fts_min_rel").await;

        insert_entry(&pool, "test_fts_min_rel", "e1", "facts", "rust programming language", 1000)
            .await;
        insert_entry(&pool, "test_fts_min_rel", "e2", "facts", "python scripting language", 2000)
            .await;
        insert_entry(&pool, "test_fts_min_rel", "e3", "facts", "go compiler toolchain", 3000).await;

        // Search for "rust" — e1 should match strongly, e2/e3 weakly or not at all
        let opts = SearchOptions { limit: 10, min_relevance: Some(0.3) };
        let results =
            searcher.search(&["facts".to_string()], "rust", &opts).await.expect("search failed");

        for se in &results {
            assert!(se.relevance >= 0.3, "relevance {} must be >= min_relevance 0.3", se.relevance);
        }
    }

    // ── Normalization edge cases ────────────────────────────────────

    #[tokio::test]
    async fn single_result_has_relevance_one() {
        let (pool, searcher) = new_searcher("test_fts_single").await;

        insert_entry(&pool, "test_fts_single", "e1", "facts", "exact keyword match", 1000).await;

        let opts = SearchOptions { limit: 10, min_relevance: None };
        let results =
            searcher.search(&["facts".to_string()], "keyword", &opts).await.expect("search failed");

        assert_eq!(results.len(), 1);
        assert!(
            (results[0].relevance - 1.0).abs() < f32::EPSILON,
            "single result should have relevance 1.0, got {}",
            results[0].relevance
        );
    }

    #[test]
    fn normalize_scores_empty() {
        let result = normalize_scores(&[]);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_scores_single() {
        let result = normalize_scores(&[-5.0]);
        assert_eq!(result, vec![1.0_f32]);
    }

    #[test]
    fn normalize_scores_uniform() {
        let result = normalize_scores(&[-3.0, -3.0, -3.0]);
        assert_eq!(result, vec![1.0_f32, 1.0_f32, 1.0_f32]);
    }

    #[test]
    fn normalize_scores_varied() {
        let result = normalize_scores(&[-10.0, -5.0, 0.0]);
        assert_eq!(result.len(), 3);
        assert!((result[0] - 0.0).abs() < 0.001, "worst → 0.0");
        assert!((result[1] - 0.5).abs() < 0.001, "middle → 0.5");
        assert!((result[2] - 1.0).abs() < 0.001, "best → 1.0");
    }
}