import ir
from asmnodes import ALabel, AComment
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)
        print('ARM code generation')
        self.outs.selectSection('data')

        for gvar in ircode.Variables:
            self.emit(ALabel(gvar.name))
            # TODO: use initial value:
            self.emit(arm.dcd_ins(0))

        self.imms = [] # list with immediates relative to PC.
        self.outs.selectSection('code')
        for f in ircode.Functions:
            # Add global variable addresses to immediate list:
            for gvar in ircode.Variables:
                pass  #self.imms.append((

            self.stack_frame = []
            self.emit(ALabel(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(ALabel(bb.name))
                for ins in bb.Instructions:
                    self.generateInstruction(ins)

            self.outs.align(4)
            while self.imms:
                l, v = self.imms.pop()
                self.emit(ALabel(l))
                self.emit(arm.dcd_ins(v))
            self.outs.align(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, ALabel(_global_address)))
    def loadStack(self, reg, val):
        self.emit(arm.ldr_sprel(reg, arm.MemSpRel(self.getStack(val))))
    def comment(self, txt):
        self.emit(AComment(txt))

    def generateInstruction(self, ins):
        if type(ins) is ir.Branch:
            tgt = ALabel(ins.target.name)
            self.emit(arm.jmp_ins(tgt))
        elif type(ins) is ir.ImmLoad:
            self.comment(str(ins))
            lname = ins.target.name + '_ivalue'
            self.emit(arm.ldr_pcrel(arm.r0, ALabel(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:
            self.comment(str(ins))
            # Load value in r0:
            self.loadStack(arm.r0, ins.value)
            # store in memory:
            self.getGlobal(arm.r1, ins.location)
            self.emit(arm.storeimm5_ins(arm.r0, arm.MemR8Rel(arm.r1, 0)))
        elif type(ins) is ir.Load:
            self.comment(str(ins))
            self.getGlobal(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:
            self.comment(str(ins))
            # 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))
            else:
                print('operation not implemented', 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.comment(str(ins))
            self.emit(arm.pop_ins(arm.RegisterSet({arm.r4, arm.r5, arm.r6, arm.r7, arm.pc})))
        elif type(ins) is ir.ConditionalBranch:
            self.comment(str(ins))
            self.loadStack(arm.r0, ins.a)
            self.loadStack(arm.r1, ins.b)
            self.emit(arm.cmp_ins(arm.r1, arm.r0))
            tgt_yes = ALabel(ins.lab1.name)
            if ins.cond == '==':
                self.emit(arm.beq_ins(tgt_yes))
            else:
                print('TODO', ins.cond)
            tgt_no = ALabel(ins.lab2.name)
            self.emit(arm.jmp_ins(tgt_no))
        elif type(ins) is ir.Alloc:
            self.addStack(ins.value)
        else:
            raise CompilerError('IR "{}" not covered'.format(ins))