ragrig 0.9.3

//! Vector store abstraction for chunk persistence and hybrid search.
//!
//! The [`VectorStore`] trait decouples the RAG pipeline from any specific
//! storage backend.  Two implementations are provided behind feature flags:
//!
//! - `internal` (default) — pure Rust, zero native deps, MessagePack on disk
//! - `lancedb` — LanceDB-backed hybrid BM25 + vector search

use crate::types::{DocumentChunk, RrfScore, SourceFile};
use anyhow::Result;
use async_trait::async_trait;
use std::collections::{HashMap, HashSet};
use std::path::Path;
#[cfg(any(feature = "internal", feature = "lancedb"))]
use std::path::PathBuf;

/// A single chunk with its embedding, ready to be stored.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "internal", derive(serde::Serialize, serde::Deserialize))]
pub struct StoredChunk {
    pub text: String,
    pub source_file: SourceFile,
    pub vector: Vec<f32>,
}

/// Result of a hybrid search.
#[derive(Clone, Debug)]
pub struct ScoredChunk {
    pub score: RrfScore,
    pub chunk: DocumentChunk,
}

// ── VectorStore trait ─────────────────────────────────────────────────────

/// Backend-agnostic chunk storage with hybrid BM25 + vector search.
#[async_trait]
pub trait VectorStore: Send + Sync + std::fmt::Debug {
    /// Clone this store into a new heap allocation.
    fn clone_box(&self) -> Box<dyn VectorStore>;

    /// Insert chunks along with their pre-computed embedding vectors.
    async fn insert(&self, chunks: Vec<StoredChunk>) -> Result<()>;

    /// Hybrid search: cosine similarity fused with BM25 via RRF.
    async fn search(
        &self,
        query_vec: &[f32],
        query_text: &str,
        top_k: usize,
        threshold: f64,
    ) -> Result<Vec<ScoredChunk>>;

    /// Remove all chunks belonging to `source_file`.
    async fn delete_by_source(&self, source: &str) -> Result<()>;

    /// Total number of stored chunks.
    fn len(&self) -> usize;

    /// All unique source file names currently in the store.
    fn sources(&self) -> HashSet<SourceFile>;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl Clone for Box<dyn VectorStore> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

// ── Internal store (feature = "internal") ─────────────────────────────────

#[cfg(feature = "internal")]
mod brute_force {
    use super::*;
    use std::collections::HashMap;
    use std::path::Path;

    #[derive(Debug)]
    pub struct BruteForceStore {
        pub(super) inner: std::sync::Mutex<BruteForceInner>,
        pub(super) path: PathBuf,
    }

    #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
    pub struct BruteForceInner {
        pub chunks: Vec<StoredChunk>,
    }

    impl BruteForceStore {
        fn store_path(folder: &Path) -> PathBuf {
            folder.join(".ragrig_store")
        }

        pub fn open_or_create(folder: &Path) -> Result<BruteForceStore> {
            let path = Self::store_path(folder);
            let inner = if path.exists() {
                let bytes = std::fs::read(&path)?;
                rmp_serde::from_slice(&bytes).map_err(|_| {
                    anyhow::anyhow!(crate::RagrigError::StoreCorrupt {
                        path: path.to_string_lossy().into_owned(),
                    })
                })?
            } else {
                BruteForceInner { chunks: Vec::new() }
            };
            Ok(BruteForceStore {
                inner: std::sync::Mutex::new(inner),
                path,
            })
        }

        pub fn save(&self) -> Result<()> {
            let inner = self.inner.lock().unwrap();
            let bytes = rmp_serde::to_vec(&*inner)?;
            std::fs::write(&self.path, &bytes)?;
            Ok(())
        }
    }

    impl Clone for BruteForceStore {
        fn clone(&self) -> Self {
            Self {
                inner: std::sync::Mutex::new(self.inner.lock().unwrap().clone()),
                path: self.path.clone(),
            }
        }
    }

    #[async_trait]
    impl VectorStore for BruteForceStore {
        fn clone_box(&self) -> Box<dyn VectorStore> {
            Box::new(self.clone())
        }

