from abc import ABC, abstractmethod
from verilog_framework.core.ast_nodes import (
    Module, Port, WireDeclaration, RegDeclaration, Assign, Always,
    Identifier, Number, BinaryOp, UnaryOp, ASTNode
)


class ASTVisitor(ABC):
    @abstractmethod
    def visit_module(self, node: Module):
        pass

    @abstractmethod
    def visit_port(self, node: Port):
        pass

    @abstractmethod
    def visit_wire_declaration(self, node: WireDeclaration):
        pass

    @abstractmethod
    def visit_reg_declaration(self, node: RegDeclaration):
        pass

    @abstractmethod
    def visit_assign(self, node: Assign):
        pass

    @abstractmethod
    def visit_always(self, node: Always):
        pass

    @abstractmethod
    def visit_identifier(self, node: Identifier):
        pass

    @abstractmethod
    def visit_number(self, node: Number):
        pass

    @abstractmethod
    def visit_binary_op(self, node: BinaryOp):
        pass

    @abstractmethod
    def visit_unary_op(self, node: UnaryOp):
        pass


class ASTTraverser:
    def __init__(self, visitor: ASTVisitor):
        self.visitor = visitor

    def traverse(self, node: ASTNode):
        if isinstance(node, Module):
            self.visitor.visit_module(node)
            for port in node.port_list:
                self.traverse(port)
            for decl in node.declarations:
                self.traverse(decl)
            for stmt in node.statements:
                self.traverse(stmt)
        elif isinstance(node, Port):
            self.visitor.visit_port(node)
        elif isinstance(node, WireDeclaration):
            self.visitor.visit_wire_declaration(node)
        elif isinstance(node, RegDeclaration):
            self.visitor.visit_reg_declaration(node)
        elif isinstance(node, Assign):
            self.visitor.visit_assign(node)
            self.traverse(node.left)
            self.traverse(node.right)
        elif isinstance(node, Always):
            self.visitor.visit_always(node)
            for stmt in node.body:
                self.traverse(stmt)
        elif isinstance(node, Identifier):
            self.visitor.visit_identifier(node)
        elif isinstance(node, Number):
            self.visitor.visit_number(node)
        elif isinstance(node, BinaryOp):
            self.visitor.visit_binary_op(node)
            self.traverse(node.left)
            self.traverse(node.right)
        elif isinstance(node, UnaryOp):
            self.visitor.visit_unary_op(node)
            self.traverse(node.operand)
