"""Python call graph parser."""

from __future__ import annotations

from typing import TYPE_CHECKING, TypedDict

if TYPE_CHECKING:
    from pathlib import Path
    from tree_sitter import Node


class CallerInfo(TypedDict):
    """Information about the calling function."""

    name: str
    filepath: str
    line: int


from vibelexity.call_graph.parsers.base import CallGraphParser
from vibelexity.models import FunctionDefinition, CallEdge
from vibelexity.patterns.common import create_query, run_captures
from vibelexity.parsers.python import PY_LANGUAGE


class PythonCallParser(CallGraphParser):
    def _collect_ast_nodes(
        self, captures: list[tuple[Node, str]]
    ) -> tuple[
        list[tuple[str, int]],
        list[tuple[str, int]],
        list[tuple[str, int]],
        dict[tuple[str, int], list[str]],
    ]:
        names: list[tuple[str, int]] = []
        params: list[tuple[str, int]] = []
        classes: list[tuple[str, int]] = []
        name_to_params: dict[tuple[str, int], list[str]] = {}

        for node, capture_name in captures:
            text = node.text.decode() if node.text else ""
            line = node.start_point[0] + 1
            if capture_name == "class_name":
                classes.append((text, line))
            elif capture_name == "func_name":
                names.append((text, line))
                name_to_params[(text, line)] = []
            elif capture_name == "param":
                params.append((text, line))

        return names, params, classes, name_to_params

    def _map_params_to_functions(
        self,
        params: list[tuple[str, int]],
        names: list[tuple[str, int]],
        name_to_params: dict[tuple[str, int], list[str]],
    ) -> None:
        for param, param_line in params:
            best_name: tuple[str, int] | None = None
            min_gap = float("inf")

            for name, name_line in names:
                if name_line <= param_line <= name_line + 5:
                    gap = param_line - name_line
                    if gap < min_gap:
                        min_gap = gap
                        best_name = (name, name_line)

            if best_name and best_name in name_to_params:
                name_to_params[best_name].append(param)

    def _create_function_definitions(
        self,
        names: list[tuple[str, int]],
        classes: list[tuple[str, int]],
        name_to_params: dict[tuple[str, int], list[str]],
        filepath: str | Path,
    ) -> list[FunctionDefinition]:
        defs: list[FunctionDefinition] = []
        for name, name_line in names:
            def_type = "function"
            class_name: str | None = None

            for cls, cls_line in classes:
                if cls_line == name_line - 1:
                    def_type = "method"
                    class_name = cls
                    break

            func_def = FunctionDefinition(
                name=name,
                filepath=str(filepath),
                line=name_line,
                def_type=def_type,
                class_name=class_name,
                params=name_to_params.get((name, name_line), []),
            )
            defs.append(func_def)

        return defs

    def extract_function_definitions(
        self, root: Node, source: str, filepath: str | Path
    ) -> list[FunctionDefinition]:
        query_str = """(function_definition name: (identifier) @func_name)
        (function_definition parameters: (parameters (identifier) @param))
        (class_definition name: (identifier) @class_name)"""

        query = create_query(PY_LANGUAGE, query_str)
        captures = run_captures(query, root)

        names, params, classes, name_to_params = self._collect_ast_nodes(captures)
        self._map_params_to_functions(params, names, name_to_params)

        return self._create_function_definitions(
            names, classes, name_to_params, filepath
        )

    def extract_calls(
        self,
        root: Node,
        source: str,
        filepath: str | Path,
        all_functions: list[FunctionDefinition],
    ) -> list[CallEdge]:
        source_lines = source.split("\n")
        all_func_map: dict[tuple[str, str], FunctionDefinition] = {
            (f.name, f.filepath): f for f in all_functions
        }

        query_str = """(call) @call"""
        query = create_query(PY_LANGUAGE, query_str)
        captures = run_captures(query, root)

        calls: list[CallEdge] = []
        for node, _ in captures:
            call_site_line = node.start_point[0] + 1

            caller_info = self._get_caller_info(
                node, source_lines, all_func_map, str(filepath)
            )
            if not caller_info:
                continue

            callee_name, receiver, is_method = self._get_callee_info(node)

            call_type = "local"
            if receiver:
                call_type = "method"
            elif callee_name not in [f.name for f in all_functions]:
                call_type = "external"

            call_edge = CallEdge(
                caller_file=caller_info["filepath"],
                caller_function=caller_info["name"],
                caller_line=caller_info["line"],
                callee_file=None,
                callee_function=callee_name,
                callee_line=None,
                call_site_line=call_site_line,
                call_type=call_type,
                receiver_object=receiver,
            )
            calls.append(call_edge)

        return calls

    def _get_caller_info(
        self,
        call_node: Node,
        source_lines: list[str],
        all_func_map: dict,
        filepath: str,
    ) -> CallerInfo | None:
        line_num = call_node.start_point[0] + 1
        filepath_str = str(filepath)

        for func_key, func_def in all_func_map.items():
            if func_def.filepath != filepath_str:
                continue
            if line_num >= func_def.line and line_num <= self._get_func_end_line(
                func_def.line, source_lines
            ):
                return {
                    "name": func_def.name,
                    "filepath": func_def.filepath,
                    "line": func_def.line,
                }

        return None

    def _get_func_end_line(self, start_line: int, source_lines: list[str]) -> int:
        for i in range(start_line, len(source_lines)):
            line = source_lines[i]
            if self._is_top_level_def(line):
                return i
        return len(source_lines)

    def _is_top_level_def(self, line: str) -> bool:
        stripped = line.lstrip()
        if stripped.startswith("def ") and not stripped.startswith("    "):
            return True
        if stripped.startswith("class ") and not stripped.startswith("    "):
            return True
        return False

    def _get_callee_info(self, call_node: Node) -> tuple[str, str | None, bool]:
        for child in call_node.children:
            if child.type == "identifier":
                return child.text.decode() if child.text else "<unknown>", None, False
            if child.type == "attribute":
                return self._get_attribute_callee(child)
            if child.type == "member":
                if child.parent and child.parent.type == "call":
                    return self._get_member_callee(child)
        return "<unknown>", None, False

    def _get_member_callee(self, member_node: Node) -> tuple[str, str | None, bool]:
        member_text = member_node.text.decode() if member_node.text else ""
        return member_text, None, False

    def _get_attribute_callee(self, attr_node: Node) -> tuple[str, str | None, bool]:
        parts: list[str] = []

        def extract_parts(node: Node) -> None:
            if node.type == "identifier":
                parts.append(node.text.decode() if node.text else "")
                return
            if node.type == "attribute":
                for c in node.children:
                    if c.type != ".":
                        extract_parts(c)

        extract_parts(attr_node)

        if not parts:
            return "<unknown>", None, False
        if len(parts) >= 2:
            receiver = ".".join(parts[:-1])
            return parts[-1], receiver, True
        return parts[0], None, False
