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:
    """
    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


def extract_import_info(
    alias: cst.ImportAlias, module_name: str = ""
) -> tuple[str, str, str]:
    """
    Extract (available_name, module_path, original_name) from an import alias.
    """
    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:
        available_name = original_name
    else:
        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.
    """
    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.
    """
    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 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 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.
    """

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

        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()

        self._is_killable_class: bool = False
        self._potential_methods: set[str] = set()

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

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:
        is_top_level = self._is_top_level_class()
        self._visit_class()

        if not is_top_level:
            return True

        class_name = node.name.value

        if self._should_remove(class_name):
            self._is_killable_class = True
            return False

        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

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

        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:
        self._visit_function()
        return True

    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

        if self._is_top_level() and self._should_remove(func_name):
            self.removed_names.add(func_name)
            return cst.RemovalSentinel.REMOVE

        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 or 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 or 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:
    """
    Convert a file path to a dotted module path.
    """
    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.
    """
    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.
    """
    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.
    """
    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 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


_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.
    """
    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 an integer, or None.
    """
    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 int(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


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

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

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


class InvalidEllipsisVisitor(cst.CSTVisitor):
    """
    Detect ellipsis in invalid locations.
    """

    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.
        """
        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 *node* is within a type annotation.
        """
        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 *node* is a function body 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 *module* contains ellipsis in invalid locations.
    """
    visitor = InvalidEllipsisVisitor()
    MetadataWrapper(module).visit(visitor)
    return visitor.invalid_found


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


def unparse_parse_source(source: str) -> str:
    """
    Parse source 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:
    """
    Search an AST for a FunctionDef node named '__init__'.
    """
    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
