import ir
from ppci import CompilerError
from .basetarget import Label, Comment, Alignment, LabelRef, DebugInfo, Nop
from .basetarget import Imm32, Imm3
import .armtarget as arm
from .instructionselector import InstructionSelector
import irmach
from irmach import AbstractInstruction as makeIns
import asm


class ArmFrame(irmach.Frame):
    """
      Arm specific frame for functions.
    """
    def __init__(self, name):
        # We use r7 as frame pointer.
        super().__init__(name)
        self.regs = [arm.r0, arm.r1, arm.r2, arm.r3, arm.r4, arm.r5, arm.r6]
        self.rv = ir.Temp('special_RV')
        self.p1 = ir.Temp('special_P1')
        self.p2 = ir.Temp('special_P2')
        self.p3 = ir.Temp('special_P3')
        self.p4 = ir.Temp('special_P4')
        self.fp = ir.Temp('special_FP')
        # Pre-colored registers:
        self.tempMap = {}
        self.tempMap[self.rv] = arm.r0
        self.tempMap[self.p1] = arm.r1
        self.tempMap[self.p2] = arm.r2
        self.tempMap[self.p3] = arm.r3
        self.tempMap[self.p4] = arm.r4
        self.tempMap[self.fp] = arm.r7
        self.locVars = {}
        self.parMap = {}
        # Literal pool:
        self.constants = []

    def argLoc(self, pos):
        """
            Gets the function parameter location in IR-code format.
        """
        if pos == 0:
            return self.p1
        elif pos == 1:
            return self.p2
        elif pos == 2:
            return self.p3
        elif pos == 3:
            return self.p4
        else:
            raise NotImplementedError('No more than 4 parameters implemented')

    def allocVar(self, lvar):
        if lvar not in self.locVars:
            self.locVars[lvar] = self.stacksize
            self.stacksize = self.stacksize + 4
        return self.locVars[lvar]

    def addConstant(self, value):
        lab_name = '{}_literal_{}'.format(self.name, len(self.constants))
        self.constants.append((lab_name, value))
        return lab_name

    def EntryExitGlue3(self):
        """
            Add code for the prologue and the epilogue. Add a label, the
            return instruction and the stack pointer adjustment for the frame.
        """
        self.instructions.insert(0, makeIns(arm.Label(self.name)))
        self.instructions.insert(1, makeIns(arm.push_ins(arm.RegisterSet({arm.lr, arm.r7}))))
        # Reserve stack space for locals:
        self.instructions.insert(2, makeIns(arm.subspsp_ins(arm.sp, arm.sp, arm.Imm7(self.stacksize))))
        # Setup frame pointer:
        self.instructions.insert(3, makeIns(arm.movregreg_ext_ins(arm.r7, arm.sp)))
        # Stack grows downwards
        self.instructions.append(makeIns(arm.addspsp_ins(arm.sp, arm.sp, arm.Imm7(self.stacksize))))
        self.instructions.append(makeIns(arm.pop_ins(arm.RegisterSet({arm.pc, arm.r7}))))
        # Add constant literals:
        self.instructions.append(makeIns(Alignment(4))) # Align at 4 bytes
        for ln, v in self.constants:
            self.instructions.append(makeIns(arm.Label(ln)))
            self.instructions.append(makeIns(arm.dcd_ins(v)))


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(arm.addregs_ins, 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(arm.subregs_ins, 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(arm.orrregs_ins, 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(arm.lslregs_ins, 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(arm.mov_imm8_ins, others=[arm.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(arm.ldr_pcrel, 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(arm.loadimm5_ins, 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(arm.loadimm5_ins, 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(arm.bl_ins(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(arm.storeimm5_ins, 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(arm.storeimm5_ins, 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(arm.b_ins(LabelRef(s.target.name)), jumps=[tgt])
        elif isinstance(s, ir.CJump):
            a = self.munchExpr(s.a)
            b = self.munchExpr(s.b)
            self.emit(arm.cmp_ins, src=[a, b])
            ntgt = self.targets[s.lab_no]
            ytgt = self.targets[s.lab_yes]
            jmp_ins = makeIns(arm.b_ins(LabelRef(s.lab_no.name)), jumps=[ntgt])
            opnames = {'<': arm.blt_ins, '>':arm.bgt_ins, '==':arm.beq_ins}
            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(arm.movregreg_ext_ins, src=[src], dst=[dst])