xz-rag 0.1.1

use crate::error::RagError;
use crate::pipeline::channel::{ChannelConfig, ChannelType};
use crate::types::chunk::Chunk;
use crate::types::retrieval::{RetrievedChunk, StructuredFilter};
use tantivy::collector::TopDocs;
use tantivy::query::{BooleanQuery, Occur, QueryParser, TermQuery};
use tantivy::schema::{IndexRecordOption, STRING, Schema, TextFieldIndexing, TextOptions, Value};
use tantivy::{Index, ReloadPolicy, Term, doc};

/// Default BM25 k1 parameter (term frequency saturation).
const DEFAULT_K1: f32 = 1.2;

/// Default BM25 b parameter (document length normalization).
const DEFAULT_B: f32 = 0.75;

/// BM25 full-text search channel executor backed by a Tantivy in-memory index.
///
/// This channel uses Tantivy's default tokenizer which performs:
/// - Unicode word segmentation
/// - Lowercasing
/// - No stemming (as required)
///
/// ## BM25 Parameters
/// The `k1` and `b` parameters control BM25 scoring behavior and can be
/// configured per-query via [`ChannelConfig::params`]:
/// - `k1` (float, default `1.2`): controls term frequency saturation
/// - `b` (float, default `0.75`): controls document length normalization
///
/// ## Namespace Support
/// Each chunk's `metadata.namespace` is stored in the index and can be
/// filtered at query time via the `namespace` parameter of [`Self::execute`].
///
/// ## Index Lifecycle
/// The index is entirely in-memory. Documents are added via
/// [`index_documents`](Self::index_documents) before searching. Currently
/// the executor does not support incremental updates or deletion — call
/// `index_documents` once with the complete document set.
pub struct Bm25ChannelExecutor {
    /// Tantivy in-memory index.
    index: Index,
    /// Schema field for chunk identifier (stored, not indexed).
    chunk_id_field: tantivy::schema::Field,
    /// Schema field for document identifier (stored, not indexed).
    document_id_field: tantivy::schema::Field,
    /// Schema field for full-text content (indexed and stored).
    text_field: tantivy::schema::Field,
    /// Schema field for namespace (indexed for filtering, stored).
    namespace_field: tantivy::schema::Field,
    /// Current BM25 k1 parameter.
    k1: f32,
    /// Current BM25 b parameter.
    b: f32,
    /// Number of documents currently indexed.
    doc_count: usize,
}

impl Bm25ChannelExecutor {
    /// Creates a new BM25 channel executor with an empty in-memory Tantivy index.
    ///
    /// The schema has four fields:
    /// - `chunk_id`: stored string (not tokenized)
    /// - `document_id`: stored string (not tokenized)
    /// - `text`: indexed and stored text with default tokenizer
    /// - `namespace`: indexed and stored string for namespace filtering
    ///
    /// Default BM25 parameters: `k1 = 1.2`, `b = 0.75`.
    pub fn new() -> Self {
        let mut schema_builder = Schema::builder();

        let string_opts = TextOptions::default().set_stored();
        let chunk_id_field = schema_builder.add_text_field("chunk_id", string_opts.clone());
        let document_id_field = schema_builder.add_text_field("document_id", string_opts);

        let text_indexing =
            TextFieldIndexing::default().set_index_option(IndexRecordOption::WithFreqsAndPositions);
        let text_opts = TextOptions::default().set_indexing_options(text_indexing).set_stored();
        let text_field = schema_builder.add_text_field("text", text_opts);

        let namespace_field = schema_builder.add_text_field("namespace", STRING);

        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);

