scryer-engine 0.2.1

use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use stack_graphs::arena::Handle;
use stack_graphs::graph::{File, Node, StackGraph};
use stack_graphs::partial::{PartialPath, PartialPaths};
use stack_graphs::stitching::{
    Database, DatabaseCandidates, ForwardCandidates, ForwardPartialPathStitcher, StitcherConfig,
};
use tree_sitter_stack_graphs::{NoCancellation, StackGraphLanguage, Variables};

pub static RUST_TSG: &str = include_str!("../../rules/rust/stack-graphs.tsg");
pub static PYTHON_TSG: &str = include_str!("../../rules/python/stack-graphs.tsg");
pub static TYPESCRIPT_TSG: &str = include_str!("../../rules/typescript/stack-graphs.tsg");

/// Compiles and initializes the Rust `StackGraphLanguage` instance from vendored TSG rules.
pub fn load_rust_language() -> anyhow::Result<StackGraphLanguage> {
    let language = tree_sitter_rust::LANGUAGE;
    let sgl = StackGraphLanguage::from_str(language.into(), RUST_TSG)
        .map_err(|e| anyhow::anyhow!("Failed to parse Rust TSG rules: {e}"))?;
    Ok(sgl)
}

/// Compiles and initializes the Python `StackGraphLanguage` instance from vendored TSG rules.
pub fn load_python_language() -> anyhow::Result<StackGraphLanguage> {
    let language = tree_sitter_python::LANGUAGE;
    let sgl = StackGraphLanguage::from_str(language.into(), PYTHON_TSG)
        .map_err(|e| anyhow::anyhow!("Failed to parse Python TSG rules: {e}"))?;
    Ok(sgl)
}

/// Compiles and initializes the TypeScript `StackGraphLanguage` instance from vendored TSG rules.
pub fn load_typescript_language() -> anyhow::Result<StackGraphLanguage> {
    let language = tree_sitter_typescript::LANGUAGE_TYPESCRIPT;
    let sgl = StackGraphLanguage::from_str(language.into(), TYPESCRIPT_TSG)
        .map_err(|e| anyhow::anyhow!("Failed to parse TypeScript TSG rules: {e}"))?;
    Ok(sgl)
}

/// Information about a symbol definition resolved through Stack Graph paths or heuristics.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolvedDefinition {
    /// Identifier name of the resolved symbol.
    pub symbol_name: String,
    /// Relative or absolute path to the file declaring the symbol.
    pub file_path: PathBuf,
    /// Semantic kind / syntax type (e.g. "function", "struct", "module", "trait").
    pub syntax_type: Option<String>,
    /// 1-indexed start line.
    pub start_line: u32,
    /// 0-indexed start column offset.
    pub start_col: u32,
    /// 1-indexed end line.
    pub end_line: u32,
    /// 0-indexed end column offset.
    pub end_col: u32,
    /// Whether this definition resides in an external dependency.
    pub is_external: bool,
    /// Name of the external crate if the symbol is in a dependency.
    pub crate_name: Option<String>,
}

/// A cancellation flag that terminates traversals if iteration steps exceed `max_steps`.
pub struct StepBoundedCancellationFlag<'a> {
    pub max_steps: usize,
    pub steps: Arc<AtomicUsize>,
    pub inner: Option<&'a dyn stack_graphs::CancellationFlag>,
}

impl<'a> StepBoundedCancellationFlag<'a> {
    pub fn new(max_steps: usize) -> Self {
        Self {
            max_steps,
            steps: Arc::new(AtomicUsize::new(0)),
            inner: None,
        }
    }

    pub fn with_inner(max_steps: usize, inner: &'a dyn stack_graphs::CancellationFlag) -> Self {
        Self {
            max_steps,
            steps: Arc::new(AtomicUsize::new(0)),
            inner: Some(inner),
        }
    }
}

