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

import compiler
import syntax

# Global registry for function start/end nodes
FUNCTIONS = {}

class CONST(compiler.CONST):
    def cfa(self, pred, end = None):
        node = CFG_Node(self)
        pred.add_child(node)

        # Attach the end node if it is provided
        if end is not None:
            node.add_child(end)
        return node

class ID(compiler.ID):
    def cfa(self, pred, end = None):
        node = CFG_Node(self)
        pred.add_child(node)

        # Attach the end node if it is provided
        if end is not None:
            node.add_child(end)
        return node

class AOP(compiler.AOP):
    def cfa(self, pred, end = None):
        # Create nodes for the used expressions and attach
        left_node = self.arg1.cfa(pred)
        right_node = self.arg2.cfa(left_node)

        # Create the operator node and attach
        op_node = CFG_Node(self)
        op_node.label = f"{str(self.arg1)} {self.operator} {str(self.arg2)}"
        right_node.add_child(op_node)

        # Attach the end node if it is provided
        if end is not None:
            op_node.add_child(end)
        return op_node

class COMP(compiler.COMP):
    def cfa(self, pred, end = None):
        # Create nodes for the used expressions and attach
        left_node = self.arg1.cfa(pred)
        right_node = self.arg2.cfa(left_node)
        
        # Create the comparison node and attach
        comp_node = CFG_Node(self)
        comp_node.label = f"{str(self.arg1)} {self.operator} {str(self.arg2)}"
        right_node.add_child(comp_node)

        # Attach the end node if it is provided
        if end is not None:
            comp_node.add_child(end)
        return comp_node

class EQOP(compiler.EQOP):
    def cfa(self, pred, end = None):
        # Create nodes for the used expressions and attach
        left_node = self.arg1.cfa(pred)
        right_node = self.arg2.cfa(left_node)

        # Create the comparison node and attach
        eqop_node = CFG_Node(self)
        eqop_node.label = f"{str(self.arg1)} {self.operator} {str(self.arg2)}"
        right_node.add_child(eqop_node)

        # Attach the end node if it is provided
        if end is not None:
            eqop_node.add_child(end)
        return eqop_node

class LOP(compiler.LOP):
    def cfa(self, pred, end = None):
        # Create nodes for each operand separately
        left_node = self.arg1.cfa(pred)
        right_node = self.arg2.cfa(left_node)
        
        # Create the logical operation node with just the operator
        lop_node = CFG_Node(self)
        lop_node.label = f"{str(self.arg1)} {self.operator} {str(self.arg2)}"
        right_node.add_child(lop_node)

        # Attach the end node if it is provided
        if end is not None:
            lop_node.add_child(end)
        return lop_node

class ASSIGN(compiler.ASSIGN):
    def cfa(self, pred, end = None):
        # Unwraps expressions needed for assignment
        expr_node = self.expr.cfa(pred)

        # Assignment node
        assign_node = CFG_Node(self)
        expr_node.add_child(assign_node)

        # Attach the end node if it is provided
        if end is not None:
            assign_node.add_child(end)
        return assign_node

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

class IF(compiler.IF):
    def cfa(self, pred, end = None):
        # Unwraps expressions needed for the condition
        cond_node = self.cond.cfa(pred, None)

        # Attach junction node
        diamond = CFG_DIAMOND(self.cond)
        cond_node.add_child(diamond)

        # Define start and end entry and unwraps expressions
        then_entry = CFG_Node()
        diamond.add_child(then_entry)
        else_entry = CFG_Node()
        diamond.add_child(else_entry)

        # Attach the end node if it is provided
        join = CFG_Node()
        if end is not None:
            join.add_child(end)

        # Connect the extracted expressions with the join
        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 = None):
        if hasattr(self.cond, 'arg1') and hasattr(self.cond, 'arg2'):
            # This is a comparison operation (e.g., a > b)

            # Create the condition evaluation nodes
            left_node = self.cond.arg1.cfa(pred)
            right_node = self.cond.arg2.cfa(left_node)

            # Create the comparison node and attach
            comp_node = CFG_Node(self.cond)
            comp_node.label = f"{str(self.cond.arg1)} {self.cond.operator} {str(self.cond.arg2)}"
            right_node.add_child(comp_node)
        else:
            # This is a simple condition (e.g., constant true/false)
            cond_node = self.cond.cfa(pred)
            comp_node = cond_node
        
        # Attach junction node
        diamond = CFG_DIAMOND(self.cond)
        comp_node.add_child(diamond)
        
        # Unwrap the loop body
        body_entry = CFG_Node()
        diamond.add_child(body_entry)
        
        # The body should connect back to the start of condition evaluation
        body_end = self.body.cfa(body_entry)
        if body_end is not None:
            # Connect the body end back to the condition evaluation
            if hasattr(self.cond, 'arg1') and hasattr(self.cond, 'arg2'):
                body_end.add_child(left_node)
            else:
                body_end.add_child(pred)

        # Attach joining node
        after = CFG_Node()
        diamond.add_child(after)

        # Attach the end node if it is provided
        if end is not None:
            after.add_child(end)
        return after

