leindex 2.1.1

use super::helpers::{extract_bool, extract_string, extract_usize, resolve_scope, wrap_with_meta};
use super::protocol::JsonRpcError;
use super::request_meta::WorkBudget;
use crate::cli::registry::ProjectRegistry;
use serde_json::Value;
use std::sync::Arc;
use std::time::Instant;

fn classify_search(
    query: &str,
    search_mode: &str,
) -> (
    crate::search::query_route::QueryRoute,
    &'static str,
    Option<crate::search::ranking::QueryType>,
) {
    use crate::search::query_route::{QueryRoute, RequestedMode};

    // When the caller explicitly requests "prose" mode, bypass the auto-routing
    // classification that may reclassify a single identifier as ExactSymbol.
    // Prose mode always uses text scoring, regardless of query shape.
    if search_mode == "prose" {
        return (
            QueryRoute::ExactText,
            "exact_text",
            Some(crate::search::ranking::QueryType::Exact),
        );
    }

    let requested = match search_mode {
        "exact" => RequestedMode::Exact,
        "semantic" | "code" => RequestedMode::Semantic,
        _ => RequestedMode::Auto,
    };
    let route = crate::search::query_route::classify(query, requested);
    let route_name = match route {
        QueryRoute::ExactSymbol => "exact_symbol",
        QueryRoute::ExactText => "exact_text",
        QueryRoute::Semantic => "semantic",
        QueryRoute::DeepPdg => "deep_pdg",
    };
    let query_type = match route {
        QueryRoute::ExactSymbol | QueryRoute::ExactText => {
            Some(crate::search::ranking::QueryType::Exact)
        }
        QueryRoute::Semantic | QueryRoute::DeepPdg => Some(if search_mode == "prose" {
            crate::search::ranking::QueryType::Text
        } else {
            crate::search::ranking::QueryType::Semantic
        }),
    };
    (route, route_name, query_type)
}

fn path_in_scope(file_path: &str, scope: Option<&str>, project_root: &std::path::Path) -> bool {
    let Some(scope) = scope else {
        return true;
    };
    let normalize = |path: &str| path.replace('\\', "/").trim_end_matches('/').to_string();
    let project_root = project_root
        .canonicalize()
        .unwrap_or_else(|_| project_root.to_path_buf());
    let scope = normalize(scope);
    let scope_path = std::path::Path::new(&scope);
    let scope_is_directory = scope.ends_with('/')
        || scope.ends_with('\\')
        || if scope_path.is_absolute() {
            scope_path.is_dir()
        } else {
            project_root.join(scope_path).is_dir()
        };
    let absolute_scope = if scope_path.is_absolute() {
        scope
    } else {
        normalize(&project_root.join(scope_path).to_string_lossy())
    };
    let file_path = normalize(file_path);
    let absolute_file_path = if std::path::Path::new(&file_path).is_absolute() {
        file_path
    } else {
        normalize(&project_root.join(&file_path).to_string_lossy())
    };
    absolute_file_path == absolute_scope
        || (scope_is_directory && absolute_file_path.starts_with(&format!("{absolute_scope}/")))
}

fn scoped_search(
    index: &mut crate::cli::leindex::LeIndex,
    query: &str,
    top_k: usize,
    offset: usize,
    query_type: Option<crate::search::ranking::QueryType>,
    ephemeral: bool,
    scope: Option<&str>,
    project_root: &std::path::Path,
) -> Result<Vec<crate::search::search::SearchResult>, JsonRpcError> {
    const MAX_FETCH_K: usize = 10_000;
    let search = |index: &mut crate::cli::leindex::LeIndex, fetch_k| {
        if ephemeral {
            index.search_ephemeral(query, fetch_k, query_type)
        } else {
            index.search(query, fetch_k, query_type)
        }
    };
    let required = offset.saturating_add(top_k);
    let mut fetch_k = required.clamp(1, MAX_FETCH_K);
    let mut all_results = search(index, fetch_k)
        .map_err(|e| JsonRpcError::search_failed(format!("Search error: {}", e)))?;
    let mut filtered: Vec<_> = all_results
        .iter()
        .filter(|result| path_in_scope(&result.file_path, scope, project_root))
        .cloned()
        .collect();

    while filtered.len() < required && all_results.len() >= fetch_k && fetch_k < MAX_FETCH_K {
        let next_fetch_k = fetch_k.saturating_mul(2).min(MAX_FETCH_K);
        if next_fetch_k == fetch_k {
            break;
        }
        fetch_k = next_fetch_k;
        all_results = search(index, fetch_k)
            .map_err(|e| JsonRpcError::search_failed(format!("Search error: {}", e)))?;
        filtered = all_results
            .iter()
            .filter(|result| path_in_scope(&result.file_path, scope, project_root))
            .cloned()
            .collect();
    }
    Ok(filtered)
}

