scryer-engine 0.3.0

use std::path::Path;
use tree_sitter::{Node, Parser};

use crate::payload::{ParsedFilePayload, RawScope, RawSymbol};

/// Tree-sitter AST parser and entity extractor for TypeScript and TSX source code.
pub struct TypeScriptAstParser;

impl TypeScriptAstParser {
    /// Parse TypeScript / TSX source code and extract all declarations, signatures, docstrings, and lexical scopes.
    pub fn parse(
        relative_path: &Path,
        content: &[u8],
        content_hash: &str,
    ) -> anyhow::Result<ParsedFilePayload> {
        let source_str = std::str::from_utf8(content)
            .map_err(|e| anyhow::anyhow!("Invalid UTF-8 in {}: {e}", relative_path.display()))?;

        let is_tsx = relative_path
            .extension()
            .and_then(|ext| ext.to_str())
            .map(|ext| ext.eq_ignore_ascii_case("tsx") || ext.eq_ignore_ascii_case("jsx"))
            .unwrap_or(false);

        let mut parser = Parser::new();
        let language = if is_tsx {
            tree_sitter_typescript::LANGUAGE_TSX
        } else {
            tree_sitter_typescript::LANGUAGE_TYPESCRIPT
        };

        parser
            .set_language(&language.into())
            .map_err(|e| anyhow::anyhow!("Failed to set TypeScript tree-sitter language: {e}"))?;

        let tree = parser.parse(content, None).ok_or_else(|| {
            anyhow::anyhow!("Tree-sitter failed to parse {}", relative_path.display())
        })?;

        let base_module = compute_base_module(relative_path);
        let mut extractor = TypeScriptAstExtractor::new(source_str, base_module);

        // Add root module scope
        let root_node = tree.root_node();
        let root_scope = extractor.add_scope("module", root_node, None);

        extractor.traverse(root_node, Some(root_scope), &[], false);

        let line_count = source_str.lines().count() as u32;

        Ok(ParsedFilePayload {
            relative_path: relative_path.to_path_buf(),
            content_hash: content_hash.to_string(),
            language: if is_tsx {
                "tsx".to_string()
            } else {
                "typescript".to_string()
            },
            line_count: if line_count == 0 { 1 } else { line_count },
            byte_size: content.len(),
            scopes: extractor.scopes,
            symbols: extractor.symbols,
            references: Vec::new(),
            edges: Vec::new(),
            syntax_errors: tree.root_node().has_error(),
        })
    }
}

/// Compute the base TypeScript module path from the relative file path.
/// E.g.: `src/components/Button.tsx` -> "src/components/Button"
pub(crate) fn compute_base_module(rel_path: &Path) -> String {
    let mut components = Vec::new();

    for part in rel_path.iter() {
        let part_str = part.to_string_lossy();
        if let Some(stem) = part_str
            .strip_suffix(".ts")
            .or_else(|| part_str.strip_suffix(".tsx"))
            .or_else(|| part_str.strip_suffix(".js"))
            .or_else(|| part_str.strip_suffix(".jsx"))
        {
            if stem != "index" {
                components.push(stem.to_string());
            }
        } else {
            components.push(part_str.to_string());
        }
    }

    if components.is_empty() {
        "index".to_string()
    } else {
        components.join("/")
    }
}

