import syntax
from vistram.tram import *

label_counter = 0

def new_label():
    global label_counter
    label_counter += 1
    return Label(text=f"L{label_counter}")

# Rules E1, E2
def elab_def(decls, rho, nl):
    r = dict(rho)
    for d in decls:
        r = d.elab(r, nl)
    return r

# Rule K6
class CONST(syntax.CONST):
    def code(self, rho, nl):
        return [const(self.value)]

# Rule K3
class ID(syntax.ID):
    def code(self, rho, nl):
        k, nl_def = rho[self.name]
        return [load(k, nl - nl_def)]

# Rule K7
class AOP(syntax.AOP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)

        match self.operator:
            case "+":
                return c1 + c2 + [add()]
            case "-":
                return c1 + c2 + [sub()]
            case "*":
                return c1 + c2 + [mul()]
            case "/":
                return c1 + c2 + [div()]
            case _:
                raise Exception(f"Unknown AOP operator.")

# Rule K5
class COMP(syntax.COMP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)

        match self.operator:
            case "<":
                return c1 + c2 + [lt()]
            case ">":
                return c1 + c2 + [gt()]
            case "<=":
                # leq is not implemented in vistra.tram
                return LOP("||",  COMP("<", self.arg1, self.arg2), EQOP("==", self.arg1, self.arg2)).code(rho, nl)
            case ">=":
                # geq is not implemented in vistra.tram
                return LOP("||",  COMP(">", self.arg1, self.arg2), EQOP("==", self.arg1, self.arg2)).code(rho, nl)
            case _:
                raise ValueError(f"Unknown COMP operator.")

# Rule K5
class EQOP(syntax.EQOP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)

        match self.operator:
            case "==":
                return c1 + c2 + [eq()]
            case "!=":
                return c1 + c2 + [neq()]
            case _:
                raise ValueError(f"Unknown EQOP operator.")

# see lecture code generation AND/OR
class LOP(syntax.LOP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)

        l1 = new_label()
        l2 = new_label()
        match self.operator:
            case "&&":
                return c1 + [ifzero(l1)] + [const(1)] +  c2 +  [mul()] +  [goto(l2)] + [nop(l1)] + [const(0)] + [nop(l2)]
            case "||":
                return c1 + [ifzero(l1)] + [const(1)] + [goto(l2)] +  [nop(l1)] + c2 + [nop(l2)]
            case _:
                raise ValueError("Unknown LOP operator.")

# Rule K4
class ASSIGN(syntax.ASSIGN):
    def code(self, rho, nl):
        c = self.expr.code(rho, nl)
        k, nl_def = rho[self.var.name]
        dif = nl - nl_def
        return c + [store(k, dif), load(k, dif)]

# Rule K9
class SEQ(syntax.SEQ):
    def code(self, rho, nl):
        c1 = self.exp1.code(rho, nl)
        c2 = self.exp2.code(rho, nl)
        return c1 + [pop()] + c2

# Rule K10
class IF(syntax.IF):
    def code(self, rho, nl):
        l1 = new_label()
        l2 = new_label()

        c1 = self.cond.code(rho, nl)
        c2 = self.exp1.code(rho, nl)
        c3 = self.exp2.code(rho, nl)

        return c1 + [ifzero(l1)] + c2 + [goto(l2)] + [nop(assigned_label=l1)] + c3 + [nop(assigned_label=l2)]

# Rule K11
class WHILE(syntax.WHILE):
    def code(self, rho, nl):
        l1 = new_label()
        l2 = new_label()
        l3 = new_label()
        l4 = new_label()

        c1 = self.cond.code(rho, nl)
        c2 = self.body.code(rho, nl)

        return c1 + [ifzero(l3), goto(l4)] + [nop(l1)] + c1 + [ifzero(l2), pop()] + [nop(l4)] + c2 + [goto(l1)] + [nop(l3), const(None), nop(l2)]

# Rule K2
class LET(syntax.LET):
    def code(self, rho, nl):
        l = new_label()

        # Normalize to list
        decls = self.decl if isinstance(self.decl, list) else [self.decl]

        shared_rho = elab_def(decls, rho, nl)
        ds = []
        for d in decls:
            ds += d.code(shared_rho, nl)

        es = self.body.code(shared_rho, nl)
        return [goto(l)] + ds + [nop(assigned_label=l)] + es

# Rule K13
class DECL(syntax.DECL):
    def elab(self, rho, nl):
        nested_rho = dict(rho)
        nested_rho[self.f_name] = (new_label(), nl)
        return nested_rho

    def code(self, rho, nl):
        l, _ = rho[self.f_name]
        nested_rho = dict(rho)
        for i, p in enumerate(self.params):
            nested_rho[p] = (i, nl + 1)
        c = self.body.code(nested_rho, nl + 1)
        return [nop(assigned_label=l)] + c + [ireturn()]

# Rule K12
class CALL(syntax.CALL):
    def code(self, rho, nl):
        c = []
        for a in self.arg:
            c += a.code(rho, nl)
        l, nl_def = rho[self.f_name]
        return c + [invoke(len(self.arg), l, nl - nl_def)]