lean_ctx/tools/

ctx_graph.rs

use std::collections::HashMap;
use std::path::Path;

use crate::core::graph_index;
use crate::core::graph_provider::{self, GraphProvider};
use crate::core::tokens::count_tokens;

pub fn handle(
    action: &str,
    path: Option<&str>,
    root: &str,
    cache: &mut crate::core::cache::SessionCache,
    crp_mode: crate::tools::CrpMode,
    depth: Option<usize>,
    kind: Option<&str>,
) -> String {
    match action {
        "build" => handle_build(root),
        "related" => handle_related(path, root),
        "symbol" => handle_symbol(path, root, cache, crp_mode),
        "impact" => handle_impact(path, root),
        "status" => handle_status(root),
        "enrich" => handle_enrich(root),
        "context" => handle_context_query(path, root),
        "diagram" => crate::tools::ctx_graph_diagram::handle(path, depth, kind, root),
        _ => {
            "Unknown action. Use: build, related, symbol, impact, status, enrich, context, diagram"
                .to_string()
        }
    }
}

fn handle_build(root: &str) -> String {
    let index = graph_index::scan(root);

    let mut by_lang: HashMap<&str, (usize, usize)> = HashMap::new();
    for entry in index.files.values() {
        let e = by_lang.entry(&entry.language).or_insert((0, 0));
        e.0 += 1;
        e.1 += entry.token_count;
    }

    let mut result = Vec::new();
    result.push(format!(
        "Project Graph: {} files, {} symbols, {} edges",
        index.file_count(),
        index.symbol_count(),
        index.edge_count()
    ));

    let mut langs: Vec<_> = by_lang.iter().collect();
    langs.sort_by_key(|(_, v)| std::cmp::Reverse(v.1));
    result.push("\nLanguages:".to_string());
    for (lang, (count, tokens)) in &langs {
        result.push(format!("  {lang}: {count} files, {tokens} tok"));
    }

    let mut import_counts: HashMap<&str, usize> = HashMap::new();
    for edge in &index.edges {
        if edge.kind == "import" {
            *import_counts.entry(&edge.to).or_insert(0) += 1;
        }
    }
    let mut hotspots: Vec<_> = import_counts.iter().collect();
    hotspots.sort_by_key(|x| std::cmp::Reverse(*x.1));

    if !hotspots.is_empty() {
        result.push(format!("\nMost imported ({}):", hotspots.len().min(10)));
        for (module, count) in hotspots.iter().take(10) {
            result.push(format!("  {module}: imported by {count} files"));
        }
    }

    if let Some(dir) = GraphProvider::index_dir(root) {
        result.push(format!(
            "\nIndex saved: {}",
            crate::core::protocol::shorten_path(&dir.to_string_lossy())
        ));
    }

    let output = result.join("\n");
    let tokens = count_tokens(&output);
    format!("{output}\n[ctx_graph build: {tokens} tok]")
}

fn handle_related(path: Option<&str>, root: &str) -> String {
    let Some(target) = path else {
        return "path is required for 'related' action".to_string();
    };

    let Some(open) = graph_provider::open_or_build(root) else {
        return "No graph index found. Run ctx_graph with action='build' first.".to_string();
    };

    let rel_target = graph_index::graph_relative_key(target, root);

    let related = open.provider.related(&rel_target, 2);
    if related.is_empty() {
        return format!(
            "No related files found for {}",
            crate::core::protocol::shorten_path(target)
        );
    }

    let mut result = format!(
        "Files related to {} ({}):\n",
        crate::core::protocol::shorten_path(target),
        related.len()
    );
    for r in &related {
        result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(r)));
    }

    let tokens = count_tokens(&result);
    format!("{result}[ctx_graph related: {tokens} tok]")
}