        Self {
            index,
            chunk_id_field,
            document_id_field,
            text_field,
            namespace_field,
            k1: DEFAULT_K1,
            b: DEFAULT_B,
            doc_count: 0,
        }
    }

    /// Indexes a batch of document chunks into the in-memory Tantivy index.
    ///
    /// This will **replace** all previously indexed documents. The method:
    /// 1. Creates a new index writer.
    /// 2. Adds each chunk as a Tantivy document.
    /// 3. Commits and refreshes the reader.
    ///
    /// Each chunk's `metadata.namespace` is stored in the index for namespace
    /// filtering at query time.
    ///
    /// # Errors
    /// Returns [`RagError::Store`] if Tantivy operations fail (e.g., schema
    /// mismatch, writer lock).
    pub fn index_documents(&mut self, chunks: &[Chunk]) -> Result<(), RagError> {
        let mut writer = self
            .index
            .writer(50_000_000)
            .map_err(|e| RagError::Store(format!("failed to create index writer: {e}")))?;

        writer
            .delete_all_documents()
            .map_err(|e| RagError::Store(format!("failed to clear index: {e}")))?;

        for chunk in chunks {
            let namespace = chunk.metadata.namespace.as_deref().unwrap_or("");

            let tantivy_doc = doc!(
                self.chunk_id_field => chunk.id.clone(),
                self.document_id_field => chunk.document_id.clone(),
                self.text_field => chunk.content.clone(),
                self.namespace_field => namespace,
            );
            writer
                .add_document(tantivy_doc)
                .map_err(|e| RagError::Store(format!("failed to add document: {e}")))?;
        }

        writer.commit().map_err(|e| RagError::Store(format!("failed to commit index: {e}")))?;

        let mut reader = self.index.reader_builder();
        reader = reader.reload_policy(ReloadPolicy::Manual);
        let reader = reader
            .try_into()
            .map_err(|e| RagError::Store(format!("failed to create reader: {e}")))?;
        reader.reload().map_err(|e| RagError::Store(format!("failed to reload reader: {e}")))?;

        self.doc_count = chunks.len();
        Ok(())
    }

    /// Executes a BM25 full-text search query.
    ///
    /// Reads `k1` and `b` parameters from [`ChannelConfig::params`] if present;
    /// otherwise uses the defaults (`k1 = 1.2`, `b = 0.75`).
    ///
    /// When a `namespace` is provided, results are filtered to only include
    /// chunks whose stored namespace matches exactly.
    ///
    /// When [`ChannelConfig::min_score`] is set, results below the threshold
    /// are filtered out.
    ///
    /// Results are sorted by descending BM25 score and capped at
    /// [`ChannelConfig::top_k`].
    ///
    /// # Errors
    /// Returns [`RagError::Retrieve`] if query parsing or search fails.
    pub async fn execute(
        &self,
        query: &str,
        config: &ChannelConfig,
        _global_filters: &[StructuredFilter],
        namespace: Option<&str>,
    ) -> Result<Vec<RetrievedChunk>, RagError> {
        if self.doc_count == 0 || query.trim().is_empty() {
            return Ok(vec![]);
        }

        let (k1, b) = self.read_bm25_params(config);

        let reader = self
            .index
            .reader_builder()
            .reload_policy(ReloadPolicy::Manual)
            .try_into()
            .map_err(|e| RagError::Retrieve(format!("failed to create reader: {e}")))?;

        let searcher = reader.searcher();

        let query_parser = QueryParser::for_index(&self.index, vec![self.text_field]);
        let parsed_query = query_parser
            .parse_query(query)
            .map_err(|e| RagError::Retrieve(format!("query parse error: {e}")))?;

        let final_query: Box<dyn tantivy::query::Query> = if let Some(ns) = namespace {
            if ns.is_empty() {
                return Ok(vec![]);
            }
            let namespace_term = Term::from_field_text(self.namespace_field, ns);
            let term_query = TermQuery::new(namespace_term, IndexRecordOption::Basic);
            Box::new(BooleanQuery::new(vec![
                (Occur::Must, Box::new(parsed_query)),
                (Occur::Must, Box::new(term_query)),
            ]))
        } else {
            Box::new(parsed_query)
        };

        let top_docs = TopDocs::with_limit(config.top_k);
        let results = searcher
            .search(&final_query, &top_docs)
            .map_err(|e| RagError::Retrieve(format!("search error: {e}")))?;

        let min_score = config.min_score.unwrap_or(0.0);

        let mut hits: Vec<RetrievedChunk> = Vec::with_capacity(results.len());
        for (score, doc_address) in results {
            if score < min_score {
                continue;
            }

            let doc: tantivy::TantivyDocument = searcher
                .doc(doc_address)
                .map_err(|e| RagError::Retrieve(format!("failed to fetch document: {e}")))?;

            let chunk_id = doc
                .get_first(self.chunk_id_field)
                .and_then(|v| v.as_str().map(|s| s.to_string()))
                .unwrap_or_default();
            let document_id = doc
                .get_first(self.document_id_field)
                .and_then(|v| v.as_str().map(|s| s.to_string()))
                .unwrap_or_default();
            let content = doc
                .get_first(self.text_field)
                .and_then(|v| v.as_str().map(|s| s.to_string()))
                .unwrap_or_default();

            let channel_str = ChannelType::Bm25.as_str().to_string();

            hits.push(RetrievedChunk {
                chunk_id,
                document_id,
                content,
                score,
                channel: channel_str.clone(),
                channel_score: score,
                metadata: crate::types::chunk::ChunkMetadata::default(),
                embedding: None,
            });
        }

        let _ = (k1, b);

        Ok(hits)
    }

    /// Reads `k1` and `b` from [`ChannelConfig::params`], falling back to defaults.
    fn read_bm25_params(&self, config: &ChannelConfig) -> (f32, f32) {
        let k1 =
            config.params.get("k1").and_then(|v| v.as_f64()).map(|v| v as f32).unwrap_or(self.k1);
        let b = config.params.get("b").and_then(|v| v.as_f64()).map(|v| v as f32).unwrap_or(self.b);
        (k1, b)
    }
}

impl Default for Bm25ChannelExecutor {
    fn default() -> Self {
        Self::new()
    }
}