python.rs

use tree_sitter::{Node, Tree};
use crate::ast_analysis::cfg::build_function_cfg;
use crate::ast_analysis::complexity::compute_all_metrics;
use crate::types::*;
use super::helpers::*;
use super::SymbolExtractor;

pub struct PythonExtractor;

impl SymbolExtractor for PythonExtractor {
    fn extract(&self, tree: &Tree, source: &[u8], file_path: &str) -> FileSymbols {
        let mut symbols = FileSymbols::new(file_path.to_string());
        walk_tree(&tree.root_node(), source, &mut symbols, match_python_node);
        walk_ast_nodes_with_config(&tree.root_node(), source, &mut symbols.ast_nodes, &PYTHON_AST_CONFIG);
        walk_tree(&tree.root_node(), source, &mut symbols, match_python_type_map);
        symbols
    }
}

fn match_python_node(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) {
    match node.kind() {
        "function_definition" => handle_function_def(node, source, symbols),
        "class_definition" => handle_class_def(node, source, symbols),
        "expression_statement" => handle_expr_stmt(node, source, symbols),
        "call" => handle_call(node, source, symbols),
        "import_statement" => handle_import_stmt(node, source, symbols),
        "import_from_statement" => handle_import_from_stmt(node, source, symbols),
        _ => {}
    }
}

// ── Per-node-kind handlers for walk_node_depth ───────────────────────────────

fn handle_function_def(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    let Some(name_node) = node.child_by_field_name("name") else { return };
    let name_text = node_text(&name_node, source);
    let mut decorators = Vec::new();
    if let Some(prev) = node.prev_sibling() {
        if prev.kind() == "decorator" {
            decorators.push(node_text(&prev, source).to_string());
        }
    }
    let parent_class = find_python_parent_class(node, source);
    let (full_name, kind) = match &parent_class {
        Some(cls) => (format!("{}.{}", cls, name_text), "method".to_string()),
        None => (name_text.to_string(), "function".to_string()),
    };
    let children = extract_python_parameters(node, source, parent_class.is_some());
    symbols.definitions.push(Definition {
        name: full_name,
        kind,
        line: start_line(node),
        end_line: Some(end_line(node)),
        decorators: if decorators.is_empty() { None } else { Some(decorators) },
        complexity: compute_all_metrics(node, source, "python"),
        cfg: build_function_cfg(node, "python", source),
        children: opt_children(children),
    });
}

fn handle_class_def(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    let Some(name_node) = node.child_by_field_name("name") else { return };
    let class_name = node_text(&name_node, source).to_string();
    let children = extract_python_class_properties(node, source);
    symbols.definitions.push(Definition {
        name: class_name.clone(),
        kind: "class".to_string(),
        line: start_line(node),
        end_line: Some(end_line(node)),
        decorators: None,
        complexity: None,
        cfg: None,
        children: opt_children(children),
    });
    let superclasses = node
        .child_by_field_name("superclasses")
        .or_else(|| find_child(node, "argument_list"));
    if let Some(superclasses) = superclasses {
        for i in 0..superclasses.child_count() {
            if let Some(child) = superclasses.child(i) {
                if child.kind() == "identifier" {
                    symbols.classes.push(ClassRelation {
                        name: class_name.clone(),
                        extends: Some(node_text(&child, source).to_string()),
                        implements: None,
                        line: start_line(node),
                    });
                }
            }
        }
    }
}

fn handle_expr_stmt(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    if !is_module_level(node) { return; }
    let Some(expr) = node.child(0) else { return };
    if expr.kind() != "assignment" { return; }
    let Some(left) = expr.child_by_field_name("left") else { return };
    if left.kind() != "identifier" { return; }
    let name = node_text(&left, source);
    if !is_upper_snake_case(name) { return; }
    symbols.definitions.push(Definition {
        name: name.to_string(),
        kind: "constant".to_string(),
        line: start_line(node),
        end_line: Some(end_line(node)),
        decorators: None,
        complexity: None,
        cfg: None,
        children: None,
    });
}

