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 ArmFrame(irmach.Frame):
    """
      Arm specific frame for functions.
    """
    pass


class ArmInstructionSelector(InstructionSelector):
    """ Instruction selector for the arm architecture """
    def newFrame(self, name):
        return ArmFrame(name)

    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.Parameter):
            offset = 1337 # TODO: determine offset in frame??
            d = self.newTmp()
            self.emit('ldr %d0, [sp + {}]'.format(offset), dst=[d])
            return d
        elif isinstance(e, ir.Call):
            args = [self.munchExpr(a) for a in e.arguments]
            frame_size = 222 # TODO: determine frame size?
            self.emit('add sp, sp, {}'.format(frame_size))
            # 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', src=[memloc, val])
        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 %l0', 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 %l0', jumps=[ntgt])
            # Explicitely add fallthrough:
            self.emit('jeq %l0', jumps=[ytgt, jmp_ins])
            self.emit2(jmp_ins)
        elif isinstance(s, ir.Terminator):
            pass
        else:
            raise NotImplementedError('Stmt --> {}'.format(s))


# TODO: this class could be target independent:
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 allocFrame(self, f):
        """
            Do register allocation for a single stack frame.
        """
        ilist = f.instructions
        self.useUnused(ilist)
        cfg = flowgraph.FlowGraph(ilist)
        f.cfg = cfg
        ig = registerallocator.InterferenceGraph(cfg)
        f.ig = ig

        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 ilist:
            i.src = tuple(regMap[t] for t in i.src)
            i.dst = tuple(regMap[t] for t in i.dst)
            #print(i)

    def generate(self, ircode):
        # Munch program into a bunch of frames. One frame per function.
        # Each frame has a flat list of abstract instructions.
        frames = self.ins_sel.munchProgram(ircode)
        self.frames = frames
        for f in frames:
            self.allocFrame(f)

        # TODO: Peep-hole here
        # TODO: Materialize assembly
        return frames