import ir
from .basetarget import Label, Comment, Alignment, LabelRef, DebugInfo, Nop
from .basetarget import Imm7
from irmach import AbstractInstruction as makeIns, Frame
from .arminstructions import Dcd, AddSp, SubSp, Push, Pop, Mov2
from .arminstructions import R0, R1, R2, R3, R4, R5, R6, R7, LR, PC, SP


class ArmFrame(Frame):
    """ Arm specific frame for functions. """
    def __init__(self, name):
        # We use r7 as frame pointer.
        super().__init__(name)
        self.regs = [R0, R1, R2, R3, R4, R5, 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] = R0
        self.tempMap[self.p1] = R1
        self.tempMap[self.p2] = R2
        self.tempMap[self.p3] = R3
        self.tempMap[self.p4] = R4
        self.tempMap[self.fp] = 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(Label(self.name)))
        self.instructions.insert(1, makeIns(Push({LR, R7})))
        # Reserve stack space for locals:
        self.instructions.insert(2, makeIns(SubSp(SP, SP, Imm7(self.stacksize))))
        # Setup frame pointer:
        self.instructions.insert(3, makeIns(Mov2(R7, SP)))
        # Stack grows downwards
        self.instructions.append(makeIns(AddSp(SP, SP, Imm7(self.stacksize))))
        self.instructions.append(makeIns(Pop({PC, R7})))
        # Add constant literals:
        self.instructions.append(makeIns(Alignment(4))) # Align at 4 bytes
        for ln, v in self.constants:
            self.instructions.append(makeIns(Label(ln)))
            self.instructions.append(makeIns(Dcd(v)))