fn handle_call(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    let Some(fn_node) = node.child_by_field_name("function") else { return };
    let (call_name, receiver) = match fn_node.kind() {
        "identifier" => (Some(node_text(&fn_node, source).to_string()), None),
        "attribute" => {
            let name = named_child_text(&fn_node, "attribute", source)
                .map(|s| s.to_string());
            let recv = named_child_text(&fn_node, "object", source)
                .map(|s| s.to_string());
            (name, recv)
        }
        _ => (None, None),
    };
    if let Some(name) = call_name {
        symbols.calls.push(Call {
            name,
            line: start_line(node),
            dynamic: None,
            receiver,
        });
    }
}

fn handle_import_stmt(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    let mut names = Vec::new();
    for i in 0..node.child_count() {
        let Some(child) = node.child(i) else { continue };
        if child.kind() != "dotted_name" && child.kind() != "aliased_import" { continue; }
        let name = if child.kind() == "aliased_import" {
            child
                .child_by_field_name("alias")
                .or_else(|| child.child_by_field_name("name"))
                .map(|n| node_text(&n, source).to_string())
        } else {
            Some(node_text(&child, source).to_string())
        };
        if let Some(name) = name {
            names.push(name);
        }
    }
    if !names.is_empty() {
        let mut imp = Import::new(names[0].clone(), names, start_line(node));
        imp.python_import = Some(true);
        symbols.imports.push(imp);
    }
}

fn handle_import_from_stmt(node: &Node, source: &[u8], symbols: &mut FileSymbols) {
    let mut source_str = String::new();
    let mut names = Vec::new();
    for i in 0..node.child_count() {
        let Some(child) = node.child(i) else { continue };
        match child.kind() {
            "dotted_name" | "relative_import" => {
                if source_str.is_empty() {
                    source_str = node_text(&child, source).to_string();
                } else {
                    names.push(node_text(&child, source).to_string());
                }
            }
            "aliased_import" => {
                let n = child
                    .child_by_field_name("name")
                    .or_else(|| child.child(0));
                if let Some(n) = n {
                    names.push(node_text(&n, source).to_string());
                }
            }
            "wildcard_import" => {
                names.push("*".to_string());
            }
            _ => {}
        }
    }
    if !source_str.is_empty() {
        let mut imp = Import::new(source_str, names, start_line(node));
        imp.python_import = Some(true);
        symbols.imports.push(imp);
    }
}

// ── Extended kinds helpers ──────────────────────────────────────────────────

fn extract_python_parameters(node: &Node, source: &[u8], is_method: bool) -> Vec<Definition> {
    let mut params = Vec::new();
    let params_node = node.child_by_field_name("parameters");
    if let Some(params_node) = params_node {
        for i in 0..params_node.child_count() {
            if let Some(child) = params_node.child(i) {
                let name = match child.kind() {
                    "identifier" => {
                        let text = node_text(&child, source);
                        Some(text.to_string())
                    }
                    "default_parameter" | "typed_default_parameter" => {
                        named_child_text(&child, "name", source)
                            .map(|s| s.to_string())
                    }
                    "typed_parameter" => {
                        // typed_parameter: first child is the identifier
                        child.child(0)
                            .filter(|c| c.kind() == "identifier")
                            .map(|c| node_text(&c, source).to_string())
                    }
                    "list_splat_pattern" | "dictionary_splat_pattern" => {
                        // *args, **kwargs
                        child.child(0)
                            .filter(|c| c.kind() == "identifier")
                            .map(|c| node_text(&c, source).to_string())
                    }
                    _ => None,
                };
                if let Some(name) = name {
                    // Skip self/cls for methods
                    if is_method && (name == "self" || name == "cls") {
                        continue;
                    }
                    params.push(child_def(name, "parameter", start_line(&child)));
                }
            }
        }
    }
    params
}

fn extract_python_class_properties(class_node: &Node, source: &[u8]) -> Vec<Definition> {
    let mut props = Vec::new();
    let body = class_node.child_by_field_name("body");
    if let Some(body) = body {
        // Look for __init__ method and scan for self.x = ... assignments
        for i in 0..body.child_count() {
            if let Some(child) = body.child(i) {
                if child.kind() == "function_definition" {
                    if let Some(name_node) = child.child_by_field_name("name") {
                        if node_text(&name_node, source) == "__init__" {
                            collect_self_assignments(&child, source, &mut props);
                        }
                    }
                }
            }
        }
    }
    props
}

