import ir
from target import Label, Comment, Alignment, LabelRef, Imm32
import cortexm3 as arm
from ppci import CompilerError

class ArmCodeGenerator:
    """
        Simple code generator
        Ad hoc implementation
    """
    def __init__(self, out):
        self.outs = out

    def emit(self, item):
        self.outs.emit(item)

    def generate(self, ircode):
        assert isinstance(ircode, ir.Module)
        self.outs.selectSection('data')

        for gvar in ircode.Variables:
            self.emit(Label(gvar.name))
            # TODO: use initial value:
            self.dcd(0)

        self.imms = [] # list with immediates relative to PC.
        self.outs.selectSection('code')

        # Manually inserted startup code:
        self.dcd(0x20000678) # initial stack ptr
        # TODO: use label here:
        #self.emit(arm.dcd_ins(LabelRef('reset'))) # reset vector
        self.dcd(0x08000009) # reset vector, lsb indicates thumb mode

        self.emit(Label('reset'))
        for f in ircode.Functions:
            self.localVars = []
            # Add global variable addresses to immediate list:
            for gvar in ircode.Variables:
                pass  #self.imms.append((

            self.stack_frame = []
            self.emit(Label(f.name))
            # Save some registers:
            self.emit(arm.push_ins(arm.RegisterSet({arm.r4, arm.r5, arm.r6,arm.r7,arm.lr})))
            for bb in f.BasicBlocks:
                self.emit(Label(bb.name))
                for ins in bb.Instructions:
                    self.generateInstruction(ins)

            self.align()
            while self.imms:
                l, v = self.imms.pop()
                self.emit(Label(l))
                self.dcd(v)
            self.align()
        self.outs.backpatch()
        self.outs.backpatch()

    def dcd(self, x):
        self.emit(arm.dcd_ins(Imm32(x)))

    def align(self):
        self.outs.emit(Alignment(4))
    # Helper functions:
    def getStack(self, v):
        off = self.stack_frame.index(v)
        return off * 4
    def addStack(self, v):
        self.stack_frame.append(v)
        return self.getStack(v)
    def getGlobal(self, r, g):
        _global_address = g.name + '__global'
        self.emit(arm.ldr_pcrel(r, LabelRef(_global_address)))
    def loadStack(self, reg, val):
        self.emit(arm.ldr_sprel(reg, arm.MemSpRel(self.getStack(val))))
    def comment(self, txt):
        self.emit(Comment(txt))

    def generateInstruction(self, ins):
        self.comment(str(ins))
        if type(ins) is ir.Branch:
            tgt = LabelRef(ins.target.name)
            self.emit(arm.b_ins(tgt))
        elif type(ins) is ir.ImmLoad:
            lname = ins.target.name + '_ivalue'
            self.emit(arm.ldr_pcrel(arm.r0, LabelRef(lname)))
            self.imms.append((lname, ins.value))
            self.emit(arm.str_sprel(arm.r0, arm.MemSpRel(self.addStack(ins.target))))
        elif type(ins) is ir.Store:
            print('STORE', ins, self.localVars)
            # Load value in r0:
            self.loadStack(arm.r0, ins.value)
            # store in memory:
            # TODO: split globals and locals??
            #self.getGlobal(arm.r1, ins.location)
            # Horrible hack with localVars
            if ins.location in self.localVars:
                print('local', ins.location)
                self.emit(arm.str_sprel(arm.r0, arm.MemSpRel(self.getStack(ins.location))))
            else:
                self.loadStack(arm.r1, ins.location)
                self.emit(arm.storeimm5_ins(arm.r0, arm.MemR8Rel(arm.r1, 0)))
        elif type(ins) is ir.Load:
            # TODO: differ global and local??
            #self.getGlobal(arm.r0, ins.location)
            if ins.location in self.localVars:
                self.emit(arm.ldr_sprel(arm.r0, arm.MemSpRel(self.getStack(ins.location))))
            else:
                self.loadStack(arm.r0, ins.location)
                self.emit(arm.loadimm5_ins(arm.r0, arm.MemR8Rel(arm.r0, 0)))
            # Store value on stack:
            self.emit(arm.str_sprel(arm.r0, arm.MemSpRel(self.addStack(ins.value))))
        elif type(ins) is ir.BinaryOperator:
            # Load operands:
            self.loadStack(arm.r0, ins.value1)
            self.loadStack(arm.r1, ins.value2)
            # do operation:
            if ins.operation == '+':
                self.emit(arm.addregs_ins(arm.r0, arm.r0, arm.r1))
            elif ins.operation == '<<':
                self.emit(arm.lslregs_ins(arm.r0, arm.r1))
            elif ins.operation == '|':
                self.emit(arm.orrregs_ins(arm.r0, arm.r1))
            else:
                raise NotImplementedError('operation {} not implemented'.format(ins.operation))
            # Store value back:
            self.emit(arm.str_sprel(arm.r0, arm.MemSpRel(self.addStack(ins.result))))
        elif type(ins) is ir.Return:
            self.emit(arm.pop_ins(arm.RegisterSet({arm.r4, arm.r5, arm.r6, arm.r7, arm.pc})))
        elif type(ins) is ir.ConditionalBranch:
            self.loadStack(arm.r0, ins.a)
            self.loadStack(arm.r1, ins.b)
            self.emit(arm.cmp_ins(arm.r1, arm.r0))
            tgt_yes = Label(ins.lab1.name)
            if ins.cond == '==':
                self.emit(arm.beq_ins(tgt_yes))
            else:
                raise NotImplementedError('"{}" not covered'.format(ins.cond))
            tgt_no = Label(ins.lab2.name)
            self.emit(arm.jmp_ins(tgt_no))
        elif type(ins) is ir.Alloc:
            # Local variables are added to stack
            self.addStack(ins.value)
            self.localVars.append(ins.value)
            # load address into variable:
        else:
            raise NotImplementedError('IR "{}" not covered'.format(ins))