"""Generic libCST and AST utilities for parsing and manipulating Python code.

This module provides commonly used CST/AST utilities that are shared across
optimizer, code_repair, and testgen modules.
"""

from __future__ import annotations

import ast
import operator
from collections import deque
from functools import lru_cache
from typing import TYPE_CHECKING

import libcst as cst
from libcst.metadata import MetadataWrapper, ParentNodeProvider, PositionProvider

if TYPE_CHECKING:
    from collections.abc import Iterator


class DepthTrackingMixin:
    """Mixin that tracks function and class nesting depth for CST transformers."""

    def __init__(self) -> None:
        self._function_depth: int = 0
        self._class_depth: int = 0

    def _visit_function(self) -> None:
        self._function_depth += 1

    def _leave_function(self) -> None:
        self._function_depth -= 1

    def _visit_class(self) -> None:
        self._class_depth += 1

    def _leave_class(self) -> None:
        self._class_depth -= 1

    def _is_top_level(self) -> bool:
        return self._function_depth == 0 and self._class_depth == 0

    def _is_top_level_function(self) -> bool:
        return self._function_depth == 0 and self._class_depth == 0

    def _is_top_level_class(self) -> bool:
        return self._class_depth == 0 and self._function_depth == 0

    def _is_inside_class(self) -> bool:
        return self._class_depth > 0

    def _is_inside_function(self) -> bool:
        return self._function_depth > 0


# ==============================================================================
# Import Extraction Utilities
# ==============================================================================


def extract_import_info(alias: cst.ImportAlias, module_name: str = "") -> tuple[str, str, str]:
    """Extract import information from a single import alias.

    Args:
        alias: The ImportAlias node (e.g., `foo` or `foo as bar`)
        module_name: The module being imported from (for ImportFrom statements)

    Returns:
        Tuple of (available_name, module_path, original_name) where:
        - available_name: The name available in the local namespace
        - module_path: The full module path for the import
        - original_name: The original name being imported

    Examples:
        `from pkg import foo` -> ("foo", "pkg", "foo")
        `from pkg import foo as bar` -> ("bar", "pkg", "foo")
        `import pkg.sub` -> ("pkg", "pkg.sub", "pkg.sub")
        `import pkg.sub as p` -> ("p", "pkg.sub", "pkg.sub")

    """
    original_name = alias.name.value if isinstance(alias.name, cst.Name) else get_dotted_name(alias.name)

    if alias.asname and isinstance(alias.asname.name, cst.Name):
        available_name = alias.asname.name.value
    elif module_name:
        # ImportFrom: available name is the imported name
        available_name = original_name
    else:
        # Import: available name is the first part of dotted name
        available_name = original_name.split(".")[0]

    return available_name, module_name, original_name


def extract_imports_from_import(node: cst.Import) -> dict[str, tuple[str, str]]:
    """Extract all imported names from an Import statement.

    Returns:
        Dict mapping available_name -> (module_path, original_name)

    Example:
        `import os, sys as system` -> {"os": ("os", "os"), "system": ("sys", "sys")}

    """
    if isinstance(node.names, cst.ImportStar):
        return {}

    result: dict[str, tuple[str, str]] = {}
    for alias in node.names:
        available_name, _, original_name = extract_import_info(alias)
        if available_name:
            result[available_name] = (original_name, original_name)
    return result


def extract_imports_from_import_from(node: cst.ImportFrom) -> dict[str, tuple[str, str]]:
    """Extract all imported names from an ImportFrom statement.

    Returns:
        Dict mapping available_name -> (module_path, original_name)

    Example:
        `from os.path import join, exists as ex` -> {"join": ("os.path", "join"), "ex": ("os.path", "exists")}

    """
    if isinstance(node.names, cst.ImportStar):
        return {}

    module_name = get_dotted_name(node.module) if node.module else ""
    result: dict[str, tuple[str, str]] = {}
    for alias in node.names:
        available_name, _, original_name = extract_import_info(alias, module_name)
        if available_name:
            result[available_name] = (module_name, original_name)
    return result


