import ir
from irmach import AbstractInstruction as makeIns
from .basetarget import Label, Comment, Alignment, LabelRef, DebugInfo, Nop
from .instructionselector import InstructionSelector
from .arminstructions import Orr, Lsl, Str2, Ldr2, Ldr3, B, Bl, Bgt, Blt, Beq
from .arminstructions import Mov2, Mov3
from .arminstructions import Add, Sub, Cmp
from .basetarget import Imm8, Imm7, Imm3


class ArmInstructionSelector(InstructionSelector):
    """ Instruction selector for the arm architecture """
    def munchExpr(self, e):
        if isinstance(e, ir.Alloc):
            return 0
        elif isinstance(e, ir.Binop) and e.operation == '+' and \
            isinstance(e.b, ir.Const) and e.b.value < 8:
            a = self.munchExpr(e.a)
            d = self.newTmp()
            c = Imm3(e.b.value)
            self.emit(arm.addregregimm3_ins, others=[c], dst=[d], src=[a])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '+':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.emit(Add, dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '-' and \
            isinstance(e.b, ir.Const) and e.b.value < 8:
            a = self.munchExpr(e.a)
            d = self.newTmp()
            c = Imm3(e.b.value)
            self.emit(arm.subregregimm3_ins, others=[c], dst=[d], src=[a])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '-':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.emit(Sub, dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '|':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            self.emit(Orr, dst=[], src=[b, d])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '<<':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            self.emit(Lsl, dst=[], src=[b, d]) # TODO: is d a source variable?
            return d
        elif isinstance(e, ir.Binop) and e.operation == '*':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            # this mul instruction has operands swapped:
            self.emit(arm.mulregreg_ins, dst=[d], src=[b, d])
            return d
        elif isinstance(e, ir.Const) and e.value < 256:
            d = self.newTmp()
            self.emit(Mov3, others=[Imm8(e.value)], dst=[d])
            return d
        elif isinstance(e, ir.Const) and e.value < (2**31):
            d = self.newTmp()
            ln = LabelRef(self.frame.addConstant(e.value))
            self.emit(Ldr3, others=[ln], dst=[d])
            return d
        elif isinstance(e, ir.Mem) and isinstance(e.e, ir.Binop) and \
                e.e.operation == '+' and isinstance(e.e.b, ir.Const):
            base = self.munchExpr(e.e.a)
            d = self.newTmp()
            c = e.e.b.value
            self.emit(Ldr2, others=[c], src=[base], dst=[d])
            return d
        elif isinstance(e, ir.Mem):
            # Load from memory
            base = self.munchExpr(e.e)
            d = self.newTmp()
            self.emit(Ldr2, others=[0], src=[base], dst=[d])
            return d
        elif isinstance(e, ir.Temp):
            return e
        elif isinstance(e, ir.Call):
            # Move arguments into proper locations:
            reguses = []
            for i, a in enumerate(e.arguments):
                loc = self.frame.argLoc(i)
                m = ir.Move(loc, a)
                self.munchStm(m)
                if isinstance(loc, ir.Temp):
                    reguses.append(loc)
            self.emit(Bl(LabelRef(e.f.name)), src=reguses, dst=[self.frame.rv])
            d = self.newTmp()
            self.move(d, self.frame.rv)
            return d
        else:
            raise NotImplementedError('Expr --> {}'.format(e))

    def munchStm(self, s):
        if isinstance(s, ir.Terminator):
            pass
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Mem) and \
            isinstance(s.dst.e, ir.Binop) and s.dst.e.operation == '+' and \
            isinstance(s.dst.e.b, ir.Const):
            a = self.munchExpr(s.dst.e.a)
            val = self.munchExpr(s.src)
            c = s.dst.e.b.value
            self.emit(Str2, others=[c], src=[a, val])
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Mem):
            memloc = self.munchExpr(s.dst.e)
            val = self.munchExpr(s.src)
            self.emit(Str2, others=[0], src=[memloc, val])
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Temp):
            val = self.munchExpr(s.src)
            dreg = s.dst
            self.move(dreg, val)
        elif isinstance(s, ir.Exp):
            # Generate expression code and discard the result.
            x = self.munchExpr(s.e)
            self.emit(Nop(), src=[x])
        elif isinstance(s, ir.Jump):
            tgt = self.targets[s.target]
            self.emit(B(LabelRef(s.target.name)), jumps=[tgt])
        elif isinstance(s, ir.CJump):
            a = self.munchExpr(s.a)
            b = self.munchExpr(s.b)
            self.emit(Cmp, src=[a, b])
            ntgt = self.targets[s.lab_no]
            ytgt = self.targets[s.lab_yes]
            jmp_ins = makeIns(B(LabelRef(s.lab_no.name)), jumps=[ntgt])
            opnames = {'<': Blt, '>':Bgt, '==':Beq}
            op = opnames[s.cond](LabelRef(s.lab_yes.name))
            self.emit(op, jumps=[ytgt, jmp_ins])  # Explicitely add fallthrough
            self.emit2(jmp_ins)
        else:
            raise NotImplementedError('Stmt --> {}'.format(s))

    def move(self, dst, src):
        self.emit(Mov2, src=[src], dst=[dst])