zesdex/apps/infrastructure/src/tools/semantic_search.rs

//! Semantic code search tool — indexes all code symbols (functions, structs,
//! enums, traits, modules) in a project and allows searching by name,
//! concept, or meaning.
//!
//! Flow: walk workspace files → parse Rust source for symbol declarations →
//! build in-memory index → search by fuzzy/prefix match on symbol names and
//! doc comments.

use crate::tools::{Tool, ToolCtx};
use anyhow::Result;
use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::Mutex;
use tracing::{debug, info, instrument, warn};

// ---------------------------------------------------------------------------
//  Symbol index types
// ---------------------------------------------------------------------------

/// The kind of a code symbol.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SymbolKind {
    Function,
    Struct,
    Enum,
    Trait,
    Module,
    Impl,
    Type,
    Constant,
    Macro,
    Other,
}

impl std::fmt::Display for SymbolKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SymbolKind::Function => write!(f, "fn"),
            SymbolKind::Struct => write!(f, "struct"),
            SymbolKind::Enum => write!(f, "enum"),
            SymbolKind::Trait => write!(f, "trait"),
            SymbolKind::Module => write!(f, "mod"),
            SymbolKind::Impl => write!(f, "impl"),
            SymbolKind::Type => write!(f, "type"),
            SymbolKind::Constant => write!(f, "const"),
            SymbolKind::Macro => write!(f, "macro"),
            SymbolKind::Other => write!(f, "symbol"),
        }
    }
}

/// A single code symbol entry in the index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeSymbol {
    /// Symbol name (e.g. "run_agent", "AppStateRest").
    pub name: String,
    /// Kind of symbol.
    pub kind: SymbolKind,
    /// File path relative to workspace root.
    pub file: String,
    /// Line number (1-indexed).
    pub line: usize,
    /// Parent symbol (e.g. struct name for impl methods).
    pub parent: Option<String>,
    /// Doc comment text, if any.
    pub doc_comment: Option<String>,
    /// Short context (the declaration line).
    pub context: String,
}

/// The in-memory symbol index, shared via a global static.
static SYMBOL_INDEX: Mutex<Option<SymbolIndex>> = Mutex::new(None);

/// Lazily compile regexes for symbol extraction.
fn compiled_regexes() -> (
    Regex,
    Regex,
    Regex,
    Regex,
    Regex,
    Regex,
    Regex,
    Regex,
    Regex,
) {
    // All regex patterns are static — compilation is infallible.
    (
        Regex::new(r"(?m)^\s*(?:pub\s+)?(?:(?:unsafe\s+)?async\s+)?fn\s+(\w+)").expect("fn regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?struct\s+(\w+)").expect("struct regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?enum\s+(\w+)").expect("enum regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?(?:(?:unsafe\s+)?)?trait\s+(\w+)").expect("trait regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?mod\s+(\w+)").expect("mod regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?(?:unsafe\s+)?impl(?:\s*<[^>]*>)?\s+(?:for\s+)?(\w+)").expect("impl regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?type\s+(\w+)").expect("type regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?const\s+(\w+)").expect("const regex"),
        Regex::new(r"(?m)^\s*(?:pub\s+)?macro_rules!\s*\(\s*(\w+)").expect("macro regex"),
    )
}

/// A symbol index cache.
#[derive(Debug, Clone)]
pub struct SymbolIndex {
    symbols: Vec<CodeSymbol>,
    workspace_path: Option<String>,
}