def collect_imported_names_from_import(node: cst.Import) -> set[str]:
    """Collect just the available names from an Import statement."""
    if isinstance(node.names, cst.ImportStar):
        return set()
    return set(extract_imports_from_import(node).keys())


def collect_imported_names_from_import_from(node: cst.ImportFrom) -> set[str]:
    """Collect just the available names from an ImportFrom statement."""
    if isinstance(node.names, cst.ImportStar):
        return set()
    return set(extract_imports_from_import_from(node).keys())


class DefinitionRemover(DepthTrackingMixin, cst.CSTTransformer):
    """Remove top-level class and function definitions by name.

    Supports:
    - Simple names: "MyClass", "my_function"
    - Qualified names: "ClassName.method_name" (removes class if method matches)
    - Protected names that should not be removed
    - Tracking which names were actually removed

    Args:
        names_to_remove: Names to remove (can be simple or qualified like "Class.method")
        protected_names: Names that should never be removed (takes precedence)

    Attributes:
        removed_names: Set of names that were actually removed (populated after visiting)

    """

    def __init__(self, names_to_remove: set[str] | list[str], protected_names: set[str] | None = None) -> None:
        DepthTrackingMixin.__init__(self)
        cst.CSTTransformer.__init__(self)

        # Parse qualified names (e.g., "ClassName.method_name")
        names_list = list(names_to_remove)
        self._qualified_names: list[list[str]] = [s.split(".")[-2:] for s in names_list]
        self._simple_names: set[str] = {pair[-1] for pair in self._qualified_names if pair}

        self.protected_names: set[str] = protected_names or set()
        self.removed_names: set[str] = set()

        # State for qualified name matching (remove class if method matches)
        self._is_killable_class: bool = False
        self._potential_methods: set[str] = set()

    def _should_remove(self, name: str) -> bool:
        """Check if a name should be removed (not protected and in removal set)."""
        return name in self._simple_names and name not in self.protected_names

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
        # Check if top-level BEFORE incrementing depth
        is_top_level = self._is_top_level_class()
        self._visit_class()

        if not is_top_level:
            return True  # Still visit children for nested classes

        class_name = node.name.value

        # Check if class itself should be removed
        if self._should_remove(class_name):
            self._is_killable_class = True
            return False

        # Check if any methods in qualified names match this class
        potential_methods = {
            pair[-1]
            for pair in self._qualified_names
            if len(pair) > 1 and pair[-2] == class_name and pair[-1] not in self.protected_names
        }
        if potential_methods:
            self._potential_methods = potential_methods
            return True

        return True

    def leave_ClassDef(
        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
    ) -> cst.ClassDef | cst.RemovalSentinel:
        self._leave_class()

        # Only remove if we're back at top level and class is marked for removal
        if self._is_top_level() and self._is_killable_class:
            self._is_killable_class = False
            self._potential_methods = set()
            self.removed_names.add(original_node.name.value)
            return cst.RemovalSentinel.REMOVE

        self._potential_methods = set()
        return updated_node

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:  # noqa: ARG002
        self._visit_function()
        return True  # Visit children

    def leave_FunctionDef(
        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
    ) -> cst.FunctionDef | cst.RemovalSentinel:
        self._leave_function()

        func_name = original_node.name.value

        # Top-level function removal (check after decrementing depth)
        if self._is_top_level() and self._should_remove(func_name):
            self.removed_names.add(func_name)
            return cst.RemovalSentinel.REMOVE

        # Method matching for qualified names (marks parent class for removal)
        if func_name in self._potential_methods:
            self._is_killable_class = True
            self._potential_methods = set()

        return updated_node


class ImportTrackingVisitor(ast.NodeVisitor):
    """Base AST visitor that tracks imported names."""

    def __init__(self) -> None:
        self.imported_names: set[str] = set()

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            name = alias.asname if alias.asname else alias.name.split(".")[0]
            self.imported_names.add(name)
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for alias in node.names:
            if alias.name == "*":
                continue
            name = alias.asname if alias.asname else alias.name
            self.imported_names.add(name)
        self.generic_visit(node)


