minni 0.1.1

mod ann;
mod bm25;
mod dense;
mod meta;
mod reranker;
mod tokenizer;

pub use ann::AnnIndex;
pub use bm25::{BM25Index, BM25SearchResult};
pub use dense::DenseRetriever;
pub use meta::{read_metadata, write_metadata, IndexMetadata};
pub use reranker::Reranker;

use anyhow::Result;
use serde::Serialize;
use std::collections::HashMap;
use std::path::Path;
use std::time::Instant;

/// Search run metadata.
#[derive(Debug, Clone, Serialize)]
pub struct SearchMeta {
    /// Search mode.
    pub mode: String,
    pub reranker_active: bool,
    pub dense_active: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fallback_reason: Option<String>,
    /// Candidate count before reranking.
    pub total_candidates: usize,
    pub result_count: usize,
    pub search_time_ms: u64,
}

/// BM25 retrieval with optional reranking.
pub struct HybridSearch {
    bm25: BM25Index,
    reranker: Option<Reranker>,
    dense: Option<DenseRetriever>,
    ann_index: Option<AnnIndex>,
    embeddings: Vec<(String, Vec<f32>)>,
    chunk_lookup: HashMap<String, BM25SearchResult>,
}

/// Optional search filters.
#[derive(Debug, Clone, Default)]
pub struct SearchFilters {
    pub path: Option<String>,
    pub language: Option<String>,
    pub chunk_type: Option<String>,
    pub symbol: Option<String>,
}

impl SearchFilters {
    pub fn is_active(&self) -> bool {
        self.path.is_some()
            || self.language.is_some()
            || self.chunk_type.is_some()
            || self.symbol.is_some()
    }

    fn matches(&self, result: &BM25SearchResult) -> bool {
        if let Some(path) = &self.path {
            if !result.file_path.contains(path) {
                return false;
            }
        }
        if let Some(lang) = &self.language {
            let lang = lang.to_lowercase();
            if result.language.to_lowercase() != lang {
                return false;
            }
        }
        if let Some(chunk_type) = &self.chunk_type {
            let chunk_type = chunk_type.to_lowercase();
            if result.chunk_type.to_lowercase() != chunk_type {
                return false;
            }
        }
        if let Some(symbol) = &self.symbol {
            let symbol = symbol.to_lowercase();
            let candidate = result.symbol_name.as_deref().unwrap_or("").to_lowercase();
            if !candidate.contains(&symbol) {
                return false;
            }
        }
        true
    }
}

impl HybridSearch {
    pub fn new(
        index_path: &Path,
        reranker_model_path: Option<&Path>,
        dense_model_path: Option<&Path>,
        ann_index: Option<AnnIndex>,
        embeddings: Vec<(String, Vec<f32>)>,
        chunk_lookup: HashMap<String, BM25SearchResult>,
    ) -> Result<Self> {
        let bm25 = BM25Index::new(index_path)?;
        let reranker = if let Some(path) = reranker_model_path {
            Some(Reranker::new(path)?)
        } else {
            None
        };
        let dense = if let Some(path) = dense_model_path {
            Some(DenseRetriever::new(path)?)
        } else {
            None
        };

        Ok(Self {
            bm25,
            reranker,
            dense,
            ann_index,
            embeddings,
            chunk_lookup,
        })
    }

    /// Search with automatic re-ranking when beneficial
    pub fn search(&mut self, query: &str, limit: usize) -> Result<(Vec<SearchResult>, SearchMeta)> {
        self.search_with_filters(query, limit, None)
    }