fn handle_symbol(
    path: Option<&str>,
    root: &str,
    cache: &mut crate::core::cache::SessionCache,
    crp_mode: crate::tools::CrpMode,
) -> String {
    let Some(spec) = path else {
        return "path is required for 'symbol' action (format: file.rs::function_name)".to_string();
    };

    let Some((file_part, symbol_name)) = spec.split_once("::") else {
        return format!("Invalid symbol spec '{spec}'. Use format: file.rs::function_name");
    };

    let Some(open) = graph_provider::open_or_build(root) else {
        return "No graph index found. Run ctx_graph with action='build' first.".to_string();
    };

    let rel_file = graph_index::graph_relative_key(file_part, root);

    let key = format!("{rel_file}::{symbol_name}");
    let Some(symbol) = open.provider.get_symbol(&key) else {
        let available = open
            .provider
            .find_symbols(symbol_name, Some(&rel_file), None);
        if available.is_empty() {
            return format!("Symbol '{symbol_name}' not found in {rel_file}. Run ctx_graph action='build' to update the index.");
        }
        let names: Vec<String> = available
            .iter()
            .take(10)
            .map(|s| format!("{}::{}", s.file, s.name))
            .collect();
        return format!(
            "Symbol '{symbol_name}' not found in {rel_file}.\nAvailable symbols:\n  {}",
            names.join("\n  ")
        );
    };

    let abs_path = if Path::new(file_part).is_absolute() {
        file_part.to_string()
    } else {
        Path::new(root)
            .join(rel_file.trim_start_matches(['/', '\\']))
            .to_string_lossy()
            .to_string()
    };

    let content = match std::fs::read_to_string(&abs_path) {
        Ok(c) => c,
        Err(e) => return format!("Cannot read {abs_path}: {e}"),
    };

    let lines: Vec<&str> = content.lines().collect();
    let start = symbol.start_line.saturating_sub(1);
    let end = symbol.end_line.min(lines.len());

    if start >= lines.len() {
        return crate::tools::ctx_read::handle(cache, &abs_path, "full", crp_mode);
    }

    let mut result = format!(
        "{}::{} ({}:{}-{})\n",
        crate::core::protocol::shorten_path(&rel_file),
        symbol_name,
        symbol.kind,
        symbol.start_line,
        symbol.end_line
    );

    for (i, line) in lines[start..end].iter().enumerate() {
        result.push_str(&format!("{:>4}|{}\n", start + i + 1, line));
    }

    let tokens = count_tokens(&result);
    let full_tokens = count_tokens(&content);
    let saved = full_tokens.saturating_sub(tokens);
    let pct = if full_tokens > 0 {
        (saved as f64 / full_tokens as f64 * 100.0).round() as usize
    } else {
        0
    };

    format!("{result}[ctx_graph symbol: {tokens} tok (full file: {full_tokens} tok, -{pct}%)]")
}

