ryo-analysis 0.1.0

//! Tantivy-based literal search index.
//!
//! This module provides full-text search capabilities for literals
//! using Tantivy as the search engine.

use std::sync::Arc;

use tantivy::collector::TopDocs;
use tantivy::query::{BooleanQuery, Occur, Query, RegexQuery, TermQuery};
use tantivy::schema::{
    Field, IndexRecordOption, Schema, TextFieldIndexing, TextOptions, Value, STORED, STRING,
};
use tantivy::{Index, ReloadPolicy, TantivyDocument, Term};

use slotmap::Key;

use crate::context::ImHashMap;
use crate::symbol::{FileId, FileRegistry, SymbolId, SymbolRegistry};
use ryo_source::pure::PureFile;

use super::{LiteralCollector, LiteralInfo, LiteralKind};

// =============================================================================
// Error Type
// =============================================================================

/// Errors that can occur during literal search operations.
#[derive(Debug)]
pub enum LiteralSearchError {
    /// Tantivy index error.
    Index(tantivy::TantivyError),
    /// Query parse error.
    QueryParse(tantivy::query::QueryParserError),
    /// Invalid query pattern.
    InvalidPattern(String),
}

impl std::fmt::Display for LiteralSearchError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Index(e) => write!(f, "Index error: {}", e),
            Self::QueryParse(e) => write!(f, "Query parse error: {}", e),
            Self::InvalidPattern(s) => write!(f, "Invalid pattern: {}", s),
        }
    }
}

impl std::error::Error for LiteralSearchError {}

impl From<tantivy::TantivyError> for LiteralSearchError {
    fn from(e: tantivy::TantivyError) -> Self {
        Self::Index(e)
    }
}

impl From<tantivy::query::QueryParserError> for LiteralSearchError {
    fn from(e: tantivy::query::QueryParserError) -> Self {
        Self::QueryParse(e)
    }
}

// =============================================================================
// LiteralQuery
// =============================================================================

/// Query parameters for literal search.
#[derive(Debug, Clone, Default)]
pub struct LiteralQuery {
    /// Pattern to search for (glob-style: `*error*`, `"hello"`, etc.).
    pub pattern: String,
    /// Filter by literal kind.
    pub kind: Option<LiteralKind>,
    /// Maximum number of results.
    pub limit: usize,
}

impl LiteralQuery {
    /// Create a new query with the given pattern.
    pub fn new(pattern: impl Into<String>) -> Self {
        Self {
            pattern: pattern.into(),
            kind: None,
            limit: 100,
        }
    }

    /// Set the literal kind filter.
    pub fn with_kind(mut self, kind: LiteralKind) -> Self {
        self.kind = Some(kind);
        self
    }

    /// Set the result limit.
    pub fn with_limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }
}

// =============================================================================
// LiteralMatch
// =============================================================================

/// A search result from the literal index.
#[derive(Debug, Clone)]
pub struct LiteralMatch {
    /// The literal value.
    pub value: String,
    /// The literal kind.
    pub kind: LiteralKind,
    /// The symbol containing this literal.
    pub symbol_id: SymbolId,
    /// The file containing this literal.
    pub file_id: FileId,
    /// The file path as a string (for display without FileRegistry lookup).
    pub file_path: String,
    /// Search relevance score.
    pub score: f32,
}

// =============================================================================
// LiteralSchema
// =============================================================================

/// Schema fields for the literal index.
struct LiteralSchema {
    /// The full schema.
    #[allow(dead_code)]
    schema: Schema,
    /// Literal value (searchable).
    value: Field,
    /// Literal kind (filterable).
    kind: Field,
    /// Symbol ID (stored).
    symbol_id: Field,
    /// File ID (stored).
    file_id: Field,
    /// File path (stored, for display).
    file_path: Field,
}

impl LiteralSchema {
    fn new() -> Self {
        let mut schema_builder = Schema::builder();

        // Value field: full-text searchable + stored
        let text_options = TextOptions::default()
            .set_indexing_options(
                TextFieldIndexing::default()
                    .set_tokenizer("raw")
                    .set_index_option(IndexRecordOption::WithFreqsAndPositions),
            )
            .set_stored();
        let value = schema_builder.add_text_field("value", text_options);

        // Kind field: keyword (exact match) + stored
        let kind = schema_builder.add_text_field("kind", STRING | STORED);

        // Symbol ID: stored only (we use u64 as string for simplicity)
        let symbol_id = schema_builder.add_text_field("symbol_id", STRING | STORED);

        // File ID: stored only
        let file_id = schema_builder.add_text_field("file_id", STRING | STORED);

        // File path: stored only (for display without registry lookup)
        let file_path = schema_builder.add_text_field("file_path", STRING | STORED);

        let schema = schema_builder.build();

        Self {
            schema,
            value,
            kind,
            symbol_id,
            file_id,
            file_path,
        }
    }
}