impl SymbolIndex {
    pub fn new() -> Self {
        SymbolIndex {
            symbols: Vec::new(),
            workspace_path: None,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.symbols.is_empty()
    }

    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    pub fn rebuild(&mut self, workspace: &str) -> Result<usize> {
        let path = std::path::Path::new(workspace);
        if !path.exists() {
            anyhow::bail!("workspace path does not exist: {workspace}");
        }

        let mut symbols = Vec::new();
        let walker = ignore::Walk::new(path);

        for entry in walker.flatten() {
            let file_path = entry.path();
            if !file_path.is_file() {
                continue;
            }

            // Only index Rust files
            let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or("");
            if ext != "rs" {
                continue;
            }

            let rel_path = file_path
                .strip_prefix(path)
                .unwrap_or(file_path)
                .display()
                .to_string();

            match std::fs::read_to_string(file_path) {
                Ok(content) => {
                    let file_symbols = extract_symbols(&content, &rel_path);
                    symbols.extend(file_symbols);
                }
                Err(e) => {
                    debug!(file = %rel_path, error = %e, "failed to read file for indexing");
                }
            }
        }

        symbols.sort_by(|a, b| a.name.cmp(&b.name));
        self.symbols = symbols;
        self.workspace_path = Some(workspace.to_string());

        let count = self.symbols.len();
        info!(symbol_count = count, workspace = %workspace, "symbol index rebuilt");
        Ok(count)
    }

    pub fn search(&self, query: &str, max_results: usize) -> Vec<&CodeSymbol> {
        if self.symbols.is_empty() {
            return Vec::new();
        }

        let query_lower = query.to_lowercase();
        let query_words: Vec<String> = query_lower.split_whitespace().map(|s| s.to_string()).collect();

        let mut scored: Vec<(i32, &CodeSymbol)> = self
            .symbols
            .iter()
            .filter_map(|sym| {
                let score = score_symbol(sym, &query_lower, &query_words);
                if score > 0 {
                    Some((score, sym))
                } else {
                    None
                }
            })
            .collect();

        // Sort by score descending, then by name ascending
        scored.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.name.cmp(&b.1.name)));

        scored
            .into_iter()
            .take(max_results)
            .map(|(_, sym)| sym)
            .collect()
    }
}

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

/// Score a symbol against a search query.
fn score_symbol(sym: &CodeSymbol, query_lower: &str, query_words: &[String]) -> i32 {
    let name_lower = sym.name.to_lowercase();
    let mut score: i32 = 0;

    // Exact match = highest score
    if name_lower == *query_lower {
        score += 1000;
    }

    // Prefix match
    if name_lower.starts_with(query_lower) {
        score += 500;
    }

    // Contains
    if name_lower.contains(query_lower) {
        score += 200;
    }

    // Word-by-word matching
    for word in query_words {
        if name_lower.contains(word) {
            score += 50;
        }
    }

    // Doc comment match
    if let Some(ref doc) = sym.doc_comment {
        let doc_lower = doc.to_lowercase();
        if doc_lower.contains(query_lower) {
            score += 30;
        }
        for word in query_words {
            if doc_lower.contains(word) {
                score += 10;
            }
        }
    }

    // Context match
    let context_lower = sym.context.to_lowercase();
    if context_lower.contains(query_lower) {
        score += 20;
    }

    score
}

/// Extract code symbols from Rust source content.
fn extract_symbols(content: &str, rel_path: &str) -> Vec<CodeSymbol> {
    let (fn_re, struct_re, enum_re, trait_re, mod_re, impl_re, type_re, const_re, macro_re) =
        compiled_regexes();

    let mut symbols = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    // Extract doc comments that precede declarations
    let doc_comments = extract_doc_comments(&lines);

    for (i, line) in lines.iter().enumerate() {
        let line_num = i + 1;
        let trimmed = line.trim();

        // Check for function declarations
        if let Some(caps) = fn_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Function,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for struct declarations
        if let Some(caps) = struct_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Struct,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for enum declarations
        if let Some(caps) = enum_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Enum,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for trait declarations
        if let Some(caps) = trait_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Trait,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for module declarations
        if let Some(caps) = mod_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Module,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: None,
                context: trimmed.to_string(),
            });
        }

        // Check for type alias declarations
        if let Some(caps) = type_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Type,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for const declarations
        if let Some(caps) = const_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            let doc = doc_comments.get(&line_num).cloned();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Constant,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: doc,
                context: trimmed.to_string(),
            });
        }

        // Check for macro declarations
        if let Some(caps) = macro_re.captures(trimmed) {
            let name = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            symbols.push(CodeSymbol {
                name,
                kind: SymbolKind::Macro,
                file: rel_path.to_string(),
                line: line_num,
                parent: None,
                doc_comment: None,
                context: trimmed.to_string(),
            });
        }

        // Parse impl blocks for method-level indexing
        if let Some(caps) = impl_re.captures(trimmed) {
            let impl_for = caps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
            // Look for methods inside this impl block
            let mut brace_depth: i32 = 0;
            let mut started = false;
            for (j, l) in lines[i..].iter().enumerate() {
                for ch in l.chars() {
                    match ch {
                        '{' => {
                            brace_depth += 1;
                            started = true;
                        }
                        '}' => {
                            brace_depth -= 1;
                        }
                        _ => {}
                    }
                }
                if started && brace_depth <= 0 && j > 1 {
                    break; // End of impl block
                }
                if j > 0 {
                    let inner_line = l.trim();
                    if let Some(mcaps) = fn_re.captures(inner_line) {
                        let method_name = mcaps.get(1).expect("capture group 1 exists by regex").as_str().to_string();
                        let abs_line = i + j + 1;
                        let doc = doc_comments.get(&abs_line).cloned();
                        symbols.push(CodeSymbol {
                            name: format!("{impl_for}::{method_name}"),
                            kind: SymbolKind::Function,
                            file: rel_path.to_string(),
                            line: abs_line,
                            parent: Some(impl_for.clone()),
                            doc_comment: doc,
                            context: inner_line.to_string(),
                        });
                    }
                }
            }
        }
    }

    symbols
}