fn file_path_to_module_prefixes(
    rel_path: &str,
    project_root: &str,
    provider: &GraphProvider,
) -> Vec<String> {
    let rel_path_slash = graph_index::graph_match_key(rel_path);
    let without_ext = rel_path_slash
        .strip_suffix(".rs")
        .or_else(|| rel_path_slash.strip_suffix(".ts"))
        .or_else(|| rel_path_slash.strip_suffix(".tsx"))
        .or_else(|| rel_path_slash.strip_suffix(".js"))
        .or_else(|| rel_path_slash.strip_suffix(".py"))
        .or_else(|| rel_path_slash.strip_suffix(".kt"))
        .or_else(|| rel_path_slash.strip_suffix(".kts"))
        .unwrap_or(&rel_path_slash);

    let module_path = without_ext
        .strip_prefix("src/")
        .unwrap_or(without_ext)
        .replace('/', "::");

    let module_path = if module_path.ends_with("::mod") {
        module_path
            .strip_suffix("::mod")
            .unwrap_or(&module_path)
            .to_string()
    } else {
        module_path
    };

    let crate_name = std::fs::read_to_string(Path::new(project_root).join("Cargo.toml"))
        .or_else(|_| std::fs::read_to_string(Path::new(project_root).join("package.json")))
        .ok()
        .and_then(|c| {
            c.lines()
                .find(|l| l.contains("\"name\"") || l.starts_with("name"))
                .and_then(|l| l.split('"').nth(1))
                .map(|n| n.replace('-', "_"))
        })
        .unwrap_or_default();

    let mut prefixes = vec![
        format!("crate::{module_path}"),
        format!("super::{module_path}"),
        module_path.clone(),
    ];
    if !crate_name.is_empty() {
        prefixes.insert(0, format!("{crate_name}::{module_path}"));
    }

    let ext = Path::new(rel_path)
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("");
    if matches!(ext, "kt" | "kts") {
        let abs_path = Path::new(project_root).join(rel_path.trim_start_matches(['/', '\\']));
        if let Ok(content) = std::fs::read_to_string(abs_path) {
            if let Some(package_name) = content.lines().map(str::trim).find_map(|line| {
                line.strip_prefix("package ")
                    .map(|rest| rest.trim().trim_end_matches(';').to_string())
            }) {
                prefixes.push(package_name.clone());
                if let Some(entry) = provider.get_file_entry(rel_path) {
                    for export in &entry.exports {
                        prefixes.push(format!("{package_name}.{export}"));
                    }
                }
                if let Some(file_stem) = Path::new(rel_path).file_stem().and_then(|s| s.to_str()) {
                    prefixes.push(format!("{package_name}.{file_stem}"));
                }
            }
        }
    }

    prefixes.sort();
    prefixes.dedup();
    prefixes
}

fn edge_matches_file(edge_to: &str, module_prefixes: &[String]) -> bool {
    module_prefixes.iter().any(|prefix| {
        edge_to == *prefix
            || edge_to.starts_with(&format!("{prefix}::"))
            || edge_to.starts_with(&format!("{prefix},"))
    })
}

fn handle_impact(path: Option<&str>, root: &str) -> String {
    let Some(target) = path else {
        return "path is required for 'impact' action".to_string();
    };

    let Some(open) = graph_provider::open_or_build(root) else {
        return "No graph index found. Run ctx_graph with action='build' first.".to_string();
    };
    let gp = &open.provider;

    let rel_target = graph_index::graph_relative_key(target, root);
    let module_prefixes = file_path_to_module_prefixes(&rel_target, root, gp);

    let import_edges = gp.edges_by_kind("import");
    let direct: Vec<&str> = import_edges
        .iter()
        .filter(|e| edge_matches_file(&e.to, &module_prefixes))
        .map(|e| e.from.as_str())
        .collect();

    let mut all_dependents: Vec<String> = direct
        .iter()
        .map(std::string::ToString::to_string)
        .collect();
    for d in &direct {
        for dep in gp.dependents(d) {
            if !all_dependents.contains(&dep) && dep != rel_target {
                all_dependents.push(dep);
            }
        }
    }

    if all_dependents.is_empty() {
        return format!(
            "No files depend on {}",
            crate::core::protocol::shorten_path(target)
        );
    }

    let mut result = format!(
        "Impact of {} ({} dependents):\n",
        crate::core::protocol::shorten_path(target),
        all_dependents.len()
    );

    if !direct.is_empty() {
        result.push_str(&format!("\nDirect ({}):\n", direct.len()));
        for d in &direct {
            result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(d)));
        }
    }

    let indirect: Vec<&String> = all_dependents
        .iter()
        .filter(|d| !direct.contains(&d.as_str()))
        .collect();
    if !indirect.is_empty() {
        result.push_str(&format!("\nIndirect ({}):\n", indirect.len()));
        for d in &indirect {
            result.push_str(&format!("  {}\n", crate::core::protocol::shorten_path(d)));
        }
    }

    let tokens = count_tokens(&result);
    format!("{result}[ctx_graph impact: {tokens} tok]")
}

