from .visitor import ASTVisitor
from verilog_framework.core.ast_nodes import (
    Module, Port, WireDeclaration, RegDeclaration, Assign, Always,
    Identifier, Number, BinaryOp, UnaryOp
)


class GraphvizVisualizer(ASTVisitor):
    def __init__(self):
        self.node_counter = 0
        self.edges = []
        self.nodes = []

    def next_node_id(self):
        node_id = self.node_counter
        self.node_counter += 1
        return f"node{node_id}"

    def add_node(self, node_id: str, label: str):
        self.nodes.append(f'{node_id} [label="{label}"];')

    def add_edge(self, from_node: str, to_node: str, label: str = ""):
        if label:
            self.edges.append(f'{from_node} -> {to_node} [label="{label}"];')
        else:
            self.edges.append(f'{from_node} -> {to_node};')

    def generate_dot(self, ast: Module) -> str:
        self.node_counter = 0
        self.edges = []
        self.nodes = []
        
        root_id = self.next_node_id()
        self.add_node(root_id, f"Module: {ast.name}")
        
        for i, port in enumerate(ast.port_list):
            port_id = self.next_node_id()
            self.add_node(port_id, f"{port.name} ({port.direction})")
            self.add_edge(root_id, port_id, "port")
        
        for i, decl in enumerate(ast.declarations):
            if isinstance(decl, WireDeclaration):
                decl_id = self.next_node_id()
                self.add_node(decl_id, f"Wire: {decl.name}")
                self.add_edge(root_id, decl_id, "decl")
            elif isinstance(decl, RegDeclaration):
                decl_id = self.next_node_id()
                self.add_node(decl_id, f"Reg: {decl.name}")
                self.add_edge(root_id, decl_id, "decl")
        
        for i, stmt in enumerate(ast.statements):
            stmt_id = self._process_statement(stmt)
            self.add_edge(root_id, stmt_id, f"stmt{i}")
        
        dot = "digraph G {\n"
        dot += "  rankdir=TB;\n"
        for node in self.nodes:
            dot += f"  {node}\n"
        for edge in self.edges:
            dot += f"  {edge}\n"
        dot += "}"
        
        return dot

    def _process_statement(self, stmt):
        if isinstance(stmt, Assign):
            stmt_id = self.next_node_id()
            self.add_node(stmt_id, "Assign")
            left_id = self._process_expression(stmt.left)
            right_id = self._process_expression(stmt.right)
            self.add_edge(stmt_id, left_id, "left")
            self.add_edge(stmt_id, right_id, "right")
            return stmt_id
        elif isinstance(stmt, Always):
            stmt_id = self.next_node_id()
            sens_label = f"Always: {','.join(stmt.sensitivity_list)}"
            self.add_node(stmt_id, sens_label)
            return stmt_id
        return None

    def _process_expression(self, expr):
        if isinstance(expr, Identifier):
            expr_id = self.next_node_id()
            self.add_node(expr_id, f"Var: {expr.name}")
            return expr_id
        elif isinstance(expr, Number):
            expr_id = self.next_node_id()
            self.add_node(expr_id, f"Num: {expr.value}")
            return expr_id
        elif isinstance(expr, BinaryOp):
            expr_id = self.next_node_id()
            self.add_node(expr_id, f"Op: {expr.operator}")
            left_id = self._process_expression(expr.left)
            right_id = self._process_expression(expr.right)
            self.add_edge(expr_id, left_id, "left")
            self.add_edge(expr_id, right_id, "right")
            return expr_id
        elif isinstance(expr, UnaryOp):
            expr_id = self.next_node_id()
            self.add_node(expr_id, f"UnOp: {expr.operator}")
            operand_id = self._process_expression(expr.operand)
            self.add_edge(expr_id, operand_id, "operand")
            return expr_id
        return None

    def visit_module(self, node: Module):
        pass

    def visit_port(self, node: Port):
        pass

    def visit_wire_declaration(self, node: WireDeclaration):
        pass

    def visit_reg_declaration(self, node: RegDeclaration):
        pass

    def visit_assign(self, node: Assign):
        pass

    def visit_always(self, node: Always):
        pass

    def visit_identifier(self, node: Identifier):
        pass

    def visit_number(self, node: Number):
        pass

    def visit_binary_op(self, node: BinaryOp):
        pass

    def visit_unary_op(self, node: UnaryOp):
        pass
