"""
Python-specific code extractor using Tree-sitter.
"""

from tree_sitter import Node

from .types import (
    ExtractedFunction,
    ExtractedClass,
    ExtractedCall,
    ExtractedImport,
)


class PythonExtractor:
    """Extractor for Python code."""

    def __init__(self, source: bytes, module_qn: str):
        self.source = source
        self.module_qn = module_qn

    def get_node_text(self, node: Node) -> str:
        """Get text content of a node."""
        return self.source[node.start_byte : node.end_byte].decode("utf-8")

    def get_node_lines(self, node: Node) -> tuple[int, int]:
        """Get start and end line numbers (1-indexed)."""
        return node.start_point[0] + 1, node.end_point[0] + 1

    def extract_functions(self, root: Node) -> list[ExtractedFunction]:
        """Extract all top-level functions (not methods)."""
        functions: list[ExtractedFunction] = []
        self._extract_functions_recursive(root, functions, is_top_level=True)
        return functions

    def _extract_functions_recursive(
        self,
        node: Node,
        functions: list[ExtractedFunction],
        is_top_level: bool = False,
    ) -> None:
        """Recursively extract functions, skipping class methods."""
        for child in node.children:
            if child.type == "function_definition":
                # Skip if inside a class (methods are handled separately)
                if not self._is_inside_class(child):
                    func = self._extract_function(child, class_name=None)
                    if func:
                        functions.append(func)
            elif child.type == "decorated_definition":
                # Handle decorated functions
                func_node = self._get_decorated_function(child)
                if func_node and not self._is_inside_class(func_node):
                    func = self._extract_function(
                        func_node, class_name=None, decorator_parent=child
                    )
                    if func:
                        functions.append(func)
            elif child.type not in ("class_definition",):
                # Recurse into other nodes but not classes
                self._extract_functions_recursive(child, functions)

    def _is_inside_class(self, node: Node) -> bool:
        """Check if a node is inside a class definition."""
        parent = node.parent
        while parent:
            if parent.type == "class_definition":
                return True
            if parent.type == "decorated_definition":
                parent = parent.parent
                continue
            if parent.type == "block":
                parent = parent.parent
                continue
            parent = parent.parent
        return False

    def _get_decorated_function(self, decorated_node: Node) -> Node | None:
        """Get the function node from a decorated_definition."""
        for child in decorated_node.children:
            if child.type == "function_definition":
                return child
        return None

    def _get_decorated_class(self, decorated_node: Node) -> Node | None:
        """Get the class node from a decorated_definition."""
        for child in decorated_node.children:
            if child.type == "class_definition":
                return child
        return None

    def _extract_function(
        self,
        node: Node,
        class_name: str | None,
        decorator_parent: Node | None = None,
    ) -> ExtractedFunction | None:
        """Extract information from a function_definition node."""
        # Get function name
        name_node = node.child_by_field_name("name")
        if not name_node:
            return None
        name = self.get_node_text(name_node)

        # Build qualified name
        if class_name:
            qualified_name = f"{self.module_qn}.{class_name}.{name}"
        else:
            qualified_name = f"{self.module_qn}.{name}"

        # Get line numbers
        start_line, end_line = self.get_node_lines(node)

        # Get parameters
        parameters = self._extract_parameters(node)

        # Build signature
        params_node = node.child_by_field_name("parameters")
        params_text = self.get_node_text(params_node) if params_node else "()"

        return_type_node = node.child_by_field_name("return_type")
        return_type = self.get_node_text(return_type_node) if return_type_node else None

        if return_type:
            signature = f"def {name}{params_text} -> {return_type}"
        else:
            signature = f"def {name}{params_text}"

        # Check if async
        is_async = self._is_async_function(node)
        if is_async:
            signature = "async " + signature

        # Get docstring
        docstring = self._extract_docstring(node)

        # Get decorators
        decorators = self._extract_decorators(node, decorator_parent)

        # Get source code
        source_code = self.get_node_text(node)

        return ExtractedFunction(
            name=name,
            qualified_name=qualified_name,
            start_line=start_line,
            end_line=end_line,
            signature=signature,
            docstring=docstring,
            decorators=decorators,
            is_async=is_async,
            is_method=class_name is not None,
            class_name=class_name,
            parameters=parameters,
            return_type=return_type,
            source_code=source_code,
            node=node,
        )

    def _extract_parameters(self, func_node: Node) -> list[str]:
        """Extract parameter names from function."""
        parameters: list[str] = []
        params_node = func_node.child_by_field_name("parameters")
        if not params_node:
            return parameters

        for child in params_node.children:
            if child.type == "identifier":
                parameters.append(self.get_node_text(child))
            elif child.type in (
                "typed_parameter",
                "default_parameter",
                "typed_default_parameter",
            ):
                # Get the name part
                name_node = child.child_by_field_name("name")
                if name_node:
                    parameters.append(self.get_node_text(name_node))
                elif child.children and child.children[0].type == "identifier":
                    parameters.append(self.get_node_text(child.children[0]))
            elif child.type == "list_splat_pattern":
                # *args
                for c in child.children:
                    if c.type == "identifier":
                        parameters.append("*" + self.get_node_text(c))
            elif child.type == "dictionary_splat_pattern":
                # **kwargs
                for c in child.children:
                    if c.type == "identifier":
                        parameters.append("**" + self.get_node_text(c))

        return parameters

    def _is_async_function(self, node: Node) -> bool:
        """Check if function is async."""
        if node.children:
            first_child = node.children[0]
            if first_child.type == "async":
                return True
        return False

    def _extract_docstring(self, node: Node) -> str | None:
        """Extract docstring from function or class."""
        body_node = node.child_by_field_name("body")
        if not body_node or not body_node.children:
            return None

        # First statement in body
        first_stmt = body_node.children[0]
        if first_stmt.type == "expression_statement":
            if first_stmt.children:
                expr = first_stmt.children[0]
                if expr.type == "string":
                    docstring = self.get_node_text(expr)
                    # Strip quotes
                    if docstring.startswith('"""') or docstring.startswith("'''"):
                        return docstring[3:-3].strip()
                    elif docstring.startswith('"') or docstring.startswith("'"):
                        return docstring[1:-1].strip()
        return None

    def _extract_decorators(
        self, node: Node, decorator_parent: Node | None
    ) -> list[str]:
        """Extract decorators from a function or class."""
        decorators: list[str] = []

        # If we have a decorated_definition parent, get decorators from there
        parent = decorator_parent or node.parent
        if parent and parent.type == "decorated_definition":
            for child in parent.children:
                if child.type == "decorator":
                    dec_text = self.get_node_text(child)
                    decorators.append(dec_text)

        return decorators

    def extract_classes(self, root: Node) -> list[ExtractedClass]:
        """Extract all classes from the AST."""
        classes: list[ExtractedClass] = []
        self._extract_classes_recursive(root, classes)
        return classes

    def _extract_classes_recursive(
        self, node: Node, classes: list[ExtractedClass]
    ) -> None:
        """Recursively extract classes."""
        for child in node.children:
            if child.type == "class_definition":
                cls = self._extract_class(child)
                if cls:
                    classes.append(cls)
            elif child.type == "decorated_definition":
                class_node = self._get_decorated_class(child)
                if class_node:
                    cls = self._extract_class(class_node, decorator_parent=child)
                    if cls:
                        classes.append(cls)
            else:
                self._extract_classes_recursive(child, classes)

    def _extract_class(
        self, node: Node, decorator_parent: Node | None = None
    ) -> ExtractedClass | None:
        """Extract information from a class_definition node."""
        # Get class name
        name_node = node.child_by_field_name("name")
        if not name_node:
            return None
        name = self.get_node_text(name_node)

        qualified_name = f"{self.module_qn}.{name}"
        start_line, end_line = self.get_node_lines(node)

        # Get base classes
        bases = self._extract_bases(node)

        # Get docstring
        docstring = self._extract_docstring(node)

        # Get decorators
        decorators = self._extract_decorators(node, decorator_parent)

        # Get methods
        methods = self._extract_methods(node, name)

        # Get source code
        source_code = self.get_node_text(node)

        return ExtractedClass(
            name=name,
            qualified_name=qualified_name,
            start_line=start_line,
            end_line=end_line,
            docstring=docstring,
            decorators=decorators,
            bases=bases,
            methods=methods,
            source_code=source_code,
            node=node,
        )

    def _extract_bases(self, class_node: Node) -> list[str]:
        """Extract base class names."""
        bases: list[str] = []
        superclasses = class_node.child_by_field_name("superclasses")
        if superclasses:
            # argument_list contains the base classes
            for child in superclasses.children:
                if child.type == "identifier":
                    bases.append(self.get_node_text(child))
                elif child.type == "attribute":
                    bases.append(self.get_node_text(child))
        return bases

    def _extract_methods(
        self, class_node: Node, class_name: str
    ) -> list[ExtractedFunction]:
        """Extract methods from a class."""
        methods: list[ExtractedFunction] = []
        body_node = class_node.child_by_field_name("body")
        if not body_node:
            return methods

        for child in body_node.children:
            if child.type == "function_definition":
                method = self._extract_function(child, class_name=class_name)
                if method:
                    methods.append(method)
            elif child.type == "decorated_definition":
                func_node = self._get_decorated_function(child)
                if func_node:
                    method = self._extract_function(
                        func_node, class_name=class_name, decorator_parent=child
                    )
                    if method:
                        methods.append(method)

        return methods

    def extract_calls(self, node: Node, caller_qn: str) -> list[ExtractedCall]:
        """Extract function calls from within a function/method node."""
        calls: list[ExtractedCall] = []
        self._extract_calls_recursive(node, caller_qn, calls)
        return calls

    def _extract_calls_recursive(
        self, node: Node, caller_qn: str, calls: list[ExtractedCall]
    ) -> None:
        """Recursively extract call expressions."""
        if node.type == "call":
            call = self._extract_single_call(node, caller_qn)
            if call:
                calls.append(call)

        for child in node.children:
            # Don't recurse into nested function definitions
            if child.type not in ("function_definition", "class_definition"):
                self._extract_calls_recursive(child, caller_qn, calls)

    def _extract_single_call(
        self, call_node: Node, caller_qn: str
    ) -> ExtractedCall | None:
        """Extract a single call expression."""
        func_node = call_node.child_by_field_name("function")
        if not func_node:
            return None

        line = call_node.start_point[0] + 1
        call_text = self.get_node_text(func_node)

        # Determine if it's a method call (obj.method)
        receiver = None
        name = call_text

        if func_node.type == "attribute":
            # Method call: obj.method()
            obj_node = func_node.child_by_field_name("object")
            attr_node = func_node.child_by_field_name("attribute")
            if obj_node and attr_node:
                receiver = self.get_node_text(obj_node)
                name = self.get_node_text(attr_node)

        # Extract arguments (simplified - just the text)
        arguments: list[str] = []
        args_node = call_node.child_by_field_name("arguments")
        if args_node:
            for child in args_node.children:
                if child.type not in ("(", ")", ","):
                    arguments.append(self.get_node_text(child))

        return ExtractedCall(
            name=name,
            line=line,
            caller_qualified_name=caller_qn,
            receiver=receiver,
            arguments=arguments,
        )

    def extract_imports(self, root: Node) -> list[ExtractedImport]:
        """Extract all import statements."""
        imports: list[ExtractedImport] = []
        self._extract_imports_recursive(root, imports)
        return imports

    def _extract_imports_recursive(
        self, node: Node, imports: list[ExtractedImport]
    ) -> None:
        """Recursively extract import statements."""
        if node.type == "import_statement":
            self._extract_import_statement(node, imports)
        elif node.type == "import_from_statement":
            self._extract_import_from_statement(node, imports)
        else:
            for child in node.children:
                self._extract_imports_recursive(child, imports)

    def _extract_import_statement(
        self, node: Node, imports: list[ExtractedImport]
    ) -> None:
        """Extract 'import x' or 'import x as y' statements."""
        line = node.start_point[0] + 1

        for child in node.children:
            if child.type == "dotted_name":
                module = self.get_node_text(child)
                imports.append(ExtractedImport(module=module, line=line))
            elif child.type == "aliased_import":
                # import x as y
                name_node = child.child_by_field_name("name")
                alias_node = child.child_by_field_name("alias")
                if name_node:
                    module = self.get_node_text(name_node)
                    alias = self.get_node_text(alias_node) if alias_node else None
                    imports.append(
                        ExtractedImport(module=module, alias=alias, line=line)
                    )

    def _extract_import_from_statement(
        self, node: Node, imports: list[ExtractedImport]
    ) -> None:
        """Extract 'from x import y' statements."""
        line = node.start_point[0] + 1

        # Get module name
        module_name = None
        for child in node.children:
            if child.type == "dotted_name":
                module_name = self.get_node_text(child)
                break
            elif child.type == "relative_import":
                module_name = self.get_node_text(child)
                break

        if not module_name:
            return

        # Get imported names
        for child in node.children:
            if child.type == "wildcard_import":
                imports.append(
                    ExtractedImport(module=module_name, is_wildcard=True, line=line)
                )
            elif child.type == "dotted_name" and child != node.children[1]:
                # Named import without alias
                name = self.get_node_text(child)
                imports.append(
                    ExtractedImport(module=module_name, name=name, line=line)
                )
            elif child.type == "aliased_import":
                # from x import y as z
                name_node = child.child_by_field_name("name")
                alias_node = child.child_by_field_name("alias")
                if name_node:
                    name = self.get_node_text(name_node)
                    alias = self.get_node_text(alias_node) if alias_node else None
                    imports.append(
                        ExtractedImport(
                            module=module_name, name=name, alias=alias, line=line
                        )
                    )
            elif child.type == "import_prefix":
                # This is the "from" keyword, skip
                continue