fn retrieval_meta(
    index: &crate::cli::leindex::LeIndex,
    route: crate::search::query_route::QueryRoute,
    route_name: &str,
    budget: &WorkBudget,
    started: Instant,
) -> Value {
    use crate::search::query_route::QueryRoute;

    serde_json::json!({
        "tfidf_status": "fresh",
        "pdg_status": if index.pdg().is_some() { "resident" } else { "not_loaded" },
        "neural_status": if matches!(route, QueryRoute::ExactSymbol | QueryRoute::ExactText) {
            "not_used_exact"
        } else {
            index.neural_status()
        },
        "route": route_name,
        "partial": budget.elapsed(started),
        "max_latency_ms": budget.max_latency_ms,
        "allow_partial": budget.allow_partial
    })
}

/// Query, paging and budget parameters for one \[search\] call.
struct SearchParams {
    /// The caller's query, echoed back when nothing matched.
    query: String,
    /// The query as the retriever sees it: task context appended when given.
    effective_query: String,
    /// Result window size.
    top_k: usize,
    /// Results to skip.
    offset: usize,
    /// Enrichment budget for this call.
    budget: WorkBudget,
    /// True when a task context was supplied (selects the ephemeral search).
    has_task_context: bool,
    /// Scoring mode, defaulting to `code`.
    search_mode: String,
}

/// Read the query and its paging/budget knobs.
fn parse_search_params(args: &Value) -> Result<SearchParams, JsonRpcError> {
    let query = extract_string(args, "query")?;
    let top_k = extract_usize(args, "top_k", 10)?;
    let offset = extract_usize(args, "offset", 0)?;
    let search_mode = args
        .get("search_mode")
        .and_then(|v| v.as_str())
        .unwrap_or("code")
        .to_string();
    let task_context = args
        .get("task_context")
        .and_then(Value::as_str)
        .map(|context| context.chars().take(2000).collect::<String>());
    let budget = WorkBudget {
        max_latency_ms: extract_usize(args, "max_latency_ms", 500)?.min(60000) as u64,
        allow_partial: extract_bool(args, "allow_partial", true),
    };
    let effective_query = task_context.as_deref().map_or_else(
        || query.clone(),
        |context| format!("{}\nTask context: {}", query, context),
    );
    Ok(SearchParams {
        has_task_context: task_context.is_some(),
        query,
        effective_query,
        top_k,
        offset,
        budget,
        search_mode,
    })
}

/// Project handles hydrate lazily (graph-only tools never build the TF-IDF
/// engine), so a cold MCP session reaches here with an empty engine even for
/// an indexed project. The one-shot CLI pre-loads everything, which is why
/// `leindex search` worked while the MCP tool reported "Project not indexed".
fn load_search_context(index: &mut crate::cli::leindex::LeIndex) -> Result<(), JsonRpcError> {
    index
        .ensure_analysis_context_loaded()
        .map_err(|e| JsonRpcError::indexing_failed(format!("Failed to load search context: {e}")))
}

/// A project with no loaded engine cannot answer a semantic query.
///
/// Distinguishes never-indexed (the classic actionable guidance) from
/// indexed-but-unhydrated (a hydration failure masquerading as "not
/// indexed" — the permanent lie from the multi-writer incident: persisted
/// stats say indexed, but the engine never built). The latter names the
/// actual failure and its remediation instead of sending the user to
/// re-index an already-indexed project.
fn ensure_searchable(index: &crate::cli::leindex::LeIndex) -> Result<(), JsonRpcError> {
    if index.search_engine().is_empty() {
        if index.is_indexed() {
            return Err(JsonRpcError::indexing_failed(format!(
                "Search engine failed to hydrate for {} (the index exists on disk with {} \
                 indexed nodes, but the in-memory engine is empty). Retry the tool call; if \
                 it persists, run `leindex index --force` in {} to rebuild.",
                index.project_path().display(),
                index.get_stats().indexed_nodes,
                index.project_path().display(),
            )));
        }
        return Err(JsonRpcError::project_not_indexed(
            index.project_path().display().to_string(),
        ));
    }
    Ok(())
}

/// The empty-result response: an actionable suggestion plus retrieval meta.
fn empty_search_payload(
    query: &str,
    index: &mut crate::cli::leindex::LeIndex,
    offset: usize,
    route: crate::search::query_route::QueryRoute,
    route_name: &str,
    budget: &WorkBudget,
    started: Instant,
) -> Value {
    let indexed_files = index.source_file_paths().map(|p| p.len()).unwrap_or(0);
    wrap_with_meta(
        serde_json::json!({
            "results": [],
            "offset": offset,
            "count": 0,
            "has_more": false,
            "suggestion": format!(
                "No semantic matches found for '{}'. The project contains {} indexed files. \
                Try: rephrase query, use different keywords, or try leindex_explore mode=find target=symbols for exact symbol names.",
                query,
                indexed_files
            ),
            "retrieval": retrieval_meta(index, route, route_name, budget, started)
        }),
        index,
    )
}