@lru_cache(maxsize=128)
def parse_module_to_cst(module_str: str) -> cst.Module:
    """Parse a module string into its libCST representation."""
    return cst.parse_module(module_str)


def file_path_to_module_path(file_path: str) -> str:
    r"""Convert a file path to a module path.

    Examples:
        "path/to/module.py" -> "path.to.module"
        "path\\to\\module.py" -> "path.to.module"

    """
    file_path = file_path.removesuffix(".py")
    return file_path.replace("/", ".").replace("\\", ".")


def get_dotted_name(node: cst.BaseExpression | None) -> str:
    """Extract dotted module path from a CST node."""
    if node is None:
        return ""
    if isinstance(node, cst.Name):
        return node.value
    if isinstance(node, cst.Attribute):
        base = get_dotted_name(node.value)
        if base:
            return f"{base}.{node.attr.value}"
        return node.attr.value
    return ""


def get_base_class_name(base: cst.Arg) -> str | None:
    """Get the name of a base class from a CST Arg node.

    Handles both simple names (Bar) and attribute access (module.Bar).
    Returns the class name or None if unable to extract.
    """
    expr = base.value
    if isinstance(expr, cst.Name):
        return expr.value
    if isinstance(expr, cst.Attribute):
        return expr.attr.value
    return None


def has_decorator(class_node: cst.ClassDef, decorator_name: str) -> bool:
    """Check if a class has a specific decorator (with or without arguments)."""
    for decorator in class_node.decorators:
        dec = decorator.decorator
        if isinstance(dec, cst.Name) and dec.value == decorator_name:
            return True
        if isinstance(dec, cst.Call) and isinstance(dec.func, cst.Name) and dec.func.value == decorator_name:
            return True
    return False


def get_node_source_text(
    node: cst.CSTNode, source_lines: list[str], wrapper: MetadataWrapper, fallback: str = "..."
) -> str:
    """Get the source text of a CST node using position metadata.

    This is useful when you need the original source representation of a node,
    such as for type annotations or default values.

    Args:
        node: The CST node to get source text for.
        source_lines: The original source code split by lines.
        wrapper: The MetadataWrapper that provides position information.
        fallback: Value to return if position cannot be determined.

    Returns:
        The source text of the node, or fallback if position is unavailable.

    """
    pos = wrapper.resolve(PositionProvider).get(node)
    if pos is None:
        return fallback
    return source_lines[pos.start.line - 1][pos.start.column : pos.end.column]


def build_module_path(module_path: str) -> cst.Attribute | cst.Name:
    """Build a libcst Attribute or Name node from a dotted module path string."""
    parts = module_path.split(".")
    result: cst.Attribute | cst.Name = cst.Name(parts[0])
    for part in parts[1:]:
        result = cst.Attribute(value=result, attr=cst.Name(part))
    return result


# ==============================================================================
# CST Expression Evaluation Utilities
# ==============================================================================

_BINARY_OPS = {cst.Power: operator.pow, cst.Multiply: operator.mul, cst.Add: operator.add, cst.Subtract: operator.sub}


def make_number_node(value: int) -> cst.Integer | cst.UnaryOperation:
    """Create a CST node representing an integer value, handling negative numbers."""
    if value >= 0:
        return cst.Integer(str(value))
    return cst.UnaryOperation(operator=cst.Minus(), expression=cst.Integer(str(abs(value))))


def evaluate_expression(node: cst.BaseExpression) -> int | None:
    """Evaluate a CST expression node to get its integer value, or None if not evaluable."""
    if isinstance(node, cst.Integer):
        return int(node.value, 0)
    if isinstance(node, cst.Float):
        return int(float(node.value))
    if isinstance(node, cst.BinaryOperation):
        left = evaluate_expression(node.left)
        right = evaluate_expression(node.right)
        op = _BINARY_OPS.get(type(node.operator))
        if left is not None and right is not None and op is not None:
            return op(left, right)
    if isinstance(node, cst.UnaryOperation) and isinstance(node.operator, cst.Minus):
        val = evaluate_expression(node.expression)
        if val is not None:
            return -val
    return None