        async fn insert(&self, chunks: Vec<StoredChunk>) -> Result<()> {
            let n = chunks.len();
            {
                let mut inner = self.inner.lock().unwrap();
                let new_sources: HashSet<SourceFile> =
                    chunks.iter().map(|c| c.source_file.clone()).collect();
                inner.chunks.retain(|c| !new_sources.contains(&c.source_file));
                inner.chunks.extend(chunks);
            }
            self.save()?;
            log::info!("Inserted {} chunks into internal store.", n);
            Ok(())
        }

        async fn search(
            &self,
            query_vec: &[f32],
            query_text: &str,
            top_k: usize,
            threshold: f64,
        ) -> Result<Vec<ScoredChunk>> {
            let inner = self.inner.lock().unwrap();
            Ok(hybrid_search(
                &inner.chunks,
                query_vec,
                query_text,
                top_k,
                threshold,
            ))
        }

        async fn delete_by_source(&self, source: &str) -> Result<()> {
            {
                let mut inner = self.inner.lock().unwrap();
                inner.chunks.retain(|c| c.source_file != source);
            }
            self.save()?;
            Ok(())
        }

        fn len(&self) -> usize {
            self.inner.lock().unwrap().chunks.len()
        }

        fn sources(&self) -> HashSet<SourceFile> {
            self.inner
                .lock()
                .unwrap()
                .chunks
                .iter()
                .map(|c| c.source_file.clone())
                .collect()
        }
    }

    // ── Search engine ──────────────────────────────────────────────────

    fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
        let (dot, norm_a, norm_b) = a.iter().zip(b.iter()).fold(
            (0.0f64, 0.0f64, 0.0f64),
            |(d, na, nb), (&x, &y)| {
                let (x, y) = (x as f64, y as f64);
                (d + x * y, na + x * x, nb + y * y)
            },
        );
        let denom = (norm_a.sqrt() * norm_b.sqrt()).max(1e-12);
        (dot / denom).clamp(-1.0, 1.0)
    }

    fn tokenize(text: &str) -> Vec<String> {
        text.to_lowercase()
            .split(|c: char| !c.is_alphanumeric())
            .filter(|t| !t.is_empty() && t.len() >= 2)
            .map(|t| t.to_string())
            .collect()
    }

    /// Okapi BM25 index with standard parameters (k1 = 1.5, b = 0.75).
    ///
    /// Robertson, Walker, Jones, Hancock-Beaulieu & Gatford (1994).
    /// "Okapi at TREC-3."  *Proceedings of TREC-3*, NIST.
    struct Bm25Index {
        doc_freqs: HashMap<String, usize>,
        doc_tfs: Vec<HashMap<String, usize>>,
        doc_lens: Vec<usize>,
        avg_doc_len: f64,
        total_docs: usize,
    }

    impl Bm25Index {
        fn build(chunks: &[StoredChunk]) -> Self {
            let total_docs = chunks.len();
            let mut doc_freqs: HashMap<String, usize> = HashMap::new();
            let mut doc_tfs: Vec<HashMap<String, usize>> = Vec::with_capacity(total_docs);
            let mut doc_lens: Vec<usize> = Vec::with_capacity(total_docs);

            for chunk in chunks {
                let tokens = tokenize(&chunk.text);
                doc_lens.push(tokens.len());
                let mut tf: HashMap<String, usize> = HashMap::new();
                for t in &tokens {
                    *tf.entry(t.clone()).or_insert(0) += 1;
                }
                for t in tf.keys() {
                    *doc_freqs.entry(t.clone()).or_insert(0) += 1;
                }
                doc_tfs.push(tf);
            }

            let avg_doc_len = if total_docs > 0 {
                doc_lens.iter().sum::<usize>() as f64 / total_docs as f64
            } else {
                1.0
            };

            Self {
                doc_freqs,
                doc_tfs,
                doc_lens,
                avg_doc_len,
                total_docs,
            }
        }