fn collect_self_assignments(node: &Node, source: &[u8], props: &mut Vec<Definition>) {
    for i in 0..node.child_count() {
        let Some(child) = node.child(i) else { continue };
        if child.kind() == "expression_statement" {
            try_extract_self_assignment(&child, source, props);
        }
        // Recurse into blocks (if/for/etc inside __init__)
        if child.kind() == "block" || child.kind() == "if_statement"
            || child.kind() == "for_statement" || child.kind() == "while_statement"
        {
            collect_self_assignments(&child, source, props);
        }
    }
}

fn try_extract_self_assignment(stmt: &Node, source: &[u8], props: &mut Vec<Definition>) {
    let Some(expr) = stmt.child(0) else { return };
    if expr.kind() != "assignment" { return; }
    let Some(left) = expr.child_by_field_name("left") else { return };
    if left.kind() != "attribute" { return; }
    let Some(obj) = left.child_by_field_name("object") else { return };
    if node_text(&obj, source) != "self" { return; }
    let Some(attr) = left.child_by_field_name("attribute") else { return };
    let name = node_text(&attr, source);
    if !props.iter().any(|p| p.name == name) {
        props.push(child_def(name.to_string(), "property", start_line(stmt)));
    }
}

fn is_module_level(node: &Node) -> bool {
    if let Some(parent) = node.parent() {
        return parent.kind() == "module";
    }
    false
}

fn is_upper_snake_case(s: &str) -> bool {
    !s.is_empty()
        && s.chars().all(|c| c.is_ascii_uppercase() || c == '_' || c.is_ascii_digit())
        && s.chars().next().map(|c| c.is_ascii_uppercase()).unwrap_or(false)
}

// ── Existing helpers ────────────────────────────────────────────────────────

const PYTHON_CLASS_KINDS: &[&str] = &["class_definition"];

fn find_python_parent_class(node: &Node, source: &[u8]) -> Option<String> {
    find_enclosing_type_name(node, PYTHON_CLASS_KINDS, source)
}

fn extract_python_type_name<'a>(type_node: &Node<'a>, source: &'a [u8]) -> Option<&'a str> {
    match type_node.kind() {
        "identifier" | "attribute" => Some(node_text(type_node, source)),
        "subscript" => {
            // List[int] → List
            named_child_text(type_node, "value", source)
        }
        _ => None,
    }
}

/// Python builtins / stdlib classes that start with an uppercase letter and would
/// false-positive on the constructor-call heuristic.  Mirrors `BUILTIN_GLOBALS_PY`
/// in `src/extractors/python.ts`.
fn is_python_builtin(name: &str) -> bool {
    matches!(
        name,
        "Exception"
            | "BaseException"
            | "ValueError"
            | "TypeError"
            | "KeyError"
            | "IndexError"
            | "AttributeError"
            | "RuntimeError"
            | "OSError"
            | "IOError"
            | "FileNotFoundError"
            | "PermissionError"
            | "NotImplementedError"
            | "StopIteration"
            | "GeneratorExit"
            | "SystemExit"
            | "KeyboardInterrupt"
            | "ArithmeticError"
            | "LookupError"
            | "UnicodeError"
            | "UnicodeDecodeError"
            | "UnicodeEncodeError"
            | "ImportError"
            | "ModuleNotFoundError"
            | "ConnectionError"
            | "TimeoutError"
            | "OverflowError"
            | "ZeroDivisionError"
            | "NameError"
            | "SyntaxError"
            | "RecursionError"
            | "MemoryError"
            | "Path"
            | "PurePath"
            | "OrderedDict"
            | "Counter"
            | "Decimal"
            | "Fraction"
    )
}

fn match_python_type_map(node: &Node, source: &[u8], symbols: &mut FileSymbols, _depth: usize) {
    match node.kind() {
        "typed_parameter" => {
            // first child is identifier, type field is the type
            if let Some(name_node) = node.child(0) {
                if name_node.kind() == "identifier" {
                    let name = node_text(&name_node, source);
                    if name != "self" && name != "cls" {
                        if let Some(type_node) = node.child_by_field_name("type") {
                            if let Some(type_name) =
                                extract_python_type_name(&type_node, source)
                            {
                                symbols.type_map.push(TypeMapEntry {
                                    name: name.to_string(),
                                    type_name: type_name.to_string(),
                                    confidence: 0.9,
                                });
                            }
                        }
                    }
                }
            }
        }
        "typed_default_parameter" => {
            if let Some(name_node) = node.child_by_field_name("name") {
                if name_node.kind() == "identifier" {
                    if let Some(type_node) = node.child_by_field_name("type") {
                        if let Some(type_name) =
                            extract_python_type_name(&type_node, source)
                        {
                            symbols.type_map.push(TypeMapEntry {
                                name: node_text(&name_node, source).to_string(),
                                type_name: type_name.to_string(),
                                confidence: 0.9,
                            });
                        }
                    }
                }
            }
        }
        // `order = Order(...)` → seed order : Order at conf 1.0.
        // `obj = module.Class(...)` → seed obj : module at conf 0.7 (factory pattern).
        // Mirrors `handlePyAssignmentType` in `src/extractors/python.ts`.
        "assignment" => {
            infer_py_assignment_type(node, source, &mut symbols.type_map);
        }
        _ => {}
    }
}