/// Extract the members of the TypeScript class, interface or enum `type_name` declared in
/// `source` that are not symbols of their own: fields (`field:`), the call, construct, index and
/// method signatures of an interface or an abstract class (`field:` / `method:`) and the members
/// of an enum (`variant:`). A class's concrete methods are symbols and come from the index.
///
/// A field is shown without its initializer. Returns an empty list when the type isn't
/// declared in `source`.
pub fn extract_type_members(source: &str, type_name: &str, tsx: bool) -> Vec<String> {
    let mut parser = Parser::new();
    let language = if tsx {
        tree_sitter_typescript::LANGUAGE_TSX
    } else {
        tree_sitter_typescript::LANGUAGE_TYPESCRIPT
    };
    if parser.set_language(&language.into()).is_err() {
        return Vec::new();
    }
    let Some(tree) = parser.parse(source, None) else {
        return Vec::new();
    };

    let text = |node: Node| &source[node.start_byte()..node.end_byte()];
    let one_line = |raw: &str| {
        let line = raw.split_whitespace().collect::<Vec<_>>().join(" ");
        let line = line.trim_end_matches([';', ',']).trim_end();
        match line.char_indices().nth(300) {
            Some((at, _)) => format!("{}…", &line[..at]),
            None => line.to_string(),
        }
    };

    let mut stack = vec![tree.root_node()];
    while let Some(node) = stack.pop() {
        let is_type = matches!(
            node.kind(),
            "class_declaration"
                | "abstract_class_declaration"
                | "interface_declaration"
                | "enum_declaration"
        );
        let named = node
            .child_by_field_name("name")
            .is_some_and(|n| text(n) == type_name);
        if is_type && named {
            let Some(body) = node.child_by_field_name("body") else {
                return Vec::new();
            };
            let mut members = Vec::new();
            let mut cursor = body.walk();
            for member in body.named_children(&mut cursor) {
                let entry = match member.kind() {
                    "public_field_definition" | "property_signature" => {
                        // Without the initializer: `cache: Map<string, User> = new Map()`.
                        let end = member
                            .child_by_field_name("value")
                            .map_or(member.end_byte(), |v| v.start_byte());
                        let declared = one_line(&source[member.start_byte()..end]);
                        format!("field: {}", declared.trim_end_matches('=').trim_end())
                    }
                    "index_signature" => format!("field: {}", one_line(text(member))),
                    "method_signature"
                    | "abstract_method_signature"
                    | "call_signature"
                    | "construct_signature" => format!("method: {}", one_line(text(member))),
                    "enum_assignment" | "property_identifier" => {
                        format!("variant: {}", one_line(text(member)))
                    }
                    _ => continue,
                };
                members.push(entry);
            }
            return members;
        }
        let mut cursor = node.walk();
        let children: Vec<Node> = node.children(&mut cursor).collect();
        stack.extend(children.into_iter().rev());
    }
    Vec::new()
}

struct TypeScriptAstExtractor<'a> {
    source: &'a str,
    base_module: String,
    scopes: Vec<RawScope>,
    symbols: Vec<RawSymbol>,
}

impl<'a> TypeScriptAstExtractor<'a> {
    fn new(source: &'a str, base_module: String) -> Self {
        Self {
            source,
            base_module,
            scopes: Vec::new(),
            symbols: Vec::new(),
        }
    }

    fn add_scope(&mut self, kind: &str, node: Node, parent_scope: Option<usize>) -> usize {
        let local_id = self.scopes.len();
        self.scopes.push(RawScope {
            local_id,
            parent_local_id: parent_scope,
            scope_kind: kind.to_string(),
            start_byte: node.start_byte(),
            end_byte: node.end_byte(),
            start_line: (node.start_position().row + 1) as u32,
            end_line: (node.end_position().row + 1) as u32,
        });
        local_id
    }