    /// Search with optional filters applied before re-ranking
    pub fn search_with_filters(
        &mut self,
        query: &str,
        limit: usize,
        filters: Option<&SearchFilters>,
    ) -> Result<(Vec<SearchResult>, SearchMeta)> {
        let start = Instant::now();

        // Stage 1: BM25 retrieval (get more candidates for re-ranking)
        let mut candidate_limit = if self.reranker.is_some() {
            (limit * 10).max(100) // Get 10x results for re-ranking
        } else {
            limit
        };

        if filters.map(|f| f.is_active()).unwrap_or(false) {
            candidate_limit = (candidate_limit * 5).min(5000);
        }

        let mut bm25_results = self.bm25.search(query, candidate_limit)?;
        self.enrich_from_lookup(&mut bm25_results);

        if let Some(filters) = filters {
            if filters.is_active() {
                bm25_results.retain(|r| filters.matches(r));
            }
        }

        let dense_active = self.dense.is_some();
        if let Some(dense) = &mut self.dense {
            let dense_limit = candidate_limit.max(limit * 5);
            let dense_results = if let Some(ann_index) = &self.ann_index {
                match dense.embed_query(query) {
                    Ok(query_embedding) => ann_index.search(&query_embedding, dense_limit),
                    Err(_) => Vec::new(),
                }
            } else {
                dense
                    .search_embeddings(query, &self.embeddings, dense_limit)
                    .unwrap_or_default()
            };

            for dense_hit in dense_results {
                if bm25_results
                    .iter()
                    .any(|r| r.chunk_id == dense_hit.chunk_id)
                {
                    continue;
                }
                if let Some(base) = self.chunk_lookup.get(&dense_hit.chunk_id) {
                    let mut candidate = base.clone();
                    candidate.score = dense_hit.score;
                    if filters.map(|f| f.matches(&candidate)).unwrap_or(true) {
                        bm25_results.push(candidate);
                    }
                }
            }
        }

        bm25_results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let total_candidates = bm25_results.len();

        // Stage 2: Neural re-ranking (if available and beneficial)
        let reranker_threshold = limit * 2;
        let (results, reranker_active, fallback_reason) = if let Some(reranker) = &mut self.reranker
        {
            if bm25_results.len() > reranker_threshold {
                // Re-rank if we have enough candidates
                let ranked = reranker.rerank(query, bm25_results, limit)?;
                (ranked, true, None)
            } else {
                // Not enough results to warrant re-ranking, convert to SearchResult
                let n = bm25_results.len();
                let reason = format!(
                    "not enough candidates for reranking: {} < {}",
                    n,
                    reranker_threshold + 1
                );
                let converted: Vec<SearchResult> = bm25_results
                    .into_iter()
                    .take(limit)
                    .map(|r| r.into())
                    .collect();
                (converted, false, Some(reason))
            }
        } else {
            // No re-ranker, use BM25 scores, convert to SearchResult
            let converted: Vec<SearchResult> = bm25_results
                .into_iter()
                .take(limit)
                .map(|r| r.into())
                .collect();
            (converted, false, None)
        };

        let search_time_ms = start.elapsed().as_millis() as u64;
        let result_count = results.len();

        let mode = match (dense_active, reranker_active) {
            (false, false) => "bm25".to_string(),
            (false, true) => "bm25+rerank".to_string(),
            (true, false) => "hybrid".to_string(),
            (true, true) => "hybrid+rerank".to_string(),
        };

        let meta = SearchMeta {
            mode,
            reranker_active,
            dense_active,
            fallback_reason,
            total_candidates,
            result_count,
            search_time_ms,
        };

        Ok((results, meta))
    }

    /// Search with BM25 only (fast)
    #[allow(dead_code)]
    pub fn search_bm25(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>> {
        let mut results = self.bm25.search(query, limit)?;
        self.enrich_from_lookup(&mut results);
        Ok(results.into_iter().map(|r| r.into()).collect())
    }

    /// Populate the four new metadata fields on BM25 results from the chunk lookup.
    ///
    /// The Tantivy index does not store these fields, so we look them up from
    /// SQLite-backed `chunk_lookup` by `chunk_id` and copy them across.
    fn enrich_from_lookup(&self, results: &mut Vec<BM25SearchResult>) {
        for result in results.iter_mut() {
            if let Some(entry) = self.chunk_lookup.get(&result.chunk_id) {
                result.parent_symbol = entry.parent_symbol.clone();
                result.signature = entry.signature.clone();
                result.doc_comment = entry.doc_comment.clone();
                result.module_path = entry.module_path.clone();
            }
        }
    }
}

/// Final search result returned to the CLI.
#[derive(Debug, Clone)]
pub struct SearchResult {
    pub chunk_id: String,
    pub file_path: String,
    pub content: String,
    pub start_line: usize,
    pub end_line: usize,
    pub chunk_type: String,
    pub language: String,
    pub symbol_name: Option<String>,
    pub score: f32,
    pub parent_symbol: Option<String>,
    pub signature: Option<String>,
    pub doc_comment: Option<String>,
    pub module_path: Option<String>,
}

impl From<BM25SearchResult> for SearchResult {
    fn from(r: BM25SearchResult) -> Self {
        Self {
            chunk_id: r.chunk_id,
            file_path: r.file_path,
            content: r.content,
            start_line: r.start_line,
            end_line: r.end_line,
            chunk_type: r.chunk_type,
            language: r.language,
            symbol_name: r.symbol_name,
            score: r.score,
            parent_symbol: r.parent_symbol,
            signature: r.signature,
            doc_comment: r.doc_comment,
            module_path: r.module_path,
        }
    }
}