car-ast 0.14.0

use super::{extract_doc_comment, extract_signature, field_text, node_text};
use crate::types::*;

pub fn extract(tree: &tree_sitter::Tree, source: &[u8]) -> (Vec<Symbol>, Vec<Import>) {
    let root = tree.root_node();
    let mut symbols = Vec::new();
    let mut imports = Vec::new();
    let mut cursor = root.walk();

    for child in root.children(&mut cursor) {
        extract_top_level(&child, source, &mut symbols, &mut imports);
    }

    (symbols, imports)
}

fn extract_top_level(
    child: &tree_sitter::Node,
    source: &[u8],
    symbols: &mut Vec<Symbol>,
    imports: &mut Vec<Import>,
) {
    match child.kind() {
        "function_definition" => {
            if let Some(sym) = extract_function(child, source, None) {
                symbols.push(sym);
            }
        }
        "class_definition" => {
            if let Some(sym) = extract_class(child, source) {
                symbols.push(sym);
            }
        }
        "decorated_definition" => {
            // @decorator \n class/function — unwrap to the inner definition
            let mut inner_cursor = child.walk();
            for inner in child.children(&mut inner_cursor) {
                match inner.kind() {
                    "class_definition" => {
                        if let Some(sym) = extract_class(&inner, source) {
                            symbols.push(sym);
                        }
                    }
                    "function_definition" => {
                        if let Some(sym) = extract_function(&inner, source, None) {
                            symbols.push(sym);
                        }
                    }
                    _ => {}
                }
            }
        }
        "import_statement" | "import_from_statement" => {
            imports.push(Import {
                path: node_text(child, source).to_string(),
                alias: None,
                span: Span::from_node(child),
            });
        }
        "expression_statement" => {
            let mut inner_cursor = child.walk();
            for inner in child.children(&mut inner_cursor) {
                if inner.kind() == "assignment" {
                    if let Some(left) = inner.child_by_field_name("left") {
                        let name = node_text(&left, source);
                        if name.chars().all(|c| c.is_uppercase() || c == '_') && !name.is_empty() {
                            symbols.push(Symbol {
                                name: name.to_string(),
                                kind: SymbolKind::Const,
                                span: Span::from_node(&inner),
                                signature: node_text(&inner, source).to_string(),
                                doc_comment: None,
                                parent: None,
                                children: Vec::new(),
                            });
                        }
                    }
                }
            }
        }
        _ => {}
    }
}

fn extract_function(
    node: &tree_sitter::Node,
    source: &[u8],
    parent: Option<&str>,
) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    let is_method = parent.is_some();

    // Check for decorators (they appear as previous siblings or named children)
    let mut sig = extract_signature(node, "body", source);

    // Look for decorator above
    if let Some(prev) = node.prev_sibling() {
        if prev.kind() == "decorator" {
            let dec_text = node_text(&prev, source);
            sig = format!("{}\n{}", dec_text, sig);
        }
    }

    Some(Symbol {
        name: name.to_string(),
        kind: if is_method {
            SymbolKind::Method
        } else {
            SymbolKind::Function
        },
        span: Span::from_node(node),
        signature: sig,
        doc_comment: extract_docstring(node, source),
        parent: parent.map(|s| s.to_string()),
        children: Vec::new(),
    })
}

fn extract_class(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    let mut methods = Vec::new();

    if let Some(body) = node.child_by_field_name("body") {
        let mut cursor = body.walk();
        for child in body.children(&mut cursor) {
            match child.kind() {
                "function_definition" => {
                    if let Some(m) = extract_function(&child, source, Some(name)) {
                        methods.push(m);
                    }
                }
                "decorated_definition" => {
                    // @property, @staticmethod, etc.
                    let mut inner = child.walk();
                    for grandchild in child.children(&mut inner) {
                        if grandchild.kind() == "function_definition" {
                            if let Some(m) = extract_function(&grandchild, source, Some(name)) {
                                methods.push(m);
                            }
                        }
                    }
                }
                _ => {}
            }
        }
    }

    Some(Symbol {
        name: name.to_string(),
        kind: SymbolKind::Class,
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_docstring(node, source),
        parent: None,
        children: methods,
    })
}

/// Extract Python docstring (first expression statement in body that is a string).
fn extract_docstring(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
    let body = node.child_by_field_name("body")?;
    let first = body.child(0)?;
    if first.kind() == "expression_statement" {
        let expr = first.child(0)?;
        if expr.kind() == "string" {
            return Some(node_text(&expr, source).to_string());
        }
    }
    // Fall back to comment-based doc
    extract_doc_comment(node, source)
}