        fn score_all(&self, query_tokens: &[String]) -> Vec<(usize, f64)> {
            const K1: f64 = 1.5;
            const B: f64 = 0.75;
            const IDF_SMOOTH: f64 = 0.5;

            let n = self.total_docs as f64;
            let mut scores: Vec<(usize, f64)> = Vec::with_capacity(self.total_docs);

            for (doc_idx, tf_map) in self.doc_tfs.iter().enumerate() {
                let mut score = 0.0;
                let doc_len = self.doc_lens[doc_idx] as f64;
                for qt in query_tokens {
                    let df = *self.doc_freqs.get(qt).unwrap_or(&0) as f64;
                    if df == 0.0 {
                        continue;
                    }
                    let idf = ((n - df + IDF_SMOOTH) / (df + IDF_SMOOTH) + 1.0).ln();
                    let tf = *tf_map.get(qt).unwrap_or(&0) as f64;
                    let numerator = tf * (K1 + 1.0);
                    let denominator =
                        tf + K1 * (1.0 - B + B * doc_len / self.avg_doc_len);
                    score += idf * numerator / denominator;
                }
                scores.push((doc_idx, score));
            }
            scores
        }
    }

    /// Reciprocal Rank Fusion (k = 60).
    ///
    /// Cormack, Clarke & Buettcher (2009).  "Reciprocal Rank Fusion
    /// Outperforms Condorcet and Individual Rank Learning Methods."
    /// *Proceedings of SIGIR '09*, pp. 758–759.  ACM.
    fn rrf_fusion(
        vec_ranked: &[(usize, f64)],
        bm25_ranked: &[(usize, f64)],
        k: f64,
    ) -> Vec<(usize, f64)> {
        let mut fusion: HashMap<usize, f64> = HashMap::new();
        for (rank, (doc_idx, _)) in vec_ranked.iter().enumerate() {
            *fusion.entry(*doc_idx).or_insert(0.0) += 1.0 / (k + rank as f64 + 1.0);
        }
        for (rank, (doc_idx, _)) in bm25_ranked.iter().enumerate() {
            *fusion.entry(*doc_idx).or_insert(0.0) += 1.0 / (k + rank as f64 + 1.0);
        }
        let mut fused: Vec<(usize, f64)> = fusion.into_iter().collect();
        fused.sort_by(|a, b| b.1.total_cmp(&a.1));
        fused
    }

    /// Hybrid retrieval pipeline: cosine similarity + BM25 fused via RRF.
    ///
    /// ```text
    /// hybrid_search(query_vec, query_text, top_k, threshold)
    /// │
    /// ├─ 1. Vector scores (cosine) ──────────────────────────────
    /// │   cosine_similarity(query_vec, chunk.vector)  →  filter >= threshold
    /// │   sort descending
    /// │
    /// ├─ 2. Text scores (BM25) ──────────────────────────────────
    /// │   tokenize(query_text)                       →  lowercase, split non-alnum, len ≥ 2
    /// │   Bm25Index::build(chunks)                    →  IDF + TF per doc
    /// │     k1 = 1.5, b = 0.75                        →  standard Okapi parameters
    /// │     score_all(query_tokens)                    →  Σ IDF × (TF·(k1+1)) / (...)
    /// │   sort descending
    /// │
    /// ├─ 3. RRF fusion (k = 60) ─────────────────────────────────
    /// │   for each rank r:  score += 1 / (k + r + 1)
    /// │   documents in both lists get contributions from each
    /// │   sort descending
    /// │
    /// └─ 4. Take top_k, map to ScoredChunk
    /// ```
    ///
    /// ## References
    ///
    /// BM25 parameters: Robertson et al. (1994), "Okapi at TREC-3."
    /// RRF fusion: Cormack, Clarke & Buettcher (2009), *SIGIR '09*.
    /// Hybrid retrieval survey: Bruch (2024), "Foundations of Vector
    ///   Retrieval," arXiv:2401.09350.
    ///
    /// The threshold gates the vector side only — chunks whose cosine
        /// similarity is below `threshold` are excluded from RRF fusion.
        /// BM25 always contributes regardless, so keyword-only matches can
        /// appear with low RRF scores (~0.01‑0.03).  Use `top_k` to control
        /// how many results reach the prompt; a smaller value (e.g. 5)
        /// keeps retrieved context tight.
    fn hybrid_search(
        chunks: &[StoredChunk],
        query_vec: &[f32],
        query_text: &str,
        top_k: usize,
        threshold: f64,
    ) -> Vec<ScoredChunk> {
        if chunks.is_empty() {
            return Vec::new();
        }

        let mut vec_scores: Vec<(usize, f64)> = chunks
            .iter()
            .enumerate()
            .map(|(i, c)| (i, cosine_similarity(query_vec, &c.vector)))
            .filter(|(_, s)| *s >= threshold)
            .collect();
        vec_scores.sort_by(|a, b| b.1.total_cmp(&a.1));

        let bm25 = Bm25Index::build(chunks);
        let query_tokens = tokenize(query_text);
        let mut bm25_scores = bm25.score_all(&query_tokens);
        bm25_scores.sort_by(|a, b| b.1.total_cmp(&a.1));

        let fused = rrf_fusion(&vec_scores, &bm25_scores, 60.0);

        fused
            .into_iter()
            .take(top_k)
            .map(|(idx, score)| {
                let chunk = &chunks[idx];
                ScoredChunk {
                    score: RrfScore::from(score),
                    chunk: DocumentChunk {
                        text: chunk.text.clone(),
                        source_file: chunk.source_file.clone(),
                    },
                }
            })
            .collect()
    }
}

