from __future__ import annotations

from typing import TYPE_CHECKING

from libcst import Arg, Attribute, CSTTransformer, CSTVisitor, Name, RemoveFromParent

from core.languages.python.cst_utils import DefinitionRemover

if TYPE_CHECKING:
    from libcst import (
        BaseCompoundStatement,
        Call,
        ClassDef,
        EmptyLine,
        FunctionDef,
        Module,
        RemovalSentinel,
        SimpleStatementLine,
    )


class UsedNameCollector(CSTVisitor):
    """Collects all names that are actually used/referenced in the module.

    This includes:
    - Names used as base classes
    - Names that are called/instantiated (e.g., MyClass())
    - Names referenced in expressions
    """

    def __init__(self) -> None:
        super().__init__()
        self.used_names: set[str] = set()
        # Track names being defined so we don't count definitions as usage
        self.currently_defining: str | None = None

    def visit_ClassDef(self, node: ClassDef) -> bool:
        # Track that we're defining this class (don't count it as "used")
        self.currently_defining = node.name.value
        # Collect base class names as used
        for base in node.bases:
            if isinstance(base, Arg) and isinstance(base.value, Name):
                self.used_names.add(base.value.value)
            elif isinstance(base, Arg) and isinstance(base.value, Attribute):
                self.used_names.add(base.value.attr.value)
        return True

    def leave_ClassDef(self, original_node: ClassDef) -> None:  # noqa: ARG002
        self.currently_defining = None

    def visit_FunctionDef(self, node: FunctionDef) -> bool:
        # Track that we're defining this function
        self.currently_defining = node.name.value
        return True

    def leave_FunctionDef(self, original_node: FunctionDef) -> None:  # noqa: ARG002
        self.currently_defining = None

    def visit_Call(self, node: Call) -> bool:
        # Collect names that are called (instantiation or function calls)
        if isinstance(node.func, Name):
            self.used_names.add(node.func.value)
        elif isinstance(node.func, Attribute) and isinstance(node.func.value, Name):
            # For chained calls like module.func(), collect the module name
            self.used_names.add(node.func.value.value)
        return True

    def visit_Name(self, node: Name) -> bool:
        # Collect any name reference (but not the name being defined)
        if node.value != self.currently_defining:
            self.used_names.add(node.value)
        return True


class LeadingWhitespaceRemover(CSTTransformer):
    def __init__(self) -> None:
        self.found_first_statement = False

    def leave_EmptyLine(
        self,
        original_node: EmptyLine,  # noqa: ARG002
        updated_node: EmptyLine,
    ) -> EmptyLine | RemovalSentinel:
        if not self.found_first_statement:
            return RemoveFromParent()
        return updated_node

    def visit_SimpleStatementLine(self, node: SimpleStatementLine) -> bool:  # noqa: ARG002
        self.found_first_statement = True
        return True

    def visit_BaseCompoundStatement(self, node: BaseCompoundStatement) -> bool:  # noqa: ARG002
        self.found_first_statement = True
        return True


def delete_top_def_nodes(module: Module, deletable_list: list[str]) -> Module:
    """Delete top-level definitions from a module.

    First collects all names that are actually used in the module to protect them,
    then removes the specified definitions.

    Args:
        module: The CST module to modify
        deletable_list: List of names to delete (can include qualified names like "Class.method")

    Returns:
        Modified module with definitions removed and leading whitespace cleaned up

    """
    # Collect all names that are actually used in the module
    # This protects classes/functions that are referenced in test code from being deleted
    used_name_collector = UsedNameCollector()
    module.visit(used_name_collector)

    # Remove definitions, protecting names that are actually used
    remover = DefinitionRemover(deletable_list, protected_names=used_name_collector.used_names)
    module = module.visit(remover)
    return module.visit(LeadingWhitespaceRemover())