/// Low-signal detection (F-07): expose the top composite score so agents can
/// judge confidence, and flag results whose best match is weak. The floor is
/// deliberately conservative (0.25): TF-IDF gives partial credit for a single
/// shared token, so a garbage query with one common word can still score ~0.7
/// — those are not flaggable without risking false negatives on legitimate
/// short queries, but `top_score` makes the judgment possible at the consumer.
fn low_signal_json(page: &[crate::search::search::SearchResult]) -> Value {
    let top_score = page
        .iter()
        .map(|result| result.score.overall)
        .fold(0.0_f32, f32::max);
    if !page.is_empty() && top_score < 0.25 {
        serde_json::json!({
            "low_signal": true,
            "top_score": top_score,
            "suggestion": "Top match scores below the confidence floor; results may be coincidental token overlap. Rephrase with more specific terms or use leindex_explore mode=find target=symbols for exact names.",
        })
    } else {
        serde_json::json!({ "top_score": top_score })
    }
}

/// Bimodal-latency marker: the first neural call after a server or daemon
/// start pays the ~15–20 s model load; every later call is milliseconds.
/// Without this hint agents timeout-and-retry on the first call of every
/// session (session-5 §5.1). Also merges the low-signal fields.
fn annotate_search_payload(payload: &mut Value, low_signal_json: &Value, started: Instant) {
    if started.elapsed().as_millis() > 5000 {
        if let Some(obj) = payload.as_object_mut() {
            obj.insert("cold_start".to_string(), serde_json::json!(true));
            obj.insert(
                "note".to_string(),
                serde_json::json!(
                    "This call included a one-time neural model load; subsequent searches are milliseconds"
                ),
            );
        }
    }
    if let (Some(obj), Some(target)) = (payload.as_object_mut(), low_signal_json.as_object()) {
        for (key, value) in target {
            obj.insert(key.clone(), value.clone());
        }
    }
}

/// Handler for LeIndex [search
///
/// Performs semantic search on the indexed code.
#[derive(Clone)]
pub struct SearchHandler;

impl SearchHandler {
    /// Returns the name of this MCP tool (MCP-compliant: ASCII letters, digits, underscore, hyphen, dot only)
    pub fn name(&self) -> &str {
        "leindex_search"
    }

    /// Returns the human-readable display title for this tool
    pub fn title(&self) -> &str {
        "LeIndex [Search]"
    }

    /// Returns the description of this RPC method
    pub fn description(&self) -> &str {
        "Semantic code search. Finds symbols by meaning, not just name. Returns ranked \
results with composite scores (semantic + text + structural). Accepts project_path \
to auto-switch/auto-index projects."
    }