    fn get_node_text(&self, node: Node) -> &'a str {
        &self.source[node.start_byte()..node.end_byte()]
    }

    /// Extract preceding JSDoc comment.
    fn extract_jsdoc(&self, node: Node) -> Option<String> {
        let start_byte = if let Some(parent) = node.parent() {
            if parent.kind() == "export_statement" {
                parent.start_byte()
            } else {
                node.start_byte()
            }
        } else {
            node.start_byte()
        };
        let preceding = &self.source[..start_byte];
        let lines: Vec<&str> = preceding.lines().collect();

        let mut doc_lines = Vec::new();
        let mut in_comment = false;

        for line in lines.iter().rev() {
            let trimmed = line.trim();
            if trimmed.is_empty() {
                if !doc_lines.is_empty() {
                    break;
                }
                continue;
            }

            if trimmed.ends_with("*/") {
                in_comment = true;
            }

            if in_comment {
                let clean = trimmed
                    .trim_start_matches("/**")
                    .trim_start_matches("/*")
                    .trim_end_matches("*/")
                    .trim_start_matches('*')
                    .trim();
                doc_lines.push(clean.to_string());
                if trimmed.starts_with("/**") || trimmed.starts_with("/*") {
                    break;
                }
            } else if trimmed.starts_with("//") {
                let clean = trimmed.trim_start_matches("//").trim();
                doc_lines.push(clean.to_string());
            } else {
                break;
            }
        }

        doc_lines.reverse();
        if doc_lines.is_empty() {
            None
        } else {
            Some(doc_lines.join("\n").trim().to_string())
        }
    }

    fn extract_signature(&self, node: Node) -> String {
        let text = self.get_node_text(node);
        let sig_end = text
            .find('{')
            .or_else(|| text.find(';'))
            .unwrap_or(text.len());
        let sig = &text[..sig_end];
        sig.split_whitespace().collect::<Vec<&str>>().join(" ")
    }

    fn make_qualified_name(&self, name: &str, qualifiers: &[String]) -> String {
        let mut parts = vec![self.base_module.clone()];
        parts.extend_from_slice(qualifiers);
        parts.push(name.to_string());
        parts.join(".")
    }

    fn traverse(
        &mut self,
        node: Node,
        current_scope: Option<usize>,
        qualifiers: &[String],
        is_exported: bool,
    ) {
        match node.kind() {
            "export_statement" => {
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    if child.kind() != "export" {
                        self.traverse(child, current_scope, qualifiers, true);
                    }
                }
            }

            "function_declaration" | "generator_function_declaration" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("anonymous_fn");

                let visibility = if is_exported { "public" } else { "private" };
                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                let scope_id = self.add_scope("function", node, current_scope);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name: qualified_name.clone(),
                    kind: "fn".to_string(),
                    visibility: visibility.to_string(),
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });

                if let Some(body) = node.child_by_field_name("body") {
                    let mut child_qualifiers = qualifiers.to_vec();
                    child_qualifiers.push(name.to_string());

                    let mut cursor = body.walk();
                    for child in body.children(&mut cursor) {
                        self.traverse(child, Some(scope_id), &child_qualifiers, false);
                    }
                }
            }

            "class_declaration" | "abstract_class_declaration" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("AnonymousClass");

                let visibility = if is_exported { "public" } else { "private" };
                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                let scope_id = self.add_scope("class", node, current_scope);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name: qualified_name.clone(),
                    kind: "class".to_string(),
                    visibility: visibility.to_string(),
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });

                if let Some(body) = node.child_by_field_name("body") {
                    let mut child_qualifiers = qualifiers.to_vec();
                    child_qualifiers.push(name.to_string());

                    let mut cursor = body.walk();
                    for child in body.children(&mut cursor) {
                        self.traverse(child, Some(scope_id), &child_qualifiers, false);
                    }
                }
            }

            "interface_declaration" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("AnonymousInterface");

                let visibility = if is_exported { "public" } else { "private" };
                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                let scope_id = self.add_scope("interface", node, current_scope);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name,
                    kind: "interface".to_string(),
                    visibility: visibility.to_string(),
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });

                if let Some(body) = node.child_by_field_name("body") {
                    let mut child_qualifiers = qualifiers.to_vec();
                    child_qualifiers.push(name.to_string());

                    let mut cursor = body.walk();
                    for child in body.children(&mut cursor) {
                        self.traverse(child, Some(scope_id), &child_qualifiers, false);
                    }
                }
            }

            "type_alias_declaration" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("AnonymousType");

                let visibility = if is_exported { "public" } else { "private" };
                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name,
                    kind: "type".to_string(),
                    visibility: visibility.to_string(),
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });
            }

            "enum_declaration" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("AnonymousEnum");

                let visibility = if is_exported { "public" } else { "private" };
                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                let scope_id = self.add_scope("enum", node, current_scope);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name,
                    kind: "enum".to_string(),
                    visibility: visibility.to_string(),
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });

                if let Some(body) = node.child_by_field_name("body") {
                    let mut cursor = body.walk();
                    for child in body.children(&mut cursor) {
                        self.traverse(child, Some(scope_id), qualifiers, false);
                    }
                }
            }

            "method_definition" => {
                let name = node
                    .child_by_field_name("name")
                    .map(|n| self.get_node_text(n))
                    .unwrap_or("anonymous_method");

                let mut visibility = "public".to_string();
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    if child.kind() == "accessibility_modifier" {
                        visibility = self.get_node_text(child).trim().to_string();
                    }
                }

                let signature = self.extract_signature(node);
                let docstring = self.extract_jsdoc(node);
                let qualified_name = self.make_qualified_name(name, qualifiers);

                let scope_id = self.add_scope("function", node, current_scope);

                self.symbols.push(RawSymbol {
                    scope_local_id: current_scope,
                    name: name.to_string(),
                    qualified_name: qualified_name.clone(),
                    kind: "method".to_string(),
                    visibility,
                    signature,
                    docstring,
                    start_byte: node.start_byte(),
                    end_byte: node.end_byte(),
                    start_line: (node.start_position().row + 1) as u32,
                    end_line: (node.end_position().row + 1) as u32,
                });

                if let Some(body) = node.child_by_field_name("body") {
                    let mut child_qualifiers = qualifiers.to_vec();
                    child_qualifiers.push(name.to_string());

                    let mut cursor = body.walk();
                    for child in body.children(&mut cursor) {
                        self.traverse(child, Some(scope_id), &child_qualifiers, false);
                    }
                }
            }

            "lexical_declaration" | "variable_declaration" => {
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    if child.kind() == "variable_declarator"
                        && let Some(name_node) = child.child_by_field_name("name")
                    {
                        let name = self.get_node_text(name_node);
                        let value_node = child.child_by_field_name("value");
                        let is_arrow_fn = value_node
                            .map(|v| v.kind() == "arrow_function")
                            .unwrap_or(false);

                        // What a function binds to a name is a local, not a declaration (it
                        // would bury the functions in an outline and take the blame for the
                        // calls it is assigned from). A nested arrow function stays a function.
                        let in_function = current_scope
                            .is_some_and(|id| self.scopes[id].scope_kind == "function");
                        if in_function && !is_arrow_fn {
                            continue;
                        }

                        let kind = if is_arrow_fn { "fn" } else { "const" };
                        let visibility = if is_exported { "public" } else { "private" };
                        let signature = self.extract_signature(child);
                        let docstring = self.extract_jsdoc(node);
                        let qualified_name = self.make_qualified_name(name, qualifiers);

                        let scope_id = if is_arrow_fn {
                            value_node.map(|v| self.add_scope("function", v, current_scope))
                        } else {
                            None
                        };

                        self.symbols.push(RawSymbol {
                            scope_local_id: current_scope,
                            name: name.to_string(),
                            qualified_name: qualified_name.clone(),
                            kind: kind.to_string(),
                            visibility: visibility.to_string(),
                            signature,
                            docstring,
                            start_byte: child.start_byte(),
                            end_byte: child.end_byte(),
                            start_line: (child.start_position().row + 1) as u32,
                            end_line: (child.end_position().row + 1) as u32,
                        });

                        if is_arrow_fn
                            && let Some(body) =
                                value_node.and_then(|v| v.child_by_field_name("body"))
                        {
                            let mut child_qualifiers = qualifiers.to_vec();
                            child_qualifiers.push(name.to_string());

                            let mut b_cursor = body.walk();
                            for b_child in body.children(&mut b_cursor) {
                                self.traverse(b_child, scope_id, &child_qualifiers, false);
                            }
                        }
                    }
                }
            }

            _ => {
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    self.traverse(child, current_scope, qualifiers, is_exported);
                }
            }
        }
    }
}