# ==============================================================================
# Ellipsis Detection Utilities
# ==============================================================================


class AnyEllipsisVisitor(cst.CSTVisitor):
    """Visitor to detect if any ellipsis exists in the code."""

    def __init__(self) -> None:
        self.found = False

    def visit_Ellipsis(self, node: cst.Ellipsis) -> bool:  # noqa: ARG002
        self.found = True
        return False


class InvalidEllipsisVisitor(cst.CSTVisitor):
    """Visitor to detect ellipsis in invalid locations (not in type annotations or function body statements)."""

    METADATA_DEPENDENCIES = (ParentNodeProvider,)

    def __init__(self) -> None:
        self.invalid_found = False

    def visit_Ellipsis(self, node: cst.Ellipsis) -> bool:
        if self._is_in_type_annotation_context(node):
            return False
        if self._is_function_body_ellipsis(node):
            return False
        self.invalid_found = True
        return False

    def _iter_ancestors(self, node: cst.CSTNode, max_depth: int = 10) -> Iterator[cst.CSTNode]:
        """Yield ancestors walking up the tree, up to max_depth."""
        current = node
        for _ in range(max_depth):
            parent = self.get_metadata(ParentNodeProvider, current, None)
            if parent is None:
                break
            yield parent
            current = parent

    def _is_in_type_annotation_context(self, node: cst.CSTNode) -> bool:
        """Check if the node is within a Subscript or Tuple (type annotation context)."""
        return any(isinstance(ancestor, (cst.Subscript, cst.Tuple)) for ancestor in self._iter_ancestors(node))

    def _is_function_body_ellipsis(self, node: cst.Ellipsis) -> bool:
        """Check if the ellipsis is a function body statement (valid for Protocol/abstract methods).

        Pattern: FunctionDef -> IndentedBlock -> SimpleStatementLine -> Expr -> Ellipsis
        """
        ancestors = list(self._iter_ancestors(node, max_depth=4))
        if len(ancestors) < 4:
            return False
        expected_types = (cst.Expr, cst.SimpleStatementLine, cst.IndentedBlock, cst.FunctionDef)
        return all(isinstance(ancestor, expected) for ancestor, expected in zip(ancestors, expected_types, strict=True))


def ellipsis_in_cst_not_types(module: cst.Module) -> bool:
    """Check if the module contains ellipsis in invalid locations (not type annotations or function bodies)."""
    visitor = InvalidEllipsisVisitor()
    MetadataWrapper(module).visit(visitor)
    return visitor.invalid_found


def any_ellipsis_in_cst(module: cst.Module) -> bool:
    """Check if the module contains any ellipsis."""
    visitor = AnyEllipsisVisitor()
    module.visit(visitor)
    return visitor.found


# ==============================================================================
# AST Utilities
# ==============================================================================


def unparse_parse_source(source: str) -> str:
    """Parse source code and unparse it back, normalizing formatting."""
    return ast.unparse(ast.parse(source))


def compare_unparsed_ast_to_source(unparsed_ast: str, source: str) -> bool:
    """Compare an unparsed AST string to a source string after normalizing."""
    return unparsed_ast == unparse_parse_source(source)


def find_init(node: ast.AST) -> bool:
    """Recursively search an AST for a FunctionDef node named '__init__'.

    Args:
        node: An AST node to search

    Returns:
        bool: True if an __init__ function is found, False otherwise

    """
    stack = deque([node])
    while stack:
        current = stack.pop()
        if isinstance(current, ast.FunctionDef) and current.name == "__init__":
            return True
        stack.extend(ast.iter_child_nodes(current))
    return False