// =============================================================================
// LiteralIndex
// =============================================================================

/// Tantivy-based index for literal search.
///
/// Provides fast full-text search over all literals in the codebase.
pub struct LiteralIndex {
    index: Index,
    schema: LiteralSchema,
}

impl LiteralIndex {
    /// Create a new empty index in RAM.
    pub fn new() -> Result<Self, LiteralSearchError> {
        let schema = LiteralSchema::new();
        let index = Index::create_in_ram(schema.schema.clone());

        // Register the "raw" tokenizer (no tokenization, exact matching)
        index
            .tokenizers()
            .register("raw", tantivy::tokenizer::RawTokenizer::default());

        Ok(Self { index, schema })
    }

    /// Build an index from parsed files.
    ///
    /// This collects all literals from the files and indexes them for search.
    pub fn build_from_files(
        files: &ImHashMap<FileId, Arc<PureFile>>,
        registry: &SymbolRegistry,
        file_registry: &FileRegistry,
    ) -> Result<Self, LiteralSearchError> {
        let index = Self::new()?;
        let mut writer = index.index.writer(50_000_000)?; // 50MB heap

        let default_symbol = SymbolId::from(slotmap::KeyData::from_ffi(0));

        for (&file_id, file) in files.iter() {
            let file_path_str = file_registry
                .path(file_id)
                .map(|p| p.as_relative().display().to_string())
                .unwrap_or_default();
            LiteralCollector::collect_file(
                file.as_ref(),
                file_id,
                default_symbol,
                |name| registry.lookup_by_name(name),
                |info| {
                    let doc = index.create_document(&info, &file_path_str);
                    let _ = writer.add_document(doc);
                },
            );
        }

        writer.commit()?;
        Ok(index)
    }

    /// Build an index from workspace files (WorkspaceFilePath keys).
    ///
    /// This is a convenience wrapper for build_from_files that works with
    /// the AnalysisContext's file map.
    pub fn build_from_workspace_files(
        files: &ImHashMap<ryo_symbol::WorkspaceFilePath, Arc<PureFile>>,
        registry: &SymbolRegistry,
    ) -> Result<Self, LiteralSearchError> {
        let index = Self::new()?;
        let mut writer = index.index.writer(50_000_000)?; // 50MB heap

        let default_symbol = SymbolId::from(slotmap::KeyData::from_ffi(0));
        let default_file_id = FileId::from(slotmap::KeyData::from_ffi(0));

        for (path, file) in files.iter() {
            let file_path_str = path.as_relative().display().to_string();
            LiteralCollector::collect_file(
                file.as_ref(),
                default_file_id,
                default_symbol,
                |name| registry.lookup_by_name(name),
                |info| {
                    let doc = index.create_document(&info, &file_path_str);
                    let _ = writer.add_document(doc);
                },
            );
        }

        writer.commit()?;
        Ok(index)
    }

    /// Create a document from literal info.
    fn create_document(&self, info: &LiteralInfo, file_path: &str) -> TantivyDocument {
        let mut doc = TantivyDocument::new();
        doc.add_text(self.schema.value, &info.value);
        doc.add_text(self.schema.kind, info.kind.as_str());
        doc.add_text(
            self.schema.symbol_id,
            info.symbol_id.data().as_ffi().to_string(),
        );
        doc.add_text(
            self.schema.file_id,
            info.file_id.data().as_ffi().to_string(),
        );
        doc.add_text(self.schema.file_path, file_path);
        doc
    }