impl stack_graphs::CancellationFlag for StepBoundedCancellationFlag<'_> {
    fn check(&self, at: &'static str) -> Result<(), stack_graphs::CancellationError> {
        let count = self.steps.fetch_add(1, Ordering::Relaxed);
        if count >= self.max_steps {
            return Err(stack_graphs::CancellationError(
                "Execution step limit exceeded (path explosion protection)",
            ));
        }
        if let Some(inner) = self.inner {
            inner.check(at)?;
        }
        Ok(())
    }
}

/// Core in-memory Stack Graph engine for a project workspace.
pub struct StackGraphEngine {
    pub graph: StackGraph,
    pub partials: PartialPaths,
    pub database: Database,
    pub sgl: Arc<StackGraphLanguage>,
    pub languages: HashMap<String, Arc<StackGraphLanguage>>,
    pub files: HashMap<PathBuf, Handle<File>>,
    pub max_steps: usize,
    /// Source of every added file, in insertion order, for rebuilding the graph.
    sources: Vec<(PathBuf, String)>,
    /// Files whose partial paths are already in `database`.
    precomputed: HashSet<Handle<File>>,
    /// A file changed after it was built; the graph must be rebuilt before use.
    stale: bool,
    /// Crate root definition nodes keyed by crate name.
    pub crate_roots: HashMap<String, Handle<Node>>,
    /// Handles of external dependency files.
    pub dependency_files: HashSet<Handle<File>>,
    /// Maps file handles to their declaring crate name.
    pub file_to_crate: HashMap<Handle<File>, String>,
    /// Registered crate roots for re-application across rebuilds.
    registered_crates: Vec<(String, PathBuf)>,
    /// Registered re-export aliases for re-application across rebuilds.
    registered_reexports: Vec<(PathBuf, Option<String>, String, String)>,
}

impl StackGraphEngine {
    /// Create a new engine for the given language rules.
    pub fn new(sgl: Arc<StackGraphLanguage>) -> Self {
        let mut languages = HashMap::new();
        languages.insert("rs".to_string(), Arc::clone(&sgl));

        match load_python_language() {
            Ok(py_sgl) => {
                languages.insert("py".to_string(), Arc::new(py_sgl));
            }
            Err(e) => tracing::warn!(
                "Python stack-graph rules failed to load; Python resolution falls back to heuristics: {e}"
            ),
        }

        match load_typescript_language() {
            Ok(ts_sgl) => {
                let ts_arc = Arc::new(ts_sgl);
                for ext in ["ts", "tsx", "js", "jsx"] {
                    languages.insert(ext.to_string(), Arc::clone(&ts_arc));
                }
            }
            Err(e) => tracing::warn!(
                "TypeScript stack-graph rules failed to load; TypeScript resolution falls back to heuristics: {e}"
            ),
        }

        Self {
            graph: StackGraph::new(),
            partials: PartialPaths::new(),
            database: Database::new(),
            sgl,
            languages,
            files: HashMap::new(),
            max_steps: 1000,
            sources: Vec::new(),
            precomputed: HashSet::new(),
            stale: false,
            crate_roots: HashMap::new(),
            dependency_files: HashSet::new(),
            file_to_crate: HashMap::new(),
            registered_crates: Vec::new(),
            registered_reexports: Vec::new(),
        }
    }

    /// Convenience constructor initializing the engine with vendored Rust TSG rules.
    pub fn new_rust() -> anyhow::Result<Self> {
        let sgl = Arc::new(load_rust_language()?);
        Ok(Self::new(sgl))
    }

    /// Configure maximum path resolution traversal steps (default: 1000).
    pub fn with_max_steps(mut self, max_steps: usize) -> Self {
        self.max_steps = max_steps;
        self
    }