fn handle_status(root: &str) -> String {
    let Some(open) = graph_provider::open_best_effort(root) else {
        return "No graph index. Run ctx_graph action='build' to create one.".to_string();
    };
    let gp = &open.provider;

    let file_paths = gp.file_paths();
    let mut by_lang: HashMap<String, usize> = HashMap::new();
    let mut total_tokens = 0usize;
    for path in &file_paths {
        if let Some(entry) = gp.get_file_entry(path) {
            *by_lang.entry(entry.language).or_insert(0) += 1;
            total_tokens += entry.token_count;
        }
    }

    let mut langs: Vec<_> = by_lang.iter().collect();
    langs.sort_by_key(|item| std::cmp::Reverse(*item.1));
    let lang_summary: String = langs
        .iter()
        .take(5)
        .map(|(l, c)| format!("{l}:{c}"))
        .collect::<Vec<_>>()
        .join(" ");

    format!(
        "Graph: {} files, {} symbols, {} edges ({:?}) | {} tok total\nLast scan: {}\nLanguages: {lang_summary}\nStored: {}",
        gp.file_count(),
        gp.symbol_count(),
        gp.edge_count().unwrap_or(0),
        open.source,
        total_tokens,
        gp.last_scan(),
        GraphProvider::index_dir(root)
            .map(|d| d.to_string_lossy().to_string())
            .unwrap_or_default()
    )
}

fn resolve_node_name(graph: &crate::core::property_graph::CodeGraph, node_id: i64) -> String {
    let conn = graph.connection();
    conn.query_row(
        "SELECT name FROM nodes WHERE id = ?1",
        rusqlite::params![node_id],
        |row| row.get::<_, String>(0),
    )
    .unwrap_or_else(|_| format!("node#{node_id}"))
}

fn handle_enrich(root: &str) -> String {
    let graph = match crate::core::property_graph::CodeGraph::open(root) {
        Ok(g) => g,
        Err(e) => return format!("Failed to open graph: {e}"),
    };

    match crate::core::graph_enricher::enrich_graph(&graph, Path::new(root), 500) {
        Ok(stats) => {
            let node_count = graph.node_count().unwrap_or(0);
            let edge_count = graph.edge_count().unwrap_or(0);
            format!(
                "Graph enriched.\n{}\nTotal: {node_count} nodes, {edge_count} edges",
                stats.format_summary()
            )
        }
        Err(e) => format!("Enrichment failed: {e}"),
    }
}

