import libcst as cst


class EnumCollector(cst.CSTVisitor):
    def __init__(self, module: cst.Module):
        self.enum_collection: dict[str, str] = {}
        self.module = module

    def visit_ClassDef(self, node: cst.ClassDef) -> None:
        for base in node.bases:
            base_value = base.value
            if isinstance(base_value, cst.Name) and base_value.value == "Enum":
                self.enum_collection[node.name.value] = self.module.code_for_node(node)
            elif isinstance(base_value, cst.Attribute):
                if (
                    isinstance(base_value.value, cst.Name)
                    and base_value.value.value == "enum"
                    and base_value.attr.value == "Enum"
                ):
                    self.enum_collection[node.name.value] = self.module.code_for_node(node)
