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


class InstructionSelector:
    pass


class RegisterAllocator:
    pass


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

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

    def generate(self, ircode):
        assert isinstance(ircode, ir.Module)
        self.logger.info('Generating arm code for {}'.format(ircode.name))
        self.available_regs = {arm.r3, arm.r4, arm.r5, arm.r6, arm.r7}
        self.regmap = {}
        # TODO: get these from linker descriptor?
        self.outs.getSection('code').address = 0x08000000
        self.outs.getSection('data').address = 0x20000000
        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(arm.bl_ins(LabelRef('main')))

        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.r3, 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()
        codesize = self.outs.getSection('code').Size
        self.logger.info('Generated {} bytes code'.format(codesize))

    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 getreg(self, v):
        if not v in self.regmap:
            self.regmap[v] = self.available_regs.pop()
        return self.regmap[v]

    def freereg(self, v, ins):
        if v.lastUse(ins):
            r = self.regmap.pop(v)
            assert r not in self.regmap.values()
            self.available_regs.add(r)

    def comment(self, txt):
        self.emit(Comment(txt))

    def debugInfo(self, loc):
        if loc:
            self.emit(DebugInfo(loc))

    def generateInstruction(self, ins):
        self.comment(str(ins))
        if hasattr(ins, 'debugLoc'):
            self.debugInfo(ins.debugLoc)
        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'
            r0 = self.getreg(ins.target)
            self.emit(arm.ldr_pcrel(r0, LabelRef(lname)))
            self.imms.append((lname, ins.value))
        elif type(ins) is ir.Store:
            # Load value in r0:
            r0 = self.getreg(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:
                # The value was alloc'ed
                self.emit(arm.str_sprel(r0, arm.MemSpRel(self.getStack(ins.location))))
            else:
                r1 = self.getreg(ins.location)
                self.emit(arm.storeimm5_ins(r0, arm.MemR8Rel(r1, 0)))
                self.freereg(ins.location, ins)
            self.freereg(ins.value, ins)
        elif type(ins) is ir.Load:
            # TODO: differ global and local??
            #self.getGlobal(arm.r0, ins.location)
            r0 = self.getreg(ins.value)
            if ins.location in self.localVars:
                self.emit(arm.ldr_sprel(r0, arm.MemSpRel(self.getStack(ins.location))))
            else:
                r2 = self.getreg(ins.location)
                self.emit(arm.loadimm5_ins(r0, arm.MemR8Rel(r2, 0)))
                self.freereg(ins.location, ins)
        elif type(ins) is ir.BinaryOperator:
            # Load operands:
            r0 = self.getreg(ins.value1)
            r1 = self.getreg(ins.value2)
            r2 = self.getreg(ins.result)
            # do operation:
            if ins.operation == '+':
                self.emit(arm.addregs_ins(r2, r0, r1))
            elif ins.operation == '<<':
                self.emit(arm.movregreg_ins(r2, r0))
                self.emit(arm.lslregs_ins(r2, r1))
            elif ins.operation == '|':
                self.emit(arm.movregreg_ins(r2, r0))
                self.emit(arm.orrregs_ins(r2, r1))
            else:
                raise NotImplementedError('operation {} not implemented'.format(ins.operation))
            self.freereg(ins.value1, ins)
            self.freereg(ins.value2, ins)
        elif type(ins) is ir.Call:
            # TODO: prep parameters:
            self.emit(arm.bl_ins(LabelRef(ins.callee.name)))
        elif type(ins) is ir.Return:
            self.emit(arm.pop_ins(arm.RegisterSet({arm.r3, arm.r4, arm.r5, arm.r6, arm.r7, arm.pc})))
        elif type(ins) is ir.ConditionalBranch:
            r0 = self.getreg(ins.a)
            r1 = self.getreg(ins.b)
            self.emit(arm.cmp_ins(r1, r0))
            tgt_yes = LabelRef(ins.lab1.name)
            if ins.cond == '==':
                self.emit(arm.beq_ins(tgt_yes))
            elif ins.cond == '<':
                self.emit(arm.blt_ins(tgt_yes))
            elif ins.cond == '>':
                self.emit(arm.bgt_ins(tgt_yes))
            else:
                raise NotImplementedError('"{}" not covered'.format(ins.cond))
            tgt_no = LabelRef(ins.lab2.name)
            self.emit(arm.b_ins(tgt_no))
            self.freereg(ins.a, ins)
            self.freereg(ins.b, ins)
        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))