/// Seed typeMap from plain Python assignments where the RHS is a constructor or factory call.
fn infer_py_assignment_type(node: &Node, source: &[u8], type_map: &mut Vec<TypeMapEntry>) {
    let Some(left) = node.child_by_field_name("left") else { return };
    let Some(right) = node.child_by_field_name("right") else { return };
    if left.kind() != "identifier" || right.kind() != "call" { return; }
    let var_name = node_text(&left, source).to_string();
    let Some(fn_node) = right.child_by_field_name("function") else { return };
    match fn_node.kind() {
        "identifier" => {
            // `order = Order(...)` — uppercase first char → constructor, conf 1.0.
            let name = node_text(&fn_node, source);
            if name.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) {
                type_map.push(TypeMapEntry {
                    name: var_name,
                    type_name: name.to_string(),
                    confidence: 1.0,
                });
            }
        }
        "attribute" => {
            // `obj = Module.Class(...)` — uppercase object name, not a builtin → conf 0.7.
            if let Some(obj_node) = fn_node.child_by_field_name("object") {
                if obj_node.kind() == "identifier" {
                    let obj_name = node_text(&obj_node, source);
                    if obj_name.chars().next().map(|c| c.is_uppercase()).unwrap_or(false)
                        && !is_python_builtin(obj_name)
                    {
                        type_map.push(TypeMapEntry {
                            name: var_name,
                            type_name: obj_name.to_string(),
                            confidence: 0.7,
                        });
                    }
                }
            }
        }
        _ => {}
    }
}

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

    fn parse_py(code: &str) -> FileSymbols {
        let mut parser = Parser::new();
        parser
            .set_language(&tree_sitter_python::LANGUAGE.into())
            .unwrap();
        let tree = parser.parse(code.as_bytes(), None).unwrap();
        PythonExtractor.extract(&tree, code.as_bytes(), "test.py")
    }

    #[test]
    fn finds_function() {
        let s = parse_py("def greet(name):\n    return name\n");
        assert_eq!(s.definitions.len(), 1);
        assert_eq!(s.definitions[0].name, "greet");
        assert_eq!(s.definitions[0].kind, "function");
    }

    #[test]
    fn finds_class_and_method() {
        let s = parse_py("class Foo:\n    def bar(self):\n        pass\n");
        let names: Vec<&str> = s.definitions.iter().map(|d| d.name.as_str()).collect();
        assert!(names.contains(&"Foo"));
        assert!(names.contains(&"Foo.bar"));
    }

    #[test]
    fn finds_imports() {
        let s = parse_py("from os.path import join, exists\n");
        assert_eq!(s.imports.len(), 1);
        assert_eq!(s.imports[0].source, "os.path");
        assert!(s.imports[0].names.contains(&"join".to_string()));
    }

    #[test]
    fn finds_calls() {
        let s = parse_py("print('hello')\nos.path.join('a', 'b')\n");
        let call_names: Vec<&str> = s.calls.iter().map(|c| c.name.as_str()).collect();
        assert!(call_names.contains(&"print"));
        assert!(call_names.contains(&"join"));
    }

    #[test]
    fn finds_inheritance() {
        let s = parse_py("class Dog(Animal):\n    pass\n");
        assert_eq!(s.classes.len(), 1);
        assert_eq!(s.classes[0].name, "Dog");
        assert_eq!(s.classes[0].extends, Some("Animal".to_string()));
    }

    // ── Extended kinds tests ────────────────────────────────────────────────

    #[test]
    fn extracts_function_parameters() {
        let s = parse_py("def greet(name, age=30):\n  pass");
        let greet = s.definitions.iter().find(|d| d.name == "greet").unwrap();
        let children = greet.children.as_ref().unwrap();
        assert_eq!(children.len(), 2);
        assert_eq!(children[0].name, "name");
        assert_eq!(children[0].kind, "parameter");
        assert_eq!(children[1].name, "age");
    }

    #[test]
    fn extracts_method_parameters_skips_self() {
        let s = parse_py("class Foo:\n    def bar(self, x, y):\n        pass\n");
        let bar = s.definitions.iter().find(|d| d.name == "Foo.bar").unwrap();
        let children = bar.children.as_ref().unwrap();
        assert_eq!(children.len(), 2);
        assert_eq!(children[0].name, "x");
        assert_eq!(children[1].name, "y");
    }

    #[test]
    fn extracts_class_properties_from_init() {
        let s = parse_py("class User:\n  def __init__(self, x, y):\n    self.x = x\n    self.y = y\n");
        let user = s.definitions.iter().find(|d| d.name == "User").unwrap();
        let children = user.children.as_ref().unwrap();
        let names: Vec<&str> = children.iter().map(|c| c.name.as_str()).collect();
        assert!(names.contains(&"x"));
        assert!(names.contains(&"y"));
        assert!(children.iter().all(|c| c.kind == "property"));
    }

    #[test]
    fn extracts_module_level_constant() {
        let s = parse_py("MAX_RETRIES = 3");
        let c = s.definitions.iter().find(|d| d.name == "MAX_RETRIES").unwrap();
        assert_eq!(c.kind, "constant");
    }

    // ── Assignment typeMap tests ─────────────────────────────────────────────

    #[test]
    fn infers_constructor_call_uppercase() {
        // order = Order("o1", 100.0) → order : Order at conf 1.0
        let s = parse_py("def run():\n    order = Order(\"o1\", 100.0)\n    order.validate()\n");
        let entry = s.type_map.iter().find(|e| e.name == "order");
        assert!(entry.is_some(), "expected order in type_map");
        let entry = entry.unwrap();
        assert_eq!(entry.type_name, "Order");
        assert!((entry.confidence - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn infers_module_factory_call() {
        // svc = Models.UserService(db) → svc : Models at conf 0.7
        // The object name must be uppercase to match the JS heuristic.
        let s = parse_py("def run():\n    svc = Models.UserService(db)\n    svc.create()\n");
        let entry = s.type_map.iter().find(|e| e.name == "svc");
        assert!(entry.is_some(), "expected svc in type_map for Module.Class(...)");
        let entry = entry.unwrap();
        assert_eq!(entry.type_name, "Models");
        assert!((entry.confidence - 0.7).abs() < f64::EPSILON);
    }

    #[test]
    fn does_not_infer_lowercase_module_factory() {
        // svc = models.UserService(db) — lowercase module name → no typeMap entry (matches JS)
        let s = parse_py("def run():\n    svc = models.UserService(db)\n    svc.create()\n");
        assert!(
            s.type_map.iter().all(|e| e.name != "svc"),
            "should not seed typeMap for lowercase module prefix"
        );
    }

    #[test]
    fn does_not_infer_lowercase_constructor() {
        // obj = create_thing() — lowercase, should not seed typeMap
        let s = parse_py("def run():\n    obj = create_thing()\n    obj.work()\n");
        assert!(
            s.type_map.iter().all(|e| e.name != "obj"),
            "should not seed typeMap for lowercase function call"
        );
    }

    #[test]
    fn does_not_infer_builtin_exception() {
        // err = ValueError("msg") — builtin exception, should not seed typeMap
        let s = parse_py("def run():\n    err = ValueError(\"msg\")\n");
        // Note: ValueError is uppercase so it WOULD match the heuristic — but it's a builtin.
        // The JS extractor does NOT exclude builtins from conf-1.0 uppercase constructor
        // matching (only from the attribute/factory path). We match that behaviour here.
        // This test documents the current behaviour rather than asserting exclusion.
        let entry = s.type_map.iter().find(|e| e.name == "err");
        // Builtins ARE seeded at conf 1.0 by the identifier branch (same as JS).
        // Only the attribute/factory branch (Module.Class) checks is_python_builtin.
        if let Some(e) = entry {
            assert_eq!(e.type_name, "ValueError");
        }
    }
}