#[cfg(feature = "internal")]
pub use brute_force::BruteForceStore;

// ── LanceDB store (behind "lancedb" feature) ──────────────────────────────

#[cfg(feature = "lancedb")]
pub mod lance_db_store {
    use super::*;
    use anyhow::anyhow;
    use arrow_array::builder::StringBuilder;
    use arrow_array::{
        Array, FixedSizeListArray, Float32Array, RecordBatch, StringArray,
        types::Float32Type,
    };
    use arrow_schema::{DataType, Field, Schema};
    use futures_util::TryStreamExt;
    use lance_index::scalar::FullTextSearchQuery;
    use lancedb::index::Index;
    use lancedb::index::scalar::FtsIndexBuilder;
    use lancedb::query::{QueryBase, QueryExecutionOptions};
    use std::sync::Arc;

    #[derive(Clone, Debug)]
    pub struct LanceDbStore {
        table: lancedb::Table,
    }

    impl LanceDbStore {
        pub fn table_path(folder: &Path) -> PathBuf {
            folder.join(".ragrig_lancedb")
        }

        pub async fn open_or_create(folder: &Path) -> Result<Self> {
            let path = Self::table_path(folder);
            let db = lancedb::connect(&path.to_string_lossy()).execute().await?;
            let table = match db.open_table("rag_knowledge_base").execute().await {
                Ok(t) => t,
                Err(_) => {
                    let schema = Schema::new(vec![
                        Field::new("text", DataType::Utf8, false),
                        Field::new("source_file", DataType::Utf8, false),
                        Field::new(
                            "vector",
                            DataType::FixedSizeList(
                                Arc::new(Field::new("item", DataType::Float32, true)),
                                768,
                            ),
                            false,
                        ),
                    ]);
                    let batch = RecordBatch::new_empty(Arc::new(schema));
                    let t = db
                        .create_table("rag_knowledge_base", batch)
                        .execute()
                        .await?;
                    t.create_index(&["text"], Index::FTS(FtsIndexBuilder::default()))
                        .execute()
                        .await?;
                    t
                }
            };
            Ok(Self { table })
        }
    }

    #[async_trait]
    impl VectorStore for LanceDbStore {
        fn clone_box(&self) -> Box<dyn VectorStore> {
            Box::new(self.clone())
        }

        async fn insert(&self, chunks: Vec<StoredChunk>) -> Result<()> {
            if chunks.is_empty() {
                return Ok(());
            }
            let dim = chunks[0].vector.len();
            let mut text_builder =
                StringBuilder::with_capacity(chunks.len(), chunks.len() * 256);
            let mut source_builder =
                StringBuilder::with_capacity(chunks.len(), chunks.len() * 128);
            let mut vec_flat: Vec<f32> = Vec::with_capacity(chunks.len() * dim);

            for c in &chunks {
                text_builder.append_value(&c.text);
                source_builder.append_value(&c.source_file);
                vec_flat.extend_from_slice(&c.vector);
            }

            let vector_array = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
                vec_flat
                    .chunks(dim)
                    .map(|chunk| Some(chunk.iter().map(|v| Some(*v)))),
                dim as i32,
            );

            let schema = Schema::new(vec![
                Field::new("text", DataType::Utf8, false),
                Field::new("source_file", DataType::Utf8, false),
                Field::new(
                    "vector",
                    DataType::FixedSizeList(
                        Arc::new(Field::new("item", DataType::Float32, true)),
                        dim as i32,
                    ),
                    false,
                ),
            ]);

            let batch = RecordBatch::try_new(
                Arc::new(schema),
                vec![
                    Arc::new(text_builder.finish()),
                    Arc::new(source_builder.finish()),
                    Arc::new(vector_array),
                ],
            )?;

            self.table.add(batch).execute().await?;
            Ok(())
        }

        async fn search(
            &self,
            query_vec: &[f32],
            query_text: &str,
            top_k: usize,
            threshold: f64,
        ) -> Result<Vec<ScoredChunk>> {
            let stream = self
                .table
                .query()
                .nearest_to(query_vec)?
                .full_text_search(FullTextSearchQuery::new(query_text.to_string()))
                .limit(top_k)
                .execute_hybrid(QueryExecutionOptions::default())
                .await?;

            let batches: Vec<RecordBatch> = stream.try_collect().await?;
            let mut results = Vec::new();

            for batch in &batches {
                let text_col = batch
                    .column_by_name("text")
                    .and_then(|col| col.as_any().downcast_ref::<StringArray>())
                    .ok_or_else(|| anyhow!("text column not found"))?;
                let source_col = batch
                    .column_by_name("source_file")
                    .and_then(|col| col.as_any().downcast_ref::<StringArray>())
                    .ok_or_else(|| anyhow!("source_file column not found"))?;

                let score_col: Option<&Float32Array> = batch
                    .column_by_name("_score")
                    .and_then(|col| col.as_any().downcast_ref::<Float32Array>())
                    .or_else(|| {
                        batch
                            .column_by_name("_distance")
                            .and_then(|col| col.as_any().downcast_ref::<Float32Array>())
                    });

                let has_score = batch.column_by_name("_score").is_some();

                for i in 0..batch.num_rows() {
                    let raw_score = match score_col {
                        Some(col) => col.value(i) as f64,
                        None => 1.0 / (1.0 + (results.len() + i) as f64),
                    };
                    if threshold > 0.0 {
                        if has_score && raw_score < threshold {
                            continue;
                        }
                        if !has_score && raw_score > threshold {
                            continue;
                        }
                    }
                    results.push(ScoredChunk {
                        score: RrfScore::from(raw_score),
                        chunk: DocumentChunk {
                            text: text_col.value(i).to_string(),
                            source_file: SourceFile::from(source_col.value(i).to_string()),
                        },
                    });
                }
            }

            Ok(results)
        }

        async fn delete_by_source(&self, source: &str) -> Result<()> {
            self.table
                .delete(&format!("source_file = '{}'", source))
                .await?;
            Ok(())
        }

        fn len(&self) -> usize {
            0
        }

        fn sources(&self) -> HashSet<SourceFile> {
            HashSet::new()
        }
    }
}

// ── Factory ───────────────────────────────────────────────────────────────

#[cfg(feature = "lancedb")]
pub async fn open_store(folder: &Path) -> Result<Box<dyn VectorStore>> {
    lance_db_store::LanceDbStore::open_or_create(folder)
        .await
        .map(|s| Box::new(s) as Box<dyn VectorStore>)
}

#[cfg(all(feature = "internal", not(feature = "lancedb")))]
pub async fn open_store(folder: &Path) -> Result<Box<dyn VectorStore>> {
    BruteForceStore::open_or_create(folder).map(|s| Box::new(s) as Box<dyn VectorStore>)
}

#[cfg(not(any(feature = "lancedb", feature = "internal")))]
pub async fn open_store(_folder: &Path) -> Result<Box<dyn VectorStore>> {
    anyhow::bail!(
        "No vector store backend enabled. Enable the 'internal' or 'lancedb' feature."
    )
}