/// Extract doc comments (/// or //!) that precede each line.
fn extract_doc_comments(lines: &[&str]) -> HashMap<usize, String> {
    let mut map = HashMap::new();
    let mut i = 0;
    while i < lines.len() {
        let line = lines[i].trim();
        if line.starts_with("///") {
            let mut doc = String::new();
            while i < lines.len() {
                let l = lines[i].trim();
                if l.starts_with("///") {
                    if !doc.is_empty() {
                        doc.push(' ');
                    }
                    doc.push_str(l.trim_start_matches("///").trim());
                    i += 1;
                } else {
                    break;
                }
            }
            // Associate doc with the next non-empty, non-doc, non-attribute line
            let target_line = find_next_declaration_line(lines, i);
            if let Some(tl) = target_line {
                map.insert(tl + 1, doc);
            }
        } else {
            i += 1;
        }
    }
    map
}

/// Find the next line that looks like a declaration (not doc, not attr).
fn find_next_declaration_line(lines: &[&str], start: usize) -> Option<usize> {
    lines[start..].iter().position(|line| {
        let trimmed = line.trim();
        !trimmed.is_empty()
            && !trimmed.starts_with("///")
            && !trimmed.starts_with("//!")
            && !trimmed.starts_with('#')
    }).map(|pos| start + pos)
}

// ---------------------------------------------------------------------------
//  Tool: SemanticSearch — search the symbol index
// ---------------------------------------------------------------------------

/// Search for code symbols by name, concept, or semantic meaning.
///
/// Flow: ensure index is built → search by query → return formatted results.
pub struct SemanticSearch;

