from typing import List, Tuple
from MiniCVisitor import MiniCVisitor
from MiniCParser import MiniCParser
from Lib.LinearCode import LinearCode
from Lib import RiscV
from Lib.RiscV import Condition
from Lib import Operands
from antlr4.tree.Trees import Trees
from Lib.Errors import MiniCInternalError, MiniCUnsupportedError

"""
CAP, MIF08, three-address code generation + simple alloc
This visitor constructs an object of type "LinearCode".
"""


class MiniCCodeGen3AVisitor(MiniCVisitor):

    _current_function: LinearCode

    def __init__(self, debug, parser):
        super().__init__()
        self._parser = parser
        self._debug = debug
        self._functions = []
        self._lastlabel = ""

    def get_functions(self) -> List[LinearCode]:
        return self._functions

    def printSymbolTable(self):  # pragma: no cover
        print("--variables to temporaries map--")
        for keys, values in self._symbol_table.items():
            print(keys + '-->' + str(values))

    # handle variable decl

    def visitVarDecl(self, ctx) -> None:
        type_str = ctx.typee().getText()
        vars_l = self.visit(ctx.id_l())
        for name in vars_l:
            if name in self._symbol_table:
                raise MiniCInternalError(
                    "Variable {} has already been declared".format(name))
            else:
                tmp = self._current_function.fdata.fresh_tmp()
                self._symbol_table[name] = tmp
                if type_str not in ("int", "bool"):
                    raise MiniCUnsupportedError("Unsupported type " + type_str)
                # Initialization to 0 or False, both represented with 0
                self._current_function.add_instruction(
                    RiscV.li(tmp, Operands.Immediate(0)))

    def visitIdList(self, ctx) -> Operands.Temporary:
        t = self.visit(ctx.id_l())
        t.append(ctx.ID().getText())
        return t

    def visitIdListBase(self, ctx) -> List[str]:
        return [ctx.ID().getText()]

    # expressions

    def visitParExpr(self, ctx) -> Operands.Temporary:
        return self.visit(ctx.expr())

    def visitIntAtom(self, ctx) -> Operands.Temporary:
        val = Operands.Immediate(int(ctx.getText()))
        dest_temp = self._current_function.fdata.fresh_tmp()
        self._current_function.add_instruction(RiscV.li(dest_temp, val))
        return dest_temp

    def visitFloatAtom(self, ctx) -> Operands.Temporary:
        raise MiniCUnsupportedError("float literal")

    def visitBooleanAtom(self, ctx) -> Operands.Temporary:
        # true is 1 false is 0
        dtemp = self._current_function.fdata.fresh_tmp()
        if(ctx.getText() == "true"):
            val = Operands.Immediate(1)
        else:
            val = Operands.Immediate(0)
        self._current_function.add_instruction(RiscV.li(dtemp,val))
        return dtemp

    def visitIdAtom(self, ctx) -> Operands.Temporary:
        try:
            # get the temporary associated to id
            return self._symbol_table[ctx.getText()]
        except KeyError:  # pragma: no cover
            raise MiniCInternalError(
                "Undefined variable {}, this should have failed to typecheck."
                .format(ctx.getText())
            )

    def visitStringAtom(self, ctx) -> Operands.Temporary:
        raise MiniCUnsupportedError("string atom")

    # now visit expressions

    def visitAtomExpr(self, ctx) -> Operands.Temporary:
        return self.visit(ctx.atom())

    def visitAdditiveExpr(self, ctx) -> Operands.Temporary:
        assert ctx.myop is not None
        if self._debug:
                print("additive expression, between:",
                    Trees.toStringTree(ctx.expr(0), None, self._parser),
                    "and",
                    Trees.toStringTree(ctx.expr(1), None, self._parser))
        ltemp = self.visit(ctx.expr(0))
        rtemp = self.visit(ctx.expr(1))
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        if(ctx.myop.type==MiniCParser.PLUS):
            self._current_function.add_instruction(RiscV.add(dtemp,ltemp,rtemp))
        elif(ctx.myop.type==MiniCParser.MINUS):
            self._current_function.add_instruction(RiscV.sub(dtemp,ltemp,rtemp))
        else:
            raise MiniCInternalError("Unknown additive operator from the parser:",ctx.myop)
        return dtemp

    def visitOrExpr(self, ctx) -> Operands.Temporary:
        if self._debug:
                print("or expression, between:",
                    Trees.toStringTree(ctx.expr(0), None, self._parser),
                    "and",
                    Trees.toStringTree(ctx.expr(1), None, self._parser))
        ltemp = self.visit(ctx.expr(0))
        rtemp = self.visit(ctx.expr(1))
        # Getting a fresh temporary for the result of the opreation
        # We could do only two instructions with slt
        d0temp = self._current_function.fdata.fresh_tmp()
        d1temp = self._current_function.fdata.fresh_tmp()
        self._current_function.add_instruction(RiscV.add(d0temp,ltemp,rtemp))
        self._current_function.add_instruction(RiscV.mul(d1temp,ltemp,rtemp))
        self._current_function.add_instruction(RiscV.sub(d0temp,d0temp,d1temp))
        return d0temp

    def visitAndExpr(self, ctx) -> Operands.Temporary:
        if self._debug:
                print("or expression, between:",
                    Trees.toStringTree(ctx.expr(0), None, self._parser),
                    "and",
                    Trees.toStringTree(ctx.expr(1), None, self._parser))
        ltemp = self.visit(ctx.expr(0))
        rtemp = self.visit(ctx.expr(1))
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        self._current_function.add_instruction(RiscV.mul(dtemp,ltemp,rtemp))
        return dtemp

    def visitEqualityExpr(self, ctx) -> Operands.Temporary:
        return self.visitRelationalExpr(ctx)

    def visitRelationalExpr(self, ctx) -> Operands.Temporary:
        assert ctx.myop is not None
        c = Condition(ctx.myop.type)
        if self._debug:
            print("relational expression:")
            print(Trees.toStringTree(ctx, None, self._parser))
            print("Condition:", c)
        ltemp = self.visit(ctx.expr(0))
        rtemp = self.visit(ctx.expr(1))
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        endlabel = self._current_function.fdata.fresh_label("ifverified")
        
        self._current_function.add_instruction(RiscV.li(dtemp,Operands.Immediate(1)))
        self._current_function.add_comment("If the result of the comparison is true, branch")
        self._current_function.add_instruction(RiscV.conditional_jump(endlabel,ltemp,Operands.Condition(ctx.myop.type),rtemp))
        self._current_function.add_instruction(RiscV.li(dtemp,Operands.Immediate(0)))
        self._current_function.add_instruction(RiscV.jump(endlabel))
        self._current_function.add_label(endlabel)
        return dtemp

    def visitMultiplicativeExpr(self, ctx) -> Operands.Temporary:
        assert ctx.myop is not None
        div_by_zero_lbl = self._current_function.fdata.get_label_div_by_zero()
        
        ltemp = self.visit(ctx.expr(0))
        rtemp = self.visit(ctx.expr(1))
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        
        if(ctx.myop.type==MiniCParser.MULT):
            self._current_function.add_instruction(RiscV.mul(dtemp,ltemp,rtemp))
        else:
            self._current_function.add_instruction(RiscV.conditional_jump(div_by_zero_lbl,rtemp,Operands.Condition('beq'),Operands.ZERO))
            if(ctx.myop.type==MiniCParser.DIV):
                self._current_function.add_instruction(RiscV.div(dtemp,ltemp,rtemp))
            elif(ctx.myop.type==MiniCParser.MOD):
                self._current_function.add_instruction(RiscV.rem(dtemp,ltemp,rtemp))
            else:
                raise MiniCInternalError("Unknown multiplicative operator from the parser:",ctx.myop)
        return dtemp

    def visitNotExpr(self, ctx) -> Operands.Temporary:
        if self._debug:
            print("unitary not expression on expression:",
                    Trees.toStringTree(ctx.expr(), None, self._parser))
        vtemp = self.visit(ctx.expr())
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        
        # (not a) is (1 xor a)
        self._current_function.add_instruction(RiscV.li(dtemp,Operands.Immediate(1)))
        self._current_function.add_instruction(RiscV.xor(dtemp,dtemp,vtemp))
        return dtemp

    def visitUnaryMinusExpr(self, ctx) -> Operands.Temporary:
        if self._debug:
            print("unitary minus expression on expression:",
                    Trees.toStringTree(ctx.expr(), None, self._parser))
        vtemp = self.visit(ctx.expr())
        # Getting a fresh temporary for the result of the opreation
        dtemp = self._current_function.fdata.fresh_tmp()
        self._current_function.add_instruction(RiscV.sub(dtemp,Operands.ZERO,vtemp))
        return dtemp

    def visitProgRule(self, ctx) -> None:
        self.visitChildren(ctx)

    def visitFuncDef(self, ctx) -> None:
        funcname = ctx.ID().getText()
        self._current_function = LinearCode(funcname)
        self._symbol_table = dict()

        self.visit(ctx.vardecl_l())
        self.visit(ctx.block())
        self._current_function.add_comment("Return at end of function:")
        # This skeleton doesn't deal properly with functions, and
        # hardcodes a "return 0;" at the end of function. Generate
        # code for this "return 0;".
        self._current_function.add_instruction(
        	RiscV.li(Operands.A0, Operands.Immediate(0)))
        self._functions.append(self._current_function)
        del self._current_function

    def visitAssignStat(self, ctx) -> None:
        if self._debug:
            print("assign statement, rightexpression is:")
            print(Trees.toStringTree(ctx.expr(), None, self._parser))
        expr_temp = self.visit(ctx.expr())
        name = ctx.ID().getText()
        self._current_function.add_instruction(RiscV.mv(self._symbol_table[name], expr_temp))

    def visitIfStat(self, ctx) -> None:
        if self._debug:
            print("if statement")
        
        lendif = self._current_function.fdata.fresh_label("endif")
        if(ctx.else_block!=None):
            lelse = self._current_function.fdata.fresh_label("else")
            dval = self.visit(ctx.expr())
            self._current_function.add_instruction(RiscV.conditional_jump(lelse, dval, Operands.Condition('beq'), Operands.ZERO))
            self.visit(ctx.then_block)
            self._current_function.add_instruction(RiscV.jump(lendif))
            self._current_function.add_label(lelse)
            self.visit(ctx.else_block)
            self._current_function.add_instruction(RiscV.jump(lendif))
            self._current_function.add_label(lendif)
        else:
            dval = self.visit(ctx.expr())
            self._current_function.add_instruction(RiscV.conditional_jump(lendif, dval, Operands.Condition('beq'), Operands.ZERO))
            self.visit(ctx.then_block)
            self._current_function.add_instruction(RiscV.jump(lendif))
            self._current_function.add_label(lendif)
            

    def visitWhileStat(self, ctx) -> None:
        if self._debug:
            print("while statement, condition is:")
            print(Trees.toStringTree(ctx.expr(), None, self._parser))
            print("and block is:")
            print(Trees.toStringTree(ctx.stat_block(), None, self._parser))
        ltest = self._current_function.fdata.fresh_label("testcond")
        lendwhile = self._current_function.fdata.fresh_label("endwhile")
        self._current_function.add_instruction(RiscV.jump(ltest))
        self._current_function.add_label(ltest)
        dcond = self.visit(ctx.expr())
        self._current_function.add_instruction(RiscV.conditional_jump(lendwhile, dcond, Operands.Condition('beq'), Operands.ZERO))
        self.visit(ctx.body)
        self._current_function.add_instruction(RiscV.jump(ltest))
        self._current_function.add_label(lendwhile)
            
    def visitForStat(self, ctx):
        init_stat = ctx.init_stat
        cond = ctx.cond
        loop_stat = ctx.loop_stat
        body = ctx.stat_block()
        
        ltest = self._current_function.fdata.fresh_label("testcond")
        lendfor = self._current_function.fdata.fresh_label("endfor")
        
        if(init_stat != None):
            self.visit(init_stat)
        self._current_function.add_instruction(RiscV.jump(ltest))
        self._current_function.add_label(ltest)
        dcond = self.visit(ctx.expr())
        self._current_function.add_instruction(RiscV.conditional_jump(lendfor, dcond, Operands.Condition('beq'), Operands.ZERO))
        self.visit(body)
        if(loop_stat != None):
            self.visit(loop_stat)
        self._current_function.add_instruction(RiscV.jump(ltest))
        self._current_function.add_label(lendfor)
    
    # visit statements

    def visitPrintlnintStat(self, ctx) -> None:
        expr_loc = self.visit(ctx.expr())
        if self._debug:
            print("print_int statement, expression is:")
            print(Trees.toStringTree(ctx.expr(), None, self._parser))
        self._current_function.add_instruction_PRINTLN_INT(expr_loc)

    def visitPrintlnboolStat(self, ctx) -> None:
        expr_loc = self.visit(ctx.expr())
        self._current_function.add_instruction_PRINTLN_INT(expr_loc)

    def visitPrintlnfloatStat(self, ctx) -> None:
        raise MiniCUnsupportedError("Unsupported type float")

    def visitPrintlnstringStat(self, ctx) -> None:
        raise MiniCUnsupportedError("Unsupported type string")

    def visitStatList(self, ctx) -> None:
        for stat in ctx.stat():
            self._current_function.add_comment(Trees.toStringTree(stat, None, self._parser))
            self.visit(stat)
