import ast
from typing import Any


class Signature:
    """
    A class for analyzing Python function or class signatures using the AST (Abstract Syntax Tree) module.

    This enhanced version automatically detects whether the given body is a function or class.
    If it's a class, it attempts to extract parameters and type hints from its __init__ method.

    Example
    -------
    >>> code = '''
    class Example:
            def __init__(self, name: str, age: int = 20):
                    self.name = name
                    self.age = age
    '''
    >>> sig = Signature(code)
    >>> print(sig.signature())
    {
            'name': 'Example',
            'parameters': ['self', 'name', 'age'],
            'annotations': {'name': 'str', 'age': 'int'}
    }
    """

    def __init__(self, defbody: str) -> None:
        if not isinstance(defbody, str):
            raise TypeError(
                f"Expected 'defbody' to be a string, got {type(defbody).__name__}"
            )
        self.body = defbody.strip()
        self._tree = ast.parse(defbody)
        self._nodes = list(ast.walk(self._tree))

    # ------------------------------------------------------

    def _find_main_node(self) -> ast.AST | None:
        """Return the main node (FunctionDef or ClassDef) in the body."""
        for node in self._nodes:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                return node
        return None

    # ------------------------------------------------------

    def _find_init_node(self, class_node: ast.ClassDef) -> ast.FunctionDef | None:
        """If class has __init__, return it."""
        for node in class_node.body:
            if isinstance(node, ast.FunctionDef) and node.name == "__init__":
                return node
        return None

    # ------------------------------------------------------

    def _extract_params(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> list[str]:
        """Extract all parameter names from a given function node."""
        params = []

        def flatten(obj: Any) -> list:
            if isinstance(obj, list):
                result = []
                for o in obj:
                    result.extend(flatten(o))
                return result
            else:
                return [obj]

        if isinstance(node.args, ast.arguments):
            subparams = []

            if node.args.args:
                subparams.append([arg.arg for arg in node.args.args])
            if node.args.vararg:
                subparams.append(node.args.vararg.arg)
            if node.args.kwarg:
                subparams.append(node.args.kwarg.arg)

            params.extend(flatten(subparams))

        return params

    # ------------------------------------------------------

    def _extract_annotations(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> dict[str, str]:
        """Extract type annotations for parameters and return value."""
        annotations = {}

        if node.args.args:
            for arg in node.args.args:
                if arg.annotation:
                    type_name = getattr(
                        arg.annotation, "id", getattr(arg.annotation, "attr", None)
                    )
                    if type_name:
                        annotations[arg.arg] = type_name

        if node.returns:
            return_type = getattr(
                node.returns, "id", getattr(node.returns, "attr", None)
            )
            if return_type:
                annotations["return"] = return_type

        return annotations

    # ------------------------------------------------------

    def signature(self) -> dict[str, Any]:
        """
        Generate a complete structured signature.

        If the body defines a class with an __init__ method, extracts parameters and annotations from it.
        Otherwise, extracts from the top-level function definition.
        """
        main_node = self._find_main_node()
        if not main_node:
            return {"name": None, "parameters": [], "annotations": {}}

        # CASE 1: Function or async function
        if isinstance(main_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            return {
                "name": main_node.name,
                "parameters": self._extract_params(main_node),
                "annotations": self._extract_annotations(main_node),
            }

        # CASE 2: Class
        elif isinstance(main_node, ast.ClassDef):
            init_node = self._find_init_node(main_node)
            if init_node:
                return {
                    "name": main_node.name,
                    "parameters": self._extract_params(init_node),
                    "annotations": self._extract_annotations(init_node),
                }
            else:
                return {
                    "name": main_node.name,
                    "parameters": [],
                    "annotations": {},
                }