impl Tool for SemanticSearch {
    fn name(&self) -> &'static str {
        "semantic_search"
    }

    fn description(&self) -> &'static str {
        "Search for code symbols (functions, structs, enums, traits) by name, concept, or meaning"
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query — function name, struct name, or concept (e.g. 'payment handler', 'auth middleware', 'user repository')"
                },
                "kind": {
                    "type": "string",
                    "enum": ["fn", "struct", "enum", "trait", "mod", "all"],
                    "description": "Filter by symbol kind (default: all)",
                    "default": "all"
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum results (default 10, max 30)",
                    "default": 10
                },
                "rebuild_index": {
                    "type": "boolean",
                    "description": "Force rebuild the symbol index before searching (default false)",
                    "default": false
                }
            },
            "required": ["query"]
        })
    }

    #[instrument(skip(self, ctx, args))]
    fn run(&self, ctx: &ToolCtx, args: &Value) -> Result<String> {
        let query = crate::tools::arg_str(args, "query")?;
        let kind_filter = args
            .get("kind")
            .and_then(|v| v.as_str())
            .unwrap_or("all");
        let max_results = args
            .get("max_results")
            .and_then(|v| v.as_u64())
            .unwrap_or(10)
            .min(30) as usize;
        let rebuild = args
            .get("rebuild_index")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        let workspace = ctx
            .workspaces
            .first()
            .map(|p| p.to_string_lossy().to_string())
            .unwrap_or_else(|| ".".to_string());

        info!(query = %query, kind = %kind_filter, max_results, rebuild, "semantic search");

        // Get or rebuild the index
        let mut guard = SYMBOL_INDEX.lock().map_err(|e| anyhow::anyhow!("index lock failed: {e}"))?;
        let index = guard.get_or_insert_with(SymbolIndex::new);

        if rebuild || index.is_empty() {
            let count = index.rebuild(&workspace)?;
            debug!(symbol_count = count, "symbol index rebuilt");
        }

        let results = index.search(&query, max_results * 2); // Get extra for filtering

        // Apply kind filter
        let filtered: Vec<&&CodeSymbol> = if kind_filter != "all" {
            let target_kind = match kind_filter {
                "fn" => SymbolKind::Function,
                "struct" => SymbolKind::Struct,
                "enum" => SymbolKind::Enum,
                "trait" => SymbolKind::Trait,
                "mod" => SymbolKind::Module,
                _ => SymbolKind::Other,
            };
            results
                .iter()
                .filter(|s| s.kind == target_kind)
                .take(max_results)
                .collect()
        } else {
            results.iter().take(max_results).collect()
        };

        if filtered.is_empty() {
            return Ok(format!(
                "No symbols found matching '{query}'.\n\
                 Try a different query, or use `rebuild_index: true` to rebuild the index first."
            ));
        }

        let total = index.len();
        info!(matched = filtered.len(), total_indexed = total, "semantic search completed");

        // Group results by file for cleaner output
        let mut by_file: std::collections::BTreeMap<String, Vec<&&CodeSymbol>> =
            std::collections::BTreeMap::new();
        for sym in &filtered {
            by_file.entry(sym.file.clone()).or_default().push(*sym);
        }

        let mut output = format!(
            "## Semantic Search Results\n\n**Query:** {query}\n**Index size:** {total} symbols\n**Matches:** {}\n\n",
            filtered.len()
        );

        for (file, symbols) in &by_file {
            output.push_str(&format!("### `{file}`\n\n"));
            for sym in symbols {
                let kind_str = sym.kind.to_string();
                let parent_str = sym
                    .parent
                    .as_ref()
                    .map(|p| format!(" [{p}]"))
                    .unwrap_or_default();
                let doc_str = sym
                    .doc_comment
                    .as_ref()
                    .map(|d| {
                        let truncated: String = d.chars().take(100).collect();
                        format!(" — {truncated}")
                    })
                    .unwrap_or_default();

                let context_trimmed = sym.context.trim();
                let context_ellipsis = if context_trimmed.len() > 80 { "…" } else { "" };

                output.push_str(&format!(
                    "- `{kind_str}` **{}**{} at line {} `{}`{}{}\n",
                    sym.name,
                    parent_str,
                    sym.line,
                    context_trimmed,
                    doc_str,
                    context_ellipsis
                ));
            }
            output.push('\n');
        }

        output.push_str(&format!(
            "---\n*{} symbols indexed. Use `rebuild_index: true` to refresh.*\n",
            total
        ));

        Ok(output)
    }
}

/// Tool: Rebuild the symbol index explicitly.
pub struct RebuildIndex;

impl Tool for RebuildIndex {
    fn name(&self) -> &'static str {
        "rebuild_index"
    }

    fn description(&self) -> &'static str {
        "Rebuild the code symbol index for semantic search"
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {}
        })
    }

    #[instrument(skip(self, ctx, _args))]
    fn run(&self, ctx: &ToolCtx, _args: &Value) -> Result<String> {
        let workspace = ctx
            .workspaces
            .first()
            .map(|p| p.to_string_lossy().to_string())
            .unwrap_or_else(|| ".".to_string());

        info!("rebuilding symbol index");

        let mut guard = SYMBOL_INDEX.lock().map_err(|e| anyhow::anyhow!("index lock failed: {e}"))?;
        let index = guard.get_or_insert_with(SymbolIndex::new);
        let count = index.rebuild(&workspace)?;

        Ok(format!(
            "Symbol index rebuilt successfully. {} symbols indexed.",
            count
        ))
    }
}