    /// Search for literals matching the query.
    pub fn search(&self, query: &LiteralQuery) -> Result<Vec<LiteralMatch>, LiteralSearchError> {
        let reader = self
            .index
            .reader_builder()
            .reload_policy(ReloadPolicy::Manual)
            .try_into()?;
        let searcher = reader.searcher();

        // Build the query
        let tantivy_query = self.build_query(query)?;

        // Execute search
        let top_docs = searcher.search(
            &tantivy_query,
            &TopDocs::with_limit(query.limit).order_by_score(),
        )?;

        // Convert results
        let mut results = Vec::new();
        for (score, doc_address) in top_docs {
            let doc: tantivy::TantivyDocument = searcher.doc(doc_address)?;
            if let Some(m) = self.doc_to_match(&doc, score) {
                results.push(m);
            }
        }

        Ok(results)
    }

    /// Build a Tantivy query from our query parameters.
    fn build_query(&self, query: &LiteralQuery) -> Result<Box<dyn Query>, LiteralSearchError> {
        let mut subqueries: Vec<(Occur, Box<dyn Query>)> = Vec::new();

        // Pattern query on value field
        if !query.pattern.is_empty() {
            let pattern_query = self.build_pattern_query(&query.pattern)?;
            subqueries.push((Occur::Must, pattern_query));
        }

        // Kind filter
        if let Some(kind) = &query.kind {
            let kind_term = Term::from_field_text(self.schema.kind, kind.as_str());
            let kind_query = TermQuery::new(kind_term, IndexRecordOption::Basic);
            subqueries.push((Occur::Must, Box::new(kind_query)));
        }

        if subqueries.is_empty() {
            // Match all if no constraints
            Ok(Box::new(tantivy::query::AllQuery))
        } else if subqueries.len() == 1 {
            Ok(subqueries.pop().unwrap().1)
        } else {
            Ok(Box::new(BooleanQuery::new(subqueries)))
        }
    }

    /// Build a query for the pattern (supports glob-style wildcards).
    fn build_pattern_query(&self, pattern: &str) -> Result<Box<dyn Query>, LiteralSearchError> {
        // If pattern is just "*", match all
        if pattern == "*" {
            return Ok(Box::new(tantivy::query::AllQuery));
        }

        // Convert glob pattern to regex
        // Tantivy doesn't support .* at start/end, so we use .+ instead
        let regex_pattern = glob_to_regex_tantivy(pattern);

        // Use RegexQuery for pattern matching
        let regex_query = RegexQuery::from_pattern(&regex_pattern, self.schema.value)
            .map_err(|e| LiteralSearchError::InvalidPattern(format!("{}: {}", pattern, e)))?;

        Ok(Box::new(regex_query))
    }

    /// Convert a document to a LiteralMatch.
    fn doc_to_match(&self, doc: &tantivy::TantivyDocument, score: f32) -> Option<LiteralMatch> {
        let value = doc.get_first(self.schema.value)?.as_str()?.to_string();
        let kind_str = doc.get_first(self.schema.kind)?.as_str()?;
        let kind = LiteralKind::parse_kind(kind_str)?;

        // Parse symbol_id and file_id from debug format
        // This is a simplification; in production we'd store them differently
        let symbol_id_str = doc.get_first(self.schema.symbol_id)?.as_str()?;
        let file_id_str = doc.get_first(self.schema.file_id)?.as_str()?;
        let file_path = doc
            .get_first(self.schema.file_path)
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();

        let symbol_id = parse_symbol_id(symbol_id_str)?;
        let file_id = parse_file_id(file_id_str)?;

        Some(LiteralMatch {
            value,
            kind,
            symbol_id,
            file_id,
            file_path,
            score,
        })
    }

    /// Get statistics about the index.
    pub fn stats(&self) -> LiteralIndexStats {
        let reader = self
            .index
            .reader_builder()
            .reload_policy(ReloadPolicy::Manual)
            .try_into()
            .ok();

        let doc_count = reader
            .as_ref()
            .map(|r| r.searcher().num_docs() as usize)
            .unwrap_or(0);

        LiteralIndexStats { doc_count }
    }
}

/// Statistics about the literal index.
#[derive(Debug, Clone)]
pub struct LiteralIndexStats {
    /// Number of indexed literals.
    pub doc_count: usize,
}

// =============================================================================
// Helper Functions
// =============================================================================