    /// Add a source file to the graph and build its partial stack graph using TSG rules.
    ///
    /// Stack graphs can't drop a file's nodes, so re-adding a file with new content marks
    /// the graph stale; it is rebuilt from all stored sources before its next use. File
    /// handles are stable across rebuilds.
    pub fn add_file(&mut self, rel_path: &Path, content: &str) -> anyhow::Result<Handle<File>> {
        if let Some(&handle) = self.files.get(rel_path) {
            if let Some((_, source)) = self.sources.iter_mut().find(|(p, _)| p == rel_path)
                && source != content
            {
                *source = content.to_string();
                self.stale = true;
            }
            return Ok(handle);
        }
        self.sources
            .push((rel_path.to_path_buf(), content.to_string()));
        self.build_file(rel_path, content)
    }

    /// Drop a file from the graph; takes effect at the next rebuild.
    pub fn remove_file(&mut self, rel_path: &Path) {
        if self.files.remove(rel_path).is_some() {
            self.sources.retain(|(p, _)| p != rel_path);
            self.stale = true;
        }
    }

    fn build_file(&mut self, rel_path: &Path, content: &str) -> anyhow::Result<Handle<File>> {
        let path_str = rel_path.to_string_lossy().to_string();
        let file_handle = self.graph.get_or_create_file(&path_str);
        self.files.insert(rel_path.to_path_buf(), file_handle);

        let globals = Variables::new();
        let cancellation = NoCancellation;

        // Dependency crate roots are registered under synthetic paths; those are Rust.
        let ext = rel_path
            .extension()
            .and_then(|e| e.to_str())
            .unwrap_or("rs")
            .to_ascii_lowercase();

        let Some(sgl) = self.languages.get(&ext) else {
            anyhow::bail!(
                "no stack-graph rules for '.{ext}' files ({}); not building a graph for it",
                rel_path.display()
            );
        };

        sgl.build_stack_graph_into(
            &mut self.graph,
            file_handle,
            content,
            &globals,
            &cancellation,
        )
        .map_err(|e| anyhow::anyhow!("TSG build error in {}: {e:?}", rel_path.display()))?;

        Ok(file_handle)
    }

    /// Rebuild the graph from stored sources if any file changed since it was built.
    fn refresh(&mut self, cancellation: &dyn stack_graphs::CancellationFlag) {
        if !self.stale {
            return;
        }
        self.stale = false;
        self.graph = StackGraph::new();
        self.partials = PartialPaths::new();
        self.database = Database::new();
        self.files.clear();
        self.precomputed.clear();
        self.crate_roots.clear();
        self.dependency_files.clear();
        self.file_to_crate.clear();
        for (path, content) in std::mem::take(&mut self.sources) {
            if let Err(e) = self.build_file(&path, &content) {
                tracing::debug!("{e}");
            }
            self.sources.push((path, content));
        }
        let reg_crates = self.registered_crates.clone();
        for (c_name, entry) in reg_crates {
            let _ = self.register_crate_root(&c_name, &entry);
        }
        let reg_reexports = self.registered_reexports.clone();
        for (file_p, _, exp_name, tgt_path) in reg_reexports {
            if let Some(&handle) = self.files.get(&file_p) {
                let _ = self.add_reexport_alias(handle, None, &exp_name, &tgt_path);
            }
        }
        let _ = self.precompute_all_files(cancellation);
    }

    /// Register an external dependency package root and create a root module definition node.
    pub fn register_crate_root(
        &mut self,
        crate_name: &str,
        entry_file_path: &Path,
    ) -> anyhow::Result<Handle<Node>> {
        let file_handle = if let Some(&handle) = self.files.get(entry_file_path) {
            handle
        } else {
            self.add_file(entry_file_path, "")?
        };
        self.dependency_files.insert(file_handle);
        self.file_to_crate
            .insert(file_handle, crate_name.to_string());

        let node_id = self.graph.new_node_id(file_handle);
        let sym = self.graph.add_symbol(crate_name);
        let crate_root_node = self
            .graph
            .add_pop_symbol_node(node_id, sym, true)
            .ok_or_else(|| anyhow::anyhow!("Failed to add crate root node for {crate_name}"))?;
        self.graph
            .add_edge(StackGraph::root_node(), crate_root_node, 0);
        self.crate_roots
            .insert(crate_name.to_string(), crate_root_node);

        if !self.registered_crates.iter().any(|(c, _)| c == crate_name) {
            self.registered_crates
                .push((crate_name.to_string(), entry_file_path.to_path_buf()));
        }

        Ok(crate_root_node)
    }