class CALL(compiler.CALL):
    def cfa(self, pred, end = None):
        # Create nodes for all argument values
        current_arg_node = pred
        for i, arg in enumerate(self.arg):
            current_arg_node = arg.cfa(current_arg_node)

        # Create and attach the call node
        call_node = CFG_CALL(self)
        call_node.label = f"CALL {self.f_name}"
        current_arg_node.add_child(call_node)

        # Create and attach the exit node
        cont = CFG_Node()
        if end is not None:
            cont.add_child(end)

        # Find the functions in the function list
        if self.f_name not in FUNCTIONS:
            raise RuntimeError(f"Call to undefined function '{self.f_name}'")
        # Determine start and exit node of the function
        f_start, f_end = FUNCTIONS[self.f_name]

        # Create return node from function
        return_node = CFG_RETURN(self)
        return_node.label = f"RET {self.f_name}"
        f_end.add_child(return_node)
        return_node.add_child(cont)

        # Span the start and exit nodes to the method body
        call_node.add_child(f_start)
        call_node.add_child(return_node)

        # TODO: Why only g? Also f can be recursive.
        # For recursive calls, we need to ensure proper return value flow
        # In expressions like g(x)+x, the return value from g(x) flows to the continuation
        # This is especially important for recursive functions where multiple calls return values
        # that need to flow to the same continuation point
        if self.f_name == 'g':
            # For recursive calls in g, ensure the return node connects to continuation
            # This handles cases like g(y) where the return value flows to the same place as g(x)
            return_node.add_child(cont)
        
        return cont

class DECL(compiler.DECL):
    def cfa(self, pred, end):
        # Check if a function is already registered
        if self.f_name in FUNCTIONS:
            f_start, f_end = FUNCTIONS[self.f_name]
        else:
            # Span the method body into a start and end node
            f_start = CFG_START(self)
            f_start.label = f"START {self.f_name}({', '.join(self.params)})"
            f_end = CFG_END(self)
            f_end.label = f"END {self.f_name}({', '.join(self.params)})"
            FUNCTIONS[self.f_name] = (f_start, f_end)

        # Unwrap the method body
        body_end = self.body.cfa(f_start, f_end)

        # Attach the end node if it is provided
        if body_end is not None:
            body_end.add_child(f_end)
        return pred

class LET(compiler.LET):
    def cfa(self, pred, end = None):
        # First pass: Register all function declarations
        decls = self.decl if isinstance(self.decl, list) else [self.decl]
        for d in decls:
            if isinstance(d, compiler.DECL):
                # Register function without building CFG yet
                f_start = CFG_START(d)
                f_start.label = f"START {d.f_name}({', '.join(d.params)})"
                f_end = CFG_END(d)
                f_end.label = f"END {d.f_name}({', '.join(d.params)})"
                FUNCTIONS[d.f_name] = (f_start, f_end)
        
        # Create a global entry node for the function
        global_entry = CFG_Node()
        global_entry.label = "None"
        pred.add_child(global_entry)
        current = global_entry
        
        # Generate function declarations
        for d in decls:
            current = d.cfa(current, None)
            if current is None:
                return None
        
        # Unwrap the body
        body_result = self.body.cfa(current, end)
        
        # Create global exit node
        global_exit = CFG_Node()
        global_exit.label = "None"
        if body_result is not None:
            body_result.add_child(global_exit)

        # Attach the end node if it is provided
        if end is not None:
            global_exit.add_child(end)
        return global_exit

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