/// Convert a glob pattern to a regex pattern for Tantivy.
///
/// Tantivy's regex engine doesn't support empty matches (like `.*` at boundaries),
/// so we use a different approach:
/// - `*foo` → `.*foo` (contains "foo" at end)
/// - `foo*` → `foo.*` (starts with "foo")
/// - `*foo*` → `.*foo.*` (contains "foo")
///
/// For patterns that would produce empty matches, we use `.+` instead.
fn glob_to_regex_tantivy(glob: &str) -> String {
    let mut regex = String::with_capacity(glob.len() * 2);

    let chars: Vec<char> = glob.chars().collect();
    let len = chars.len();

    for (i, &c) in chars.iter().enumerate() {
        match c {
            '*' => {
                // Use .+ to avoid empty matches at boundaries
                // But only if there's content before/after
                let at_start = i == 0;
                let at_end = i == len - 1;

                if at_start && at_end {
                    // Just "*" - handled separately
                    regex.push_str(".*");
                } else if at_start {
                    // "*foo" - match anything before foo
                    regex.push_str(".*");
                } else if at_end {
                    // "foo*" - match anything after foo
                    regex.push_str(".*");
                } else {
                    // "fo*o" - match anything in middle
                    regex.push_str(".*");
                }
            }
            '?' => regex.push('.'),
            '.' | '+' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '^' | '$' | '\\' => {
                regex.push('\\');
                regex.push(c);
            }
            _ => regex.push(c),
        }
    }

    regex
}

/// Parse a SymbolId from its FFI u64 string representation.
fn parse_symbol_id(s: &str) -> Option<SymbolId> {
    let ffi: u64 = s.parse().ok()?;
    Some(SymbolId::from(slotmap::KeyData::from_ffi(ffi)))
}

/// Parse a FileId from its FFI u64 string representation.
fn parse_file_id(s: &str) -> Option<FileId> {
    let ffi: u64 = s.parse().ok()?;
    Some(FileId::from(slotmap::KeyData::from_ffi(ffi)))
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_glob_to_regex_tantivy() {
        // Tantivy doesn't use anchors, patterns match anywhere in the value
        assert_eq!(glob_to_regex_tantivy("*error*"), ".*error.*");
        assert_eq!(glob_to_regex_tantivy("hello"), "hello");
        assert_eq!(glob_to_regex_tantivy("test?"), "test.");
        assert_eq!(glob_to_regex_tantivy("a.b"), "a\\.b");
    }

    #[test]
    fn test_create_empty_index() {
        let index = LiteralIndex::new().expect("Failed to create index");
        let stats = index.stats();
        assert_eq!(stats.doc_count, 0);
    }

    #[test]
    fn test_search_empty_index() {
        let index = LiteralIndex::new().expect("Failed to create index");
        let query = LiteralQuery::new("*");
        let results = index.search(&query).expect("Search failed");
        assert!(results.is_empty());
    }

    #[test]
    fn test_index_and_search() {
        let index = LiteralIndex::new().expect("Failed to create index");
        let mut writer = index
            .index
            .writer(50_000_000)
            .expect("Failed to create writer");

        // Add some test documents
        let literals = vec![
            ("\"hello world\"", LiteralKind::String),
            ("\"error: connection failed\"", LiteralKind::String),
            ("42", LiteralKind::Int),
            ("3.14", LiteralKind::Float),
            ("true", LiteralKind::Bool),
        ];

        let dummy_symbol = SymbolId::from(slotmap::KeyData::from_ffi(1));
        let dummy_file = FileId::from(slotmap::KeyData::from_ffi(1));

        for (value, kind) in literals {
            let info = LiteralInfo::with_kind(value.to_string(), kind, dummy_symbol, dummy_file);
            let doc = index.create_document(&info, "test/file.rs");
            writer.add_document(doc).expect("Failed to add document");
        }
        writer.commit().expect("Failed to commit");

        // Test search
        let stats = index.stats();
        assert_eq!(stats.doc_count, 5);

        // Search for error
        let query = LiteralQuery::new("*error*");
        let results = index.search(&query).expect("Search failed");
        assert_eq!(results.len(), 1);
        assert!(results[0].value.contains("error"));

        // Search with kind filter
        let query = LiteralQuery::new("*").with_kind(LiteralKind::Int);
        let results = index.search(&query).expect("Search failed");
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].value, "42");
    }

    #[test]
    fn test_literal_query_builder() {
        let query = LiteralQuery::new("*test*")
            .with_kind(LiteralKind::String)
            .with_limit(50);

        assert_eq!(query.pattern, "*test*");
        assert_eq!(query.kind, Some(LiteralKind::String));
        assert_eq!(query.limit, 50);
    }
}