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

import compiler
import syntax


# -------------------------------------------------
# Global function environment
# -------------------------------------------------

FUNCTIONS = {}  # name -> (func_start, func_end)


# -------------------------------------------------
# Expressions — NO CFG NODES
# -------------------------------------------------

class CONST(compiler.CONST):
    def cfa(self, pred, end):
        return pred


class ID(compiler.ID):
    def cfa(self, pred, end):
        return pred


class AOP(compiler.AOP):
    def cfa(self, pred, end):
        return pred


class COMP(compiler.COMP):
    def cfa(self, pred, end):
        return pred


class EQOP(compiler.EQOP):
    def cfa(self, pred, end):
        return pred


class LOP(compiler.LOP):
    def cfa(self, pred, end):
        return pred


# -------------------------------------------------
# Statements
# -------------------------------------------------

class ASSIGN(compiler.ASSIGN):
    def cfa(self, pred, end):
        n = CFG_Node(self)
        pred.add_child(n)
        return n


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


class IF(compiler.IF):
    def cfa(self, pred, end):
        # decision node (condition only)
        cond = CFG_DIAMOND(self.cond)
        pred.add_child(cond)

        # explicit branch entries
        then_entry = CFG_Node()
        else_entry = CFG_Node()
        cond.add_child(then_entry)
        cond.add_child(else_entry)

        # join node
        join = CFG_Node()

        then_end = self.exp1.cfa(then_entry, end)
        else_end = self.exp2.cfa(else_entry, end)

        if then_end is not None:
            then_end.add_child(join)
        if else_end is not None:
            else_end.add_child(join)

        return join


class WHILE(compiler.WHILE):
    def cfa(self, pred, end):
        # loop condition
        cond = CFG_DIAMOND(self.cond)
        pred.add_child(cond)

        # body entry
        body_entry = CFG_Node()
        cond.add_child(body_entry)

        body_end = self.body.cfa(body_entry, end)
        if body_end is not None:
            body_end.add_child(cond)  # back-edge

        # loop exit
        after = CFG_Node()
        cond.add_child(after)

        return after


# -------------------------------------------------
# Functions / calls (interprocedural CFG)
# -------------------------------------------------

class CALL(compiler.CALL):
    def cfa(self, pred, end):
        call = CFG_CALL(self)
        pred.add_child(call)

        # continuation after return
        cont = CFG_Node()

        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]

        # call → function entry
        call.add_child(f_start)

        # function exit → continuation
        f_end.add_child(cont)

        return cont


class DECL(compiler.DECL):
    def cfa(self, pred, end):
        # function entry / exit
        f_start = CFG_Node(self)
        f_end = CFG_Node(self)

        # register function
        FUNCTIONS[self.f_name] = (f_start, f_end)

        # build function body
        body_end = self.body.cfa(f_start, f_end)
        if body_end is not None:
            body_end.add_child(f_end)

        # function declaration does not alter current control flow
        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, end)
            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 to function exit
        return None        # no fallthrough