    /// Add a public module or child item definition under a parent node in the stack graph.
    pub fn add_module_item(
        &mut self,
        file_handle: Handle<File>,
        parent_node: Handle<Node>,
        item_name: &str,
        is_definition: bool,
    ) -> anyhow::Result<Handle<Node>> {
        let node_id = self.graph.new_node_id(file_handle);
        let sym = self.graph.add_symbol(item_name);
        let item_node = self
            .graph
            .add_pop_symbol_node(node_id, sym, is_definition)
            .ok_or_else(|| anyhow::anyhow!("Failed to add module item node {item_name}"))?;
        self.graph.add_edge(parent_node, item_node, 0);
        Ok(item_node)
    }

    /// Add a `pub use` re-export alias node connecting `exported_name` to `target_path`.
    pub fn add_reexport_alias(
        &mut self,
        file_handle: Handle<File>,
        parent_node: Option<Handle<Node>>,
        exported_name: &str,
        target_path: &str,
    ) -> anyhow::Result<Handle<Node>> {
        let pop_id = self.graph.new_node_id(file_handle);
        let pop_sym = self.graph.add_symbol(exported_name);
        let pop_node = self
            .graph
            .add_pop_symbol_node(pop_id, pop_sym, true)
            .ok_or_else(|| anyhow::anyhow!("Failed to add reexport pop node {exported_name}"))?;

        if let Some(parent) = parent_node {
            self.graph.add_edge(parent, pop_node, 0);
        } else {
            self.graph.add_edge(StackGraph::root_node(), pop_node, 0);
        }

        // Split target path into segments, e.g. "http::Request" -> ["http", "Request"]
        let segments: Vec<&str> = target_path
            .trim_start_matches("::")
            .split("::")
            .filter(|s| !s.is_empty())
            .collect();

        if segments.is_empty() {
            return Ok(pop_node);
        }

        let mut prev = pop_node;
        // Pushing in reverse order ensures top of stack matches the first segment
        for &seg in segments.iter().rev() {
            let push_id = self.graph.new_node_id(file_handle);
            let push_sym = self.graph.add_symbol(seg);
            let push_node = self
                .graph
                .add_push_symbol_node(push_id, push_sym, false)
                .ok_or_else(|| anyhow::anyhow!("Failed to add push symbol node {seg}"))?;
            self.graph.add_edge(prev, push_node, 0);
            prev = push_node;
        }

        self.graph.add_edge(prev, StackGraph::root_node(), 0);

        let file_path = PathBuf::from(self.graph[file_handle].name());
        if !self
            .registered_reexports
            .iter()
            .any(|(p, _, e, t)| p == &file_path && e == exported_name && t == target_path)
        {
            self.registered_reexports.push((
                file_path,
                None,
                exported_name.to_string(),
                target_path.to_string(),
            ));
        }

        Ok(pop_node)
    }

    /// Add an import stitching alias connecting local reference `local_name` to `target_path`.
    pub fn add_import_stitching(
        &mut self,
        file_handle: Handle<File>,
        local_name: &str,
        target_path: &str,
    ) -> anyhow::Result<Handle<Node>> {
        let pop_id = self.graph.new_node_id(file_handle);
        let pop_sym = self.graph.add_symbol(local_name);
        let pop_node = self
            .graph
            .add_pop_symbol_node(pop_id, pop_sym, true)
            .ok_or_else(|| anyhow::anyhow!("Failed to add import pop node {local_name}"))?;

        self.graph.add_edge(StackGraph::root_node(), pop_node, 0);

        let segments: Vec<&str> = target_path
            .trim_start_matches("::")
            .split("::")
            .filter(|s| !s.is_empty())
            .collect();

        if segments.is_empty() {
            return Ok(pop_node);
        }

        let mut prev = pop_node;
        for &seg in segments.iter().rev() {
            let push_id = self.graph.new_node_id(file_handle);
            let push_sym = self.graph.add_symbol(seg);
            let push_node = self
                .graph
                .add_push_symbol_node(push_id, push_sym, false)
                .ok_or_else(|| anyhow::anyhow!("Failed to add push node {seg}"))?;
            self.graph.add_edge(prev, push_node, 0);
            prev = push_node;
        }

        self.graph.add_edge(prev, StackGraph::root_node(), 0);
        Ok(pop_node)
    }

