from typing import List
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) -> List[str]:
        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:
        dest = self._current_function.fdata.fresh_tmp()
        self._current_function.add_instruction(
            RiscV.li(
                dest, Operands.Immediate(1 if ctx.getText() == "true" else 0)
            )
        )
        return dest  # true is 1 false is 0

    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
        tmpl: Operands.Temporary = self.visit(ctx.expr(0))
        tmpr: Operands.Temporary = self.visit(ctx.expr(1))
        tmp = self._current_function.fdata.fresh_tmp()

        if ctx.myop.type == MiniCParser.PLUS:
            self._current_function.add_instruction(
                RiscV.add(tmp, tmpl, tmpr)
            )
        elif ctx.myop.type == MiniCParser.MINUS:
            self._current_function.add_instruction(
                RiscV.sub(tmp, tmpl, tmpr)
            )
        else:
            raise MiniCInternalError(
                f"Unknown additive operator '{ctx.myop}'")

        return tmp

    def visitOrExpr(self, ctx) -> Operands.Temporary:
        tmp = self._current_function.fdata.fresh_tmp()
        t1, t2 = self.visit(ctx.expr(0)), self.visit(ctx.expr(1))
        self._current_function.add_instruction(
            RiscV.lor(tmp, t1, t2)
        )
        return tmp

    def visitAndExpr(self, ctx) -> Operands.Temporary:
        tmp = self._current_function.fdata.fresh_tmp()
        t1, t2 = self.visit(ctx.expr(0)), self.visit(ctx.expr(1))
        self._current_function.add_instruction(
            RiscV.land(tmp, t1, t2)
        )
        return tmp

    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, [], self._parser))
            print("Condition:", c)

        dest = self._current_function.fdata.fresh_tmp()
        t1 = self.visit(ctx.expr(0))
        t2 = self.visit(ctx.expr(1))

        endrel = self._current_function.fdata.fresh_label("end_rel")

        #   dest <- 1
        #   t1 cond t2 ? jmp end
        #   dest <- 0
        # end:
        self._current_function.add_instruction(
            RiscV.li(dest, Operands.Immediate(1)))

        self._current_function.add_instruction(
            RiscV.conditional_jump(endrel, t1, Condition(ctx.myop.type), t2))
        self._current_function.add_instruction(
            RiscV.li(dest, Operands.Immediate(0)))

        self._current_function.add_label(endrel)
        return dest

    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()
        t1 = self.visit(ctx.expr(0))
        t2 = self.visit(ctx.expr(1))
        dest = self._current_function.fdata.fresh_tmp()

        if ctx.myop.type == MiniCParser.DIV or ctx.myop.type == MiniCParser.MOD:
            self._current_function.add_instruction(
                RiscV.conditional_jump(div_by_zero_lbl, t2, Condition("beq"), Operands.ZERO))

        if ctx.myop.type == MiniCParser.MULT:
            self._current_function.add_instruction(
                RiscV.mul(dest, t1, t2))
        elif ctx.myop.type == MiniCParser.DIV:
            self._current_function.add_instruction(
                RiscV.div(dest, t1, t2))
        elif ctx.myop.type == MiniCParser.MOD:
            self._current_function.add_instruction(
                RiscV.rem(dest, t1, t2))
        else:
            raise MiniCInternalError(
                f"Unknown additive operator '{ctx.myop}'")

        return dest

    def visitNotExpr(self, ctx) -> Operands.Temporary:
        dest = self._current_function.fdata.fresh_tmp()
        label_neq_zero = self._current_function.fdata.fresh_label("neq_zero")
        t1 = self.visit(ctx.expr())

        #   dest <- 0
        #   t1 != dest ? jmp neq_zero
        #   dest <- 1
        # neq_zero:
        self._current_function.add_instruction(
            RiscV.li(dest, Operands.Immediate(0))
        )
        self._current_function.add_instruction(
            RiscV.conditional_jump(label_neq_zero, t1, Condition("bne"), Operands.ZERO)
        )
        self._current_function.add_instruction(
            RiscV.li(dest, Operands.Immediate(1))
        )
        
        self._current_function.add_label(label_neq_zero)
        return dest

    def visitUnaryMinusExpr(self, ctx) -> Operands.Temporary:
        dest = self._current_function.fdata.fresh_tmp()

        self._current_function.add_instruction(
            RiscV.sub(dest, Operands.ZERO, self.visit(ctx.expr()))
        )
        return dest

    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(), [], 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")
        else_label = self._current_function.fdata.fresh_label("else_if")
        end_if_label = self._current_function.fdata.fresh_label("end_if")

        self._current_function.add_instruction(
            RiscV.conditional_jump(
                else_label, self.visit(ctx.expr()),
                Condition('beq'), Operands.ZERO
            )
        )

        self.visit(ctx.stat_block()[0])
        self._current_function.add_instruction(
            RiscV.jump(end_if_label)
        )

        self._current_function.add_label(else_label)
        if len(ctx.stat_block()) > 1:
            self.visit(ctx.stat_block()[1])
        
        self._current_function.add_label(end_if_label)

    def visitWhileStat(self, ctx) -> None:
        if self._debug:
            print("while statement, condition is:")
            print(Trees.toStringTree(ctx.expr(), [], self._parser))
            print("and block is:")
            print(Trees.toStringTree(ctx.stat_block(), [], self._parser))
        
        # ltest <- fresh_label()
        loop_test = self._current_function.fdata.fresh_label("while_test")
        # lendwhile <- fresh_label()
        end_while = self._current_function.fdata.fresh_label("end_while")
        # code.addLabel(ltest)
        self._current_function.add_label(loop_test)
        # t1 <- GenCodeExpr(b)
        # code.add("beq lendwhile, t1, 0")

        self._current_function.add_instruction(
            RiscV.conditional_jump(
                end_while, self.visit(ctx.expr()),
                Condition('beq'), Operands.ZERO
            )
        )
        # GenCodeSmt(S) # execute S
        self.visit(ctx.stat_block())
        # code.add("j ltest") # and jump to the test
        self._current_function.add_instruction(
            RiscV.jump(loop_test)
        )
        # code.addLabel(lendwhile) # else it is done.
        self._current_function.add_label(end_while)
    
    # 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(), [], 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, [], self._parser))
            self.visit(stat)

    def visitForStat(self, ctx) -> None:
        def immediate_to_tmp(op):
            if not isinstance(op, Operands.Immediate):
                return op
            new_op = self._current_function.fdata.fresh_tmp()
            self._current_function.add_instruction(
                RiscV.li(new_op, op))
            return new_op


        from_t = immediate_to_tmp(self.visit(ctx.expr(0)))
        to_t = immediate_to_tmp(self.visit(ctx.expr(1)))
        
        stride = Operands.Immediate(1)
        if len(ctx.expr()) > 2:
            stride = self.visit(ctx.expr(2))

        stride = immediate_to_tmp(stride)

        name = ctx.ID().getText()
        if name not in self._symbol_table:
            raise MiniCInternalError(
                "Variable {} has not been declared before for loop".format(name))

        variable = self._symbol_table[name] # variable acting as counter
        for_test = self._current_function.fdata.fresh_label("forF_test")
        end_for = self._current_function.fdata.fresh_label("end_forF")
        body_for = self._current_function.fdata.fresh_label("body_forF")
        neg_stride = self._current_function.fdata.fresh_label("neg_stride_forF")
        counter = self._current_function.fdata.fresh_tmp() # real counter not affected by var. changes

        #   counter <- from
        self._current_function.add_instruction(
            RiscV.mv(counter, from_t))
        #* for_test:
        #   var <- counter
        #   bge 0 stride neg_stride
        #   (positive stride here)
        #   ble counter to end_for
        #   j body
        self._current_function.add_label(for_test)
        self._current_function.add_instruction(
            RiscV.mv(variable, counter))
        self._current_function.add_instruction(
            RiscV.conditional_jump(
                neg_stride, Operands.ZERO,
                Condition(MiniCParser.GTEQ), stride
            ))
        self._current_function.add_instruction(
            RiscV.conditional_jump(
                end_for, counter,
                Condition(MiniCParser.GTEQ), to_t
            ))
        self._current_function.add_instruction(
            RiscV.jump(body_for))
        #* neg_stride:
        #   bl to counter end_for
        self._current_function.add_label(neg_stride)
        self._current_function.add_instruction(
            RiscV.conditional_jump(
                end_for, to_t,
                Condition(MiniCParser.GTEQ), counter
            ))
        #* body:
        #   <body>
        #   counter <- counter + stride
        #   j for_test
        self._current_function.add_label(body_for)
        self.visit(ctx.stat_block())
        self._current_function.add_instruction(
            RiscV.add(counter, counter, stride))
        self._current_function.add_instruction(
            RiscV.jump(for_test))
        #* end_for:
        self._current_function.add_label(end_for)