    /// Returns the JSON schema for the arguments of this RPC method
    pub fn argument_schema(&self) -> Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query (e.g., 'authentication', 'database connection')"
                },
                "project_path": {
                    "type": "string",
                    "description": "Project directory (auto-indexes on first use; omit to use current project)"
                },
                "top_k": {
                    "type": "integer",
                    "description": "Maximum number of results to return (default: 10)",
                    "default": 10,
                    "minimum": 1,
                    "maximum": 100
                },
                "scope": {
                    "type": "string",
                    "description": "Optional path to limit results (absolute or relative to project root)"
                },
                "offset": {
                    "type": "integer",
                    "description": "Skip the first N results for pagination (default: 0)",
                    "default": 0,
                    "minimum": 0
                },
                "search_mode": {
                    "type": "string",
                    "enum": ["code", "prose", "auto", "exact", "semantic"],
                    "description": "Scoring mode: 'code' (default) emphasizes semantic/structural similarity, \
        'prose' boosts text-match weight for natural-language queries (e.g. roadmap, README content), \
        'auto' detects based on query shape, \
        'exact' prioritizes exact symbol name matches (higher text/structural weights), \
        'semantic' prioritizes conceptual relevance (higher TF-IDF semantic weights).",
                    "default": "code"
                },
                "task_context": {
                    "type": "string",
                    "description": "Optional bounded review/task context used for this retrieval only",
                    "maxLength": 2000
                },
                "max_latency_ms": {
                    "type": "integer",
                    "description": "Optional enrichment budget; never cancels the search (default: 500)",
                    "default": 500,
                    "minimum": 0,
                    "maximum": 60000
                },
                "allow_partial": {
                    "type": "boolean",
                    "description": "Return core TF-IDF results when optional enrichment exceeds the budget",
                    "default": true
                }
            },
            "required": ["query"]
        })
    }

    /// Executes the RPC method
    pub async fn execute(
        &self,
        registry: &Arc<ProjectRegistry>,
        args: Value,
    ) -> Result<Value, JsonRpcError> {
        let params = parse_search_params(&args)?;
        let started = Instant::now();
        let (route, route_name, query_type) =
            classify_search(&params.effective_query, &params.search_mode);

        let project_path = args.get("project_path").and_then(|v| v.as_str());
        let handle = registry.get_or_create(project_path).await?;
        let mut guard = handle.write().await;

        let scope = resolve_scope(&args, guard.project_path())?;

        load_search_context(&mut guard)?;
        ensure_searchable(&guard)?;

        let project_root = guard.project_path().to_path_buf();
        let filtered = scoped_search(
            &mut guard,
            &params.effective_query,
            params.top_k,
            params.offset,
            query_type,
            params.has_task_context,
            scope.as_deref(),
            &project_root,
        )?;

        let total_filtered = filtered.len();
        let page: Vec<_> = filtered
            .into_iter()
            .skip(params.offset)
            .take(params.top_k)
            .collect();
        let total_returned = page.len();

        let low_signal_json = low_signal_json(&page);

        if total_filtered == 0 {
            return Ok(empty_search_payload(
                &params.query,
                &mut guard,
                params.offset,
                route,
                route_name,
                &params.budget,
                started,
            ));
        }

        let mut payload = serde_json::json!({
            "results": serde_json::to_value(&page).map_err(|e|
                JsonRpcError::internal_error(format!("Serialization error: {}", e)))?,
            "offset": params.offset,
            "count": total_returned,
            "has_more": params.offset + total_returned < total_filtered,
            "retrieval": retrieval_meta(&guard, route, route_name, &params.budget, started)
        });
        annotate_search_payload(&mut payload, &low_signal_json, started);
        Ok(wrap_with_meta(payload, &guard))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::mcp::helpers::test_registry_for;
    use tempfile::tempdir;

    #[tokio::test]
    async fn test_search_handler_zero_results_includes_suggestion() {
        // Test that semantic search with no matches returns helpful suggestion.
        // Uses a source file with content unrelated to the query.
        let dir = tempdir().unwrap();
        let src = dir.path().join("lib.rs");
        std::fs::write(&src, "pub fn alpha_beta_gamma() {}\n").unwrap();
        let registry = test_registry_for(dir.path());
        let args = serde_json::json!({ "query": "zzz_nonexistent_qqq_12345" });
        let result = SearchHandler.execute(&registry, args).await;
        // Should succeed
        assert!(result.is_ok(), "search should succeed");
        let val = result.unwrap();
        // With improved scoring, unrelated queries should return 0 results
        // due to no token overlap and no symbol name match
        if val["count"].as_i64().unwrap_or(0) == 0 {
            // Verify suggestion field is present for zero results
            assert!(
                val.get("suggestion").is_some(),
                "zero results should include suggestion"
            );
        }
    }

    fn directory_root() -> &'static std::path::Path {
        std::path::Path::new(".")
    }

    #[test]
    fn test_path_in_scope_handles_files_and_directories() {
        let separator = std::path::MAIN_SEPARATOR;
        let nested = format!("src{separator}cli{separator}main.rs");
        let sibling = format!("src{separator}lib.rs");

        assert!(path_in_scope(&nested, Some("src"), directory_root()));
        assert!(path_in_scope(
            "src\\cli\\main.rs",
            Some("src"),
            directory_root()
        ));
        assert!(!path_in_scope(&sibling, Some(&nested), directory_root()));
        assert!(path_in_scope(&nested, Some(&nested), directory_root()));
        assert!(path_in_scope(&nested, None, directory_root()));
    }

    #[test]
    fn test_path_in_scope_preserves_file_and_dotted_directory_boundaries() {
        let directory = tempdir().unwrap();
        let dotted_dir = directory.path().join("docs.v1");
        std::fs::create_dir(&dotted_dir).unwrap();
        let extensionless_file = directory.path().join("Makefile");
        std::fs::write(&extensionless_file, b"all:\n").unwrap();

        let dotted_dir = dotted_dir.to_string_lossy().replace('\\', "/");
        let extensionless_file = extensionless_file.to_string_lossy().replace('\\', "/");
        let dotted_child = format!("{dotted_dir}/guide.md");
        let false_file_child = format!("{extensionless_file}/nested.rs");

        assert!(path_in_scope(
            &dotted_child,
            Some(&dotted_dir),
            directory.path()
        ));
        assert!(!path_in_scope(
            &false_file_child,
            Some(&extensionless_file),
            directory.path()
        ));
    }

    #[test]
    fn test_search_schema_has_pagination() {
        let handler = SearchHandler;
        let schema = handler.argument_schema();
        let props = schema.get("properties").unwrap();
        assert!(props.get("offset").is_some());
    }
}