    /// Precompute minimal partial paths for a file and register them in the stitching database.
    pub fn precompute_file_paths(
        &mut self,
        file: Handle<File>,
        cancellation: &dyn stack_graphs::CancellationFlag,
    ) -> anyhow::Result<usize> {
        if !self.precomputed.insert(file) {
            return Ok(0);
        }
        let mut added = 0;
        let config = StitcherConfig::default();
        let mut discovered = Vec::new();

        // Each file gets its own step budget. Sharing one counter across files exhausts it
        // after the first few and silently drops partial paths for the rest.
        let bounded = StepBoundedCancellationFlag::with_inner(self.max_steps, cancellation);
        let res = ForwardPartialPathStitcher::find_minimal_partial_path_set_in_file(
            &self.graph,
            &mut self.partials,
            file,
            config,
            &bounded,
            |_graph, _partials, path| {
                discovered.push(path.clone());
            },
        );

        if res.is_ok() {
            for path in discovered {
                self.database
                    .add_partial_path(&self.graph, &mut self.partials, path);
                added += 1;
            }
        } else {
            // Not done: leave the file eligible for a later attempt instead of marking it
            // precomputed with no paths.
            self.precomputed.remove(&file);
            tracing::debug!(
                "Partial path precompute hit the {}-step limit for {}",
                self.max_steps,
                self.graph[file].name()
            );
        }

        Ok(added)
    }

    /// Precompute partial paths for all indexed files.
    pub fn precompute_all_files(
        &mut self,
        cancellation: &dyn stack_graphs::CancellationFlag,
    ) -> anyhow::Result<usize> {
        self.refresh(cancellation);
        let mut total = 0;
        let file_handles: Vec<Handle<File>> = self.files.values().copied().collect();
        for file in file_handles {
            total += self.precompute_file_paths(file, cancellation)?;
        }
        Ok(total)
    }

    /// Convert a definition node into structured `ResolvedDefinition`.
    fn node_to_definition(&self, node: Handle<Node>) -> Option<ResolvedDefinition> {
        let node_id = self.graph[node].id();
        let file_handle = node_id.file()?;
        let file_path = PathBuf::from(self.graph[file_handle].name());

        let symbol_name = match &self.graph[node] {
            Node::PopSymbol(n) => self.graph[n.symbol].to_string(),
            Node::PopScopedSymbol(n) => self.graph[n.symbol].to_string(),
            Node::PushSymbol(n) => self.graph[n.symbol].to_string(),
            Node::PushScopedSymbol(n) => self.graph[n.symbol].to_string(),
            _ => return None,
        };

        let (start_line, start_col, end_line, end_col, syntax_type) =
            if let Some(src) = self.graph.source_info(node) {
                let st = src
                    .syntax_type
                    .into_option()
                    .map(|h| self.graph[h].to_string());
                (
                    src.span.start.line as u32 + 1,
                    src.span.start.column.utf8_offset as u32,
                    src.span.end.line as u32 + 1,
                    src.span.end.column.utf8_offset as u32,
                    st,
                )
            } else {
                (1, 0, 1, 0, None)
            };

        let is_external = self.dependency_files.contains(&file_handle) || file_path.is_absolute();
        let crate_name = self.file_to_crate.get(&file_handle).cloned();

        Some(ResolvedDefinition {
            symbol_name,
            file_path,
            syntax_type,
            start_line,
            start_col,
            end_line,
            end_col,
            is_external,
            crate_name,
        })
    }

