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

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 == '+':
            a = self.munchExpr(e.value1)
            b = self.munchExpr(e.value2)
            d = self.newTmp()
            self.emit('add %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '-':
            a = self.munchExpr(e.value1)
            b = self.munchExpr(e.value2)
            d = self.newTmp()
            self.emit('sub %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '|':
            a = self.munchExpr(e.value1)
            b = self.munchExpr(e.value2)
            d = self.newTmp()
            self.emit('or %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '<<':
            a = self.munchExpr(e.value1)
            b = self.munchExpr(e.value2)
            d = self.newTmp()
            self.emit('lsl %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '*':
            a = self.munchExpr(e.value1)
            b = self.munchExpr(e.value2)
            d = self.newTmp()
            self.emit('mul %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Const):
            d = self.newTmp()
            if e.value < 256:
                self.emit('ldr %d0, {}'.format(e.value), dst=[d])
            else:
                self.emit('ldrpcrel TODO', dst=[d])
            return d
        elif isinstance(e, ir.Mem):
            # Load from memory
            loc = self.munchExpr(e.e)
            d = self.newTmp()
            self.emit('ldr %d0, [%s0]', src=[loc], dst=[d])
            return d
        elif isinstance(e, ir.Temp):
            return self.getTempReg(e)
        elif isinstance(e, ir.Call):
            args = [self.munchExpr(a) for a in e.arguments]
            self.emit('add sp, sp, 22')
            # TODO: save frame
            for a in args:
                self.emit('push %s0', src=[a])
            self.emit('bl {}'.format(e.f.name))
            self.emit('sub sp, sp, 22')
        else:
            raise NotImplementedError('Expr --> {}'.format(e))

    def munchStm(self, s):
        if isinstance(s, ir.Move) and isinstance(s.dst, ir.Mem):
            memloc = self.munchExpr(s.dst.e)
            val = self.munchExpr(s.src)
            self.emit('str [%s0], %s1')
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Temp):
            val = self.munchExpr(s.src)
            dreg = self.getTempReg(s.dst)
            self.emit('mov %d0, %s0', dst=[dreg], src=[val])
        elif isinstance(s, ir.Jump):
            tgt = self.targets[s.target]
            self.emit('jmp {}'.format(s), jumps=[tgt])
        elif isinstance(s, ir.CJump):
            a = self.munchExpr(s.a)
            b = self.munchExpr(s.b)
            self.emit('cmp %s0, %s1', src=[a, b])
            ntgt = self.targets[s.lab_no]
            ytgt = self.targets[s.lab_yes]
            jmp_ins = self.makeIns('jmp {}'.format(s.lab_no), jumps=[ntgt])
            # Explicitely add fallthrough:
            self.emit('jeq {}'.format(s.lab_yes), jumps=[ytgt, jmp_ins])
            self.emit2(jmp_ins)
        else:
            raise NotImplementedError('--> {}'.format(s))


class ArmCodeGenerator:
    def __init__(self, outs):
        # TODO: schedule traces in better order.
        # This is optional!
        self.ins_sel = ArmInstructionSelector()
        self.outs = outs
        self.outs.getSection('code').address = 0x08000000
        self.outs.getSection('data').address = 0x20000000

    def useUnused(self, inslist):
        # Use unused temporaries at the end of the list
        defTemps = []
        useTemps = []
        for i in inslist:
            for d in iter(i.dst):
                defTemps.append(d)
            for s in iter(i.src):
                useTemps.append(s)
        defTemps = set(defTemps)
        useTemps = set(useTemps)
        unUsed = defTemps - useTemps
        #print('Unused:', unUsed)
        for uu in unUsed:
            inslist.append(irmach.AbstractInstruction('use %s0', src=[uu]))
        #print(useTemps)

    def generate(self, ircode, cfg_file=None, ig_file=None):
        ir2 = self.ins_sel.munchProgram(ircode)
        self.useUnused(ir2)
        cfg = flowgraph.FlowGraph(ir2)
        if cfg_file:
            cfg.to_dot(cfg_file)
        ig = registerallocator.InterferenceGraph(cfg)
        if ig_file:
            ig.to_dot(ig_file)

        regs = ['r0', 'r1', 'r2', 'r3', 'r4', 'r5', 'r6', 'r7']
        ra = registerallocator.RegisterAllocator()
        regMap = ra.registerAllocate(ig, regs)
        #print(regMap)
        # Use allocated registers:
        for i in ir2:
            i.src = tuple(regMap[t] for t in i.src)
            i.dst = tuple(regMap[t] for t in i.dst)
            #print(i)
        return ir2