fn handle_context_query(query: Option<&str>, root: &str) -> String {
    let Some(query) = query else {
        return "Usage: ctx_graph action=context path=\"<query or file path>\"".to_string();
    };

    let graph = match crate::core::property_graph::CodeGraph::open(root) {
        Ok(g) => g,
        Err(e) => return format!("Failed to open graph: {e}"),
    };

    let gp = graph_provider::open_or_build(root);
    let mut result = Vec::new();

    if let Ok(Some(node)) = graph.get_node_by_path(query) {
        result.push(format!("## Context for `{query}`\n"));

        if let Some(node_id) = node.id {
            let edges_out = graph.edges_from(node_id).unwrap_or_default();
            let edges_in = graph.edges_to(node_id).unwrap_or_default();

            let mut tests: Vec<String> = Vec::new();
            let mut commits: Vec<String> = Vec::new();
            let mut knowledge: Vec<String> = Vec::new();
            let mut imports: Vec<String> = Vec::new();
            let mut dependents: Vec<String> = Vec::new();

            for edge in &edges_out {
                let target = resolve_node_name(&graph, edge.target_id);
                match edge.kind {
                    crate::core::property_graph::EdgeKind::TestedBy => tests.push(target),
                    crate::core::property_graph::EdgeKind::ChangedIn => commits.push(target),
                    crate::core::property_graph::EdgeKind::MentionedIn => {
                        knowledge.push(target);
                    }
                    crate::core::property_graph::EdgeKind::Imports => imports.push(target),
                    _ => {}
                }
            }

            for edge in &edges_in {
                let source = resolve_node_name(&graph, edge.source_id);
                if edge.kind == crate::core::property_graph::EdgeKind::Imports {
                    dependents.push(source);
                }
            }

            if !tests.is_empty() {
                result.push(format!("**Tests ({}):** {}", tests.len(), tests.join(", ")));
            }
            if !commits.is_empty() {
                result.push(format!(
                    "**Recent commits ({}):** {}",
                    commits.len(),
                    commits
                        .iter()
                        .take(5)
                        .cloned()
                        .collect::<Vec<_>>()
                        .join(", ")
                ));
            }
            if !knowledge.is_empty() {
                result.push(format!(
                    "**Knowledge ({}):** {}",
                    knowledge.len(),
                    knowledge.join(", ")
                ));
            }
            if !imports.is_empty() {
                result.push(format!(
                    "**Imports ({}):** {}",
                    imports.len(),
                    imports
                        .iter()
                        .take(10)
                        .cloned()
                        .collect::<Vec<_>>()
                        .join(", ")
                ));
            }
            if !dependents.is_empty() {
                result.push(format!(
                    "**Depended on by ({}):** {}",
                    dependents.len(),
                    dependents
                        .iter()
                        .take(10)
                        .cloned()
                        .collect::<Vec<_>>()
                        .join(", ")
                ));
            }

            if let Ok(impact) = graph.impact_analysis(query, 3) {
                if !impact.affected_files.is_empty() {
                    result.push(format!(
                        "**Impact radius:** {} files within 3 hops",
                        impact.affected_files.len()
                    ));
                }
            }
        }
    } else {
        result.push(format!("## Search: `{query}`\n"));
        let related = gp
            .as_ref()
            .map(|o| o.provider.related(query, 2))
            .unwrap_or_default();
        if related.is_empty() {
            result.push("No matching nodes found in graph.".to_string());
        } else {
            result.push(format!("**Related files ({}):**", related.len()));
            for f in related.iter().take(15) {
                result.push(format!("  - {f}"));
            }
        }
    }

    result.join("\n")
}

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

    #[test]
    fn test_edge_matches_file_crate_prefix() {
        let prefixes = vec![
            "lean_ctx::core::cache".to_string(),
            "crate::core::cache".to_string(),
            "super::core::cache".to_string(),
            "core::cache".to_string(),
        ];
        assert!(edge_matches_file(
            "lean_ctx::core::cache::SessionCache",
            &prefixes
        ));
        assert!(edge_matches_file(
            "crate::core::cache::SessionCache",
            &prefixes
        ));
        assert!(edge_matches_file("crate::core::cache", &prefixes));
        assert!(!edge_matches_file(
            "lean_ctx::core::config::Config",
            &prefixes
        ));
        assert!(!edge_matches_file("crate::core::cached_reader", &prefixes));
    }

    #[test]
    fn test_file_path_to_module_prefixes_rust() {
        let gp =
            GraphProvider::GraphIndex(crate::core::graph_index::ProjectIndex::new("/nonexistent"));
        let prefixes = file_path_to_module_prefixes("src/core/cache.rs", "/nonexistent", &gp);
        assert!(prefixes.contains(&"crate::core::cache".to_string()));
        assert!(prefixes.contains(&"core::cache".to_string()));
    }

    #[test]
    fn test_file_path_to_module_prefixes_mod_rs() {
        let gp =
            GraphProvider::GraphIndex(crate::core::graph_index::ProjectIndex::new("/nonexistent"));
        let prefixes = file_path_to_module_prefixes("src/core/mod.rs", "/nonexistent", &gp);
        assert!(prefixes.contains(&"crate::core".to_string()));
        assert!(!prefixes.iter().any(|p| p.contains("mod")));
    }
}