    /// Resolve a specific reference AST node using partial path stitching.
    pub fn resolve_node(
        &mut self,
        ref_node: Handle<Node>,
        cancellation: &dyn stack_graphs::CancellationFlag,
    ) -> anyhow::Result<Vec<ResolvedDefinition>> {
        if !self.graph[ref_node].is_reference() {
            return Ok(Vec::new());
        }

        let mut initial_path = PartialPath::from_node(&self.graph, &mut self.partials, ref_node);
        initial_path.eliminate_precondition_stack_variables(&mut self.partials);

        let mut stitcher = ForwardPartialPathStitcher::from_partial_paths(
            &self.graph,
            &mut self.partials,
            std::iter::once(initial_path),
        );
        stitcher.set_similar_path_detection(true);
        stitcher.set_check_only_join_nodes(true);

        let mut candidates =
            DatabaseCandidates::new(&self.graph, &mut self.partials, &mut self.database);

        let mut steps = 0;
        let mut complete_paths = Vec::new();

        while !stitcher.is_complete() {
            cancellation.check("resolving stack graph reference")?;
            if steps >= self.max_steps {
                tracing::warn!(
                    "Path stitching reached max_steps ({}) boundary limit, terminating search",
                    self.max_steps
                );
                break;
            }
            steps += 1;

            for path in stitcher.previous_phase_partial_paths() {
                candidates.load_forward_candidates(path, cancellation)?;
            }
            stitcher.process_next_phase(&mut candidates, |_, _, _| true);

            for path in stitcher.previous_phase_partial_paths() {
                if path.is_complete(&self.graph) {
                    complete_paths.push(path.clone());
                }
            }
        }

        let mut definitions = Vec::new();
        for path in complete_paths {
            let def_node = path.end_node;
            if !self.graph[def_node].is_definition() {
                continue;
            }
            let Some(def) = self.node_to_definition(def_node) else {
                continue;
            };
            if !definitions.contains(&def) {
                definitions.push(def);
            }
        }

        Ok(definitions)
    }

    /// Resolve the symbol reference located at `(line, col)` in `rel_path`.
    pub fn resolve_at_location(
        &mut self,
        rel_path: &Path,
        line: u32,
        col: u32,
    ) -> anyhow::Result<Option<ResolvedDefinition>> {
        self.refresh(&stack_graphs::NoCancellation);
        let file_handle = match self.files.get(rel_path) {
            Some(h) => *h,
            None => return Ok(None),
        };

        // Find reference node matching line and col
        let ref_nodes: Vec<Handle<Node>> = self
            .graph
            .nodes_for_file(file_handle)
            .filter(|&n| self.graph[n].is_reference())
            .collect();

        // Convert 1-indexed line to 0-indexed for comparison with LSP span
        let target_line_0 = line.saturating_sub(1);

        let mut best_node = None;
        for node in ref_nodes {
            if let Some(src) = self.graph.source_info(node) {
                let start_line = src.span.start.line as u32;
                let end_line = src.span.end.line as u32;
                let start_col = src.span.start.column.utf8_offset as u32;
                let end_col = src.span.end.column.utf8_offset as u32;

                if target_line_0 >= start_line && target_line_0 <= end_line {
                    if target_line_0 == start_line && target_line_0 == end_line {
                        if col >= start_col && col <= end_col {
                            best_node = Some(node);
                            break;
                        }
                    } else {
                        best_node = Some(node);
                        break;
                    }
                }
            }
        }

        let node_to_resolve = match best_node {
            Some(n) => n,
            None => return Ok(None),
        };

        let step_cancellation = StepBoundedCancellationFlag::new(self.max_steps);
        let results = self.resolve_node(node_to_resolve, &step_cancellation)?;
        Ok(results.into_iter().next())
    }
}