ragrig 0.9.5

//! Embedding generation and document ingestion.
//!
//! This module handles chunking + embedding + storage orchestration.
//! The actual embedding work is delegated to [`crate::embed::Embedder`]
//! and storage to [`crate::store::VectorStore`].

use crate::documents::build_text_to_source;
use crate::documents::build_text_to_source_with_stats;
use crate::documents::FileIndexResult;
use crate::embed::Embedder;
use crate::parsers::{DocumentParsers, self};
use crate::store::{ScoredChunk, VectorStore, embed_and_insert};
use crate::types::{ChunkConfig, DocumentType, SourceFile};
use anyhow::{Result, anyhow};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;

// --- Path Helpers ---

/// Path to the JSON file that stores file-hash metadata for incremental updates.
pub fn get_embeddings_file_path(folder: &Path) -> PathBuf {
    folder.join(".ragrig_embeddings.json")
}

// --- Shared helpers --------------------------------------------------------

/// Walk `folder` and collect all PDF / EPUB files as `DocumentType` pairs.
pub fn scan_document_files(folder: &Path) -> Vec<(DocumentType, String)> {
    WalkDir::new(folder)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter_map(|entry| {
            let path = entry.path().to_path_buf();
            if !path.is_file() {
                return None;
            }
            let ext = path.extension()?.to_str()?;
            let doc_type = DocumentType::from_extension(ext, path.clone())?;
            let name = doc_type.file_name().to_string();
            Some((doc_type, name))
        })
        .collect()
}

// --- Public API ----------------------------------------------------------

/// Chunk and embed a set of document files, then insert into the store.
pub async fn embed_documents(
    embedder: &dyn Embedder,
    parsers: &DocumentParsers,
    config: &ChunkConfig,
    document_files: Vec<(DocumentType, String)>,
    store: &dyn VectorStore,
) -> Result<()> {
    log::info!("Parsing {} documents...", document_files.len());

    let (all_texts, text_to_source) = build_text_to_source(&document_files, parsers, config)?;

    if all_texts.is_empty() {
        return Ok(());
    }

    log::info!(
        "Generating embeddings for {} total text chunks...",
        all_texts.len()
    );

    let embedded = embedder.embed(all_texts).await?;
    embed_and_insert(store, embedded, &text_to_source).await
}

/// Walk the document folder, chunk everything, embed, and populate a fresh
/// store (the caller provides the store; call `store.delete_by_source` first
/// if you want a clean slate).
pub async fn collect_documents(
    embedder: &dyn Embedder,
    parsers: &DocumentParsers,
    folder: &Path,
    config: &ChunkConfig,
    store: &dyn VectorStore,
) -> Result<()> {
    let _stats = collect_documents_with_stats(embedder, parsers, folder, config, store).await?;
    Ok(())
}

/// Like [`collect_documents`] but also returns per-file index statistics.
/// Use this when you need a result table (e.g. the REPL `/embed index` command).
pub async fn collect_documents_with_stats(
    embedder: &dyn Embedder,
    parsers: &DocumentParsers,
    folder: &Path,
    config: &ChunkConfig,
    store: &dyn VectorStore,
) -> Result<Vec<FileIndexResult>> {
    log::info!("Scanning folder recursively: {:?}", folder);

    let document_files = scan_document_files(folder);

    log::info!(
        "Found {} document files (PDF + EPUB).",
        document_files.len()
    );

    let (all_texts, text_to_source, stats) =
        build_text_to_source_with_stats(&document_files, parsers, config)?;

    if all_texts.is_empty() {
        return Err(anyhow!("No text extracted from documents."));
    }

    log::info!(
        "Generating embeddings for {} total text chunks...",
        all_texts.len()
    );

    eprintln!("Embedding {} text chunks …", all_texts.len());
    let batch_size = 50usize;
    let mut embedded: Vec<(String, Vec<f32>)> = Vec::with_capacity(all_texts.len());
    for (batch_i, batch) in all_texts.chunks(batch_size).enumerate() {
        let done = (batch_i * batch_size).min(all_texts.len());
        eprint!(
            "\r  [embedded {}/{} chunks] …                 ",
            done + batch.len(),
            all_texts.len()
        );
        let batch_embedded = embedder.embed(batch.to_vec()).await?;
        embedded.extend(batch_embedded);
    }
    eprintln!("\r  [embedded {}/{} chunks] done.", all_texts.len(), all_texts.len());
    let count = embedded.len();
    embed_and_insert(store, embedded, &text_to_source).await?;

    eprintln!("Collection complete: {} chunks stored.", count);
    Ok(stats)
}

/// One-shot indexing convenience: scan `folder`, chunk everything with
/// sensible defaults, embed, and return a populated vector store.
///
/// Equivalent to calling [`DocumentParsers::new(build_parsers())`](crate::parsers::build_parsers),
/// [`ChunkConfig::default()`], [`open_store`](crate::store::open_store), and
/// [`collect_documents`] — just wrapped into one call.
pub async fn index_folder(
    folder: &Path,
    embedder: &dyn Embedder,
) -> Result<Box<dyn VectorStore>> {
    let parsers = DocumentParsers::new(parsers::build_parsers());
    let config = ChunkConfig::default();
    let store = crate::store::open_store(folder).await?;
    collect_documents(embedder, &parsers, folder, &config, &*store).await?;
    Ok(store)
}

// --- Query & Retrieval ---------------------------------------------------

/// Embed `query` and perform hybrid search against the store.
pub async fn search_similar(
    embedder: &dyn Embedder,
    top_k: usize,
    similarity_threshold: f64,
    store: &dyn VectorStore,
    query: &str,
) -> Result<Vec<ScoredChunk>> {
    let embedded = embedder.embed(vec![query.to_string()]).await?;
    let query_vec: Vec<f32> = embedded
        .first()
        .map(|(_, v)| v.clone())
        .ok_or_else(|| anyhow!("Failed to get query embedding"))?;

    store
        .search(&query_vec, query, top_k, similarity_threshold)
        .await
}

// --- Housekeeping --------------------------------------------------------

/// Remove chunks whose source file no longer exists in `current_files`.
pub async fn remove_deleted_embeddings(
    store: &dyn VectorStore,
    current_files: &[(DocumentType, String)],
) -> Result<()> {
    let current_file_names: HashSet<SourceFile> = current_files
        .iter()
        .map(|(doc_type, _)| SourceFile::from(doc_type.file_name().to_string()))
        .collect();

    // Get all sources currently in the store and delete any not in the set.
    let stored_sources = store.sources();
    for name in &stored_sources {
        if !current_file_names.contains(name) {
            log::info!("Removing chunks for deleted file: {}", name);
            store.delete_by_source(&name.0).await?;
        }
    }

    Ok(())
}

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

    #[test]
    fn scan_picks_up_pdf_and_ignores_txt() {
        let dir = std::env::temp_dir().join(format!("ragrig-scan-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("a.pdf"), b"%PDF-1.4 fake").unwrap();
        std::fs::write(dir.join("c.txt"), b"ignored").unwrap();

        let docs = scan_document_files(&dir);
        let names: Vec<&str> = docs.iter().map(|(dt, _)| dt.file_name()).collect();
        assert!(names.contains(&"a.pdf"));
        assert!(!names.contains(&"c.txt"));
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn scan_ignores_directories() {
        let dir = std::env::temp_dir().join(format!("ragrig-scan-dir-{}", std::process::id()));
        std::fs::create_dir_all(dir.join("subdir")).unwrap();
        let docs = scan_document_files(&dir);
        assert!(docs.is_empty());
        let _ = std::fs::remove_dir_all(&dir);
    }
}