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}")

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

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

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

class AOP(syntax.AOP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)
        op = self.operator
        if op == "+":
            return c1 + c2 + [add()]
        if op == "-":
            return c1 + c2 + [sub()]
        if op == "*":
            return c1 + c2 + [mul()]
        return c1 + c2 + [div()]

class COMP(syntax.COMP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)
        if self.operator == "<":
            return c1 + c2 + [lt()]
        if self.operator == ">":
            return c1 + c2 + [gt()]
        if self.operator == "<=":
            return c1 + c2 + [gt(), nop()]  # NOT IMPLEMENTED IN TRAM
        if self.operator == ">=":
            return c1 + c2 + [lt(), nop()]  # NOT IMPLEMENTED IN TRAM
        return c1 + c2

class EQOP(syntax.EQOP):
    def code(self, rho, nl):
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)
        if self.operator == "==":
            return c1 + c2 + [eq()]
        return c1 + c2 + [neq()]

class LOP(syntax.LOP):
    def code(self, rho, nl):
        # TRIPLA: logical operators must be done by short-circuiting (not strict)
        # but your TRAM has no boolean ops, so fallback to == and pop
        c1 = self.arg1.code(rho, nl)
        c2 = self.arg2.code(rho, nl)
        if self.operator == "&&":
            l = new_label()
            return (
                c1 +
                [ifzero(l)] +
                c2 +
                [ifzero(l)] +
                [const(1)] +
                [nop(assigned_label=l)]
            )
        if self.operator == "||":
            l = new_label()
            return (
                c1 +
                [ifzero(l)] +
                [const(1)] +
                [goto(l)] +
                [nop(assigned_label=l)] +
                c2
            )
        return c1 + c2

class ASSIGN(syntax.ASSIGN):
    def code(self, rho, nl):
        c = self.expr.code(rho, nl)
        k, nlp = rho[self.var.name]
        d = nl - nlp
        return c + [store(k, d), load(k, d)]

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

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)]
        )

class WHILE(syntax.WHILE):
    def code(self, rho, nl):
        l1 = new_label()  # loop header
        l2 = new_label()  # exit
        l3 = new_label()  # repeat
        l4 = new_label()  # evaluate body
        return (
            self.cond.code(rho, nl)
            + [ifzero(l2), goto(l4)]
            + [nop(assigned_label=l1)]
            + self.cond.code(rho, nl)
            + [ifzero(l3), pop()]
            + [nop(assigned_label=l4)]
            + self.body.code(rho, nl)
            + [goto(l1)]
            + [nop(assigned_label=l2), const(None), nop(assigned_label=l3)]
        )

class LET(syntax.LET):
    def code(self, rho, nl):
        l = new_label()
        rho2 = elab_def(self.decl, rho, nl)
        cd = []
        for d in self.decl:
            cd += d.code(rho2, nl)
        ce = self.body.code(rho2, nl)
        return [goto(l)] + cd + [nop(assigned_label=l)] + ce

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

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

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