from __future__ import annotations

from libcst import (
    BaseCompoundStatement,
    ClassDef,
    CSTTransformer,
    EmptyLine,
    FunctionDef,
    Module,
    RemovalSentinel,
    RemoveFromParent,
    SimpleStatementLine,
)


class TopDefTerminator(CSTTransformer):
    def __init__(self, hit_list: list[str]) -> None:
        super().__init__()
        self.hit_list_parents_names: list[list[str]] = [s.split(".")[-2:] for s in hit_list]
        self.hit_list_names: set[str] = {pair[-1] for pair in self.hit_list_parents_names if pair}
        self.level: int = 0
        self.is_killable_class: bool = False
        self.potential_methods: set[str] = set()

    def visit_ClassDef(self, node: ClassDef) -> bool:
        self.level += 1
        if self.level > 1:
            return False
        if node.name.value in self.hit_list_names:
            self.is_killable_class = True
            return False
        if potential_methods := {
            pair[-1] for pair in self.hit_list_parents_names if len(pair) > 1 and pair[-2] == node.name.value
        }:
            self.potential_methods = potential_methods
            return True
        return False

    def leave_ClassDef(self, original_node: ClassDef, updated_node: ClassDef) -> ClassDef | RemovalSentinel:
        self.level -= 1
        self.potential_methods = set()
        if self.is_killable_class:
            self.is_killable_class = False
            return RemovalSentinel.REMOVE
        return updated_node

    def visit_FunctionDef(self, node: FunctionDef) -> bool:
        return False

    def leave_FunctionDef(self, original_node: FunctionDef, updated_node: FunctionDef) -> FunctionDef | RemovalSentinel:
        if self.level == 0:
            if original_node.name.value in self.hit_list_names:
                return RemovalSentinel.REMOVE
            return updated_node
        if original_node.name.value in self.potential_methods:
            self.is_killable_class = True
            self.potential_methods = set()
        return updated_node


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

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

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

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


def delete_top_def_nodes(module: Module, deletable_list: list[str]) -> Module:
    # First apply the TopDefTerminator
    module = module.visit(TopDefTerminator(deletable_list))
    module = module.visit(LeadingWhitespaceRemover())
    return module
