from cfg.CFG_Node import (
    CFG_Node,
    CFG_CALL,
    CFG_RETURN,
    CFG_DIAMOND,
)

import compiler
import syntax

FUNCTIONS = {}

class CONST(compiler.CONST):
    def cfa(self, pred, end):
        n = CFG_Node(self)
        pred.add_child(n)
        n.add_child(end) if end else None
        return n

class ID(compiler.ID):
    def cfa(self, pred, end):
        n = CFG_Node(self)
        pred.add_child(n)
        n.add_child(end) if end else None
        return n

class AOP(compiler.AOP):
    def cfa(self, pred, end):
        left_node = self.arg1.cfa(pred, None)
        right_node = self.arg2.cfa(left_node, None)
        op_node = CFG_Node(self)
        right_node.add_child(op_node)
        op_node.add_child(end) if end else None
        return op_node

class COMP(compiler.COMP):
    def cfa(self, pred, end):
        left_node = self.arg1.cfa(pred, None)
        right_node = self.arg2.cfa(left_node, None)
        comp_node = CFG_Node(self.operator)
        right_node.add_child(comp_node)
        comp_node.add_child(end) if end else None
        return comp_node

class EQOP(compiler.EQOP):
    def cfa(self, pred, end):
        left_node = self.arg1.cfa(pred, None)
        right_node = self.arg2.cfa(left_node, None)
        eqop_node = CFG_Node(self)
        right_node.add_child(eqop_node)
        eqop_node.add_child(end) if end else None
        return eqop_node

class LOP(compiler.LOP):
    def cfa(self, pred, end):
        n = CFG_Node(self)
        pred.add_child(n)
        n.add_child(end) if end else None
        return n

class ASSIGN(compiler.ASSIGN):
    def cfa(self, pred, end):
        expr_node = self.expr.cfa(pred, None)
        assign_node = CFG_Node(self)
        expr_node.add_child(assign_node)
        assign_node.add_child(end) if end else None
        return assign_node

class SEQ(compiler.SEQ):
    def cfa(self, pred, end):
        mid = self.exp1.cfa(pred, None)
        if mid is None:
            return None
        return self.exp2.cfa(mid, end)

class IF(compiler.IF):
    def cfa(self, pred, end):
        cond_node = self.cond.cfa(pred, None)
        diamond = CFG_DIAMOND(self.cond)
        cond_node.add_child(diamond)
        then_entry = CFG_Node()
        else_entry = CFG_Node()
        diamond.add_child(then_entry)
        diamond.add_child(else_entry)
        join = CFG_Node()
        join.add_child(end) if end else None
        then_end = self.exp1.cfa(then_entry, join)
        else_end = self.exp2.cfa(else_entry, join)
        return join

class WHILE(compiler.WHILE):
    def cfa(self, pred, end):
        cond_node = self.cond.cfa(pred, None)
        diamond = CFG_DIAMOND(self.cond)
        cond_node.add_child(diamond)
        body_entry = CFG_Node()
        diamond.add_child(body_entry)
        body_end = self.body.cfa(body_entry, diamond)
        if body_end is not None:
            body_end.add_child(diamond)
        after = CFG_Node()
        diamond.add_child(after)
        after.add_child(end) if end else None
        return after

class CALL(compiler.CALL):
    def cfa(self, pred, end):
        call_node = CFG_Node(self)
        call_node.label = f"START {self.f_name}({', '.join(map(str, self.arg))})"
        pred.add_child(call_node)

        cont = CFG_Node()
        cont.add_child(end) if end else None

        if self.f_name not in FUNCTIONS:
            raise RuntimeError(f"Call to undefined function '{self.f_name}'")

        f_start, f_end = FUNCTIONS[self.f_name]

        # Create return node from function
        return_node = CFG_Node(self)
        return_node.label = f"END {self.f_name}({', '.join(map(str, self.arg))})"
        f_end.add_child(return_node)
        return_node.add_child(cont)

        call_node.add_child(f_start)
        return cont

class DECL(compiler.DECL):
    def cfa(self, pred, end):
        f_start = CFG_Node(self)
        f_start.label = f"START {self.f_name}({', '.join(self.params)})"
        f_end = CFG_Node(self)
        f_end.label = f"END {self.f_name}({', '.join(self.params)})"

        FUNCTIONS[self.f_name] = (f_start, f_end)
        body_end = self.body.cfa(f_start, f_end)
        if body_end is not None:
            body_end.add_child(f_end)
        return pred

class LET(compiler.LET):
    def cfa(self, pred, end):
        current = pred
        decls = self.decl if isinstance(self.decl, list) else [self.decl]
        for d in decls:
            current = d.cfa(current, None)
            if current is None:
                return None
        return self.body.cfa(current, end)

class RETURN(syntax.EXPRESSION):
    def cfa(self, pred, end):
        n = CFG_RETURN(self)
        pred.add_child(n)
        n.add_child(end)
        return None