/// Helper: convert embedded `(text, Vec<f32>)` pairs into `StoredChunk`s
/// keyed by source file, then insert into the store.
pub async fn embed_and_insert(
    store: &dyn VectorStore,
    embedded: Vec<(String, Vec<f32>)>,
    text_to_source: &HashMap<String, String>,
) -> Result<()> {
    let chunks: Vec<StoredChunk> = embedded
        .into_iter()
        .map(|(text, vector)| {
            let source_file = text_to_source
                .get(&text)
                .cloned()
                .unwrap_or_else(|| "unknown".to_string());
            StoredChunk {
                text,
                source_file: SourceFile::from(source_file),
                vector,
            }
        })
        .collect();
    store.insert(chunks).await
}

#[cfg(test)]
#[cfg(feature = "internal")]
mod tests {
    use super::*;
    use std::env;

    fn temp_folder() -> PathBuf {
        use std::sync::atomic::{AtomicUsize, Ordering};
        static COUNTER: AtomicUsize = AtomicUsize::new(0);
        let n = COUNTER.fetch_add(1, Ordering::Relaxed);
        let mut dir = env::temp_dir();
        dir.push(format!("ragrig_test_{}_{}", std::process::id(), n));
        let _ = std::fs::create_dir_all(&dir);
        dir
    }

    fn cleanup(dir: &Path) {
        let _ = std::fs::remove_dir_all(dir);
    }

    fn chunk(text: &str, source: &str) -> StoredChunk {
        StoredChunk {
            text: text.into(),
            source_file: source.into(),
            vector: vec![1.0f32, 2.0, 3.0],
        }
    }

    #[tokio::test]
    async fn insert_and_len() {
        let dir = temp_folder();
        let store = BruteForceStore::open_or_create(&dir).unwrap();
        assert_eq!(store.len(), 0);
        store.insert(vec![chunk("hello", "doc1")]).await.unwrap();
        assert_eq!(store.len(), 1);
        store.insert(vec![chunk("world", "doc2")]).await.unwrap();
        assert_eq!(store.len(), 2);
        cleanup(&dir);
    }

    #[tokio::test]
    async fn insert_replaces_same_source() {
        let dir = temp_folder();
        let store = BruteForceStore::open_or_create(&dir).unwrap();
        store.insert(vec![chunk("old", "doc1")]).await.unwrap();
        store.insert(vec![chunk("new", "doc1")]).await.unwrap();
        assert_eq!(store.len(), 1);
        cleanup(&dir);
    }

    #[tokio::test]
    async fn delete_by_source() {
        let dir = temp_folder();
        let store = BruteForceStore::open_or_create(&dir).unwrap();
        store
            .insert(vec![chunk("a", "src1"), chunk("b", "src2")])
            .await
            .unwrap();
        assert_eq!(store.len(), 2);
        store.delete_by_source("src1").await.unwrap();
        assert_eq!(store.len(), 1);
        let sources = store.sources();
        assert!(sources.contains("src2"));
        assert!(!sources.contains(&SourceFile::from("src1")));
        cleanup(&dir);
    }

    #[tokio::test]
    async fn search_returns_scored_results() {
        let dir = temp_folder();
        let store = BruteForceStore::open_or_create(&dir).unwrap();
        let qv = vec![1.0f32, 2.0, 3.0];
        store
            .insert(vec![
                chunk("cat", "s1"),
                chunk("dog", "s2"),
                chunk("cat dog", "s3"),
            ])
            .await
            .unwrap();
        let hits = store.search(&qv, "cat", 3, 0.0).await.unwrap();
        assert!(!hits.is_empty());
        // Exact vector match is highest score, then BM25-boosted.
        for h in &hits {
            assert!(h.score > 0.0);
            assert!(!h.chunk.text.is_empty());
            assert!(!h.chunk.source_file.0.is_empty());
        }
        cleanup(&dir);
    }

    #[tokio::test]
    async fn persistence_round_trip() {
        let dir = temp_folder();
        let store = BruteForceStore::open_or_create(&dir).unwrap();
        store.insert(vec![chunk("persist me", "src")]).await.unwrap();
        drop(store);

        let reopened = BruteForceStore::open_or_create(&dir).unwrap();
        assert_eq!(reopened.len(), 1);
        assert!(reopened.sources().contains(&SourceFile::from("src")));
        cleanup(&dir);
    }
}