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

class ArmFrame(irmach.Frame):
    """
      Arm specific frame for functions.
    """
    def __init__(self, name):
        # We use r7 as frame pointer.
        super().__init__(name)
        self.regs = ['r0', 'r1', 'r2', 'r3', 'r4', 'r5', 'r6']
        self.rv = ir.Temp('special_RV')
        self.p1 = ir.Temp('special_P1')
        self.p2 = ir.Temp('special_P2')
        self.p3 = ir.Temp('special_P3')
        self.p4 = ir.Temp('special_P4')
        self.fp = ir.Temp('special_FP')
        # Pre-colored registers:
        self.tempMap = {}
        self.tempMap[self.rv] = 'r0'
        self.tempMap[self.p1] = 'r1'
        self.tempMap[self.p2] = 'r2'
        self.tempMap[self.p3] = 'r3'
        self.tempMap[self.p4] = 'r4'
        self.tempMap[self.fp] = 'r7'
        self.locVars = {}
        self.parMap = {}
        # Literal pool:
        self.constants = []

    def argLoc(self, pos):
        """
            Gets the function parameter location in IR-code format.
        """
        if pos == 0:
            return self.p1
        elif pos == 1:
            return self.p2
        elif pos == 2:
            return self.p3
        elif pos == 3:
            return self.p4
        else:
            raise NotImplementedError('No more than 4 parameters implemented')

    def allocVar(self, lvar):
        if lvar not in self.locVars:
            self.locVars[lvar] = self.stacksize
            self.stacksize = self.stacksize + 4
        return self.locVars[lvar]

    def addConstant(self, value):
        lab_name = '{}_literal_{}'.format(self.name, len(self.constants))
        self.constants.append((lab_name, value))
        return lab_name

    def EntryExitGlue3(self):
        """
            Add code for the prologue and the epilogue. Add a label, the
            return instruction and the stack pointer adjustment for the frame.
        """
        self.instructions.insert(0, makeIns('{}:'.format(self.name)))
        self.instructions.insert(1, makeIns('push {lr, r7}'))
        # Reserve stack space for locals:
        self.instructions.insert(2, makeIns('sub sp, sp, {}'.format(self.stacksize)))
        # Setup frame pointer:
        self.instructions.insert(3, makeIns('mov r7, sp'))
        # Stack grows downwards
        self.instructions.append(makeIns('add sp, sp, {}'.format(self.stacksize)))
        self.instructions.append(makeIns('pop {pc,r7}'))
        # Add constant literals:
        self.instructions.append(makeIns('align 4')) # Align at 4 bytes
        for ln, v in self.constants:
            self.instructions.append(makeIns('{}:'.format(ln)))
            self.instructions.append(makeIns('dcd {}'.format(v)))


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 == '+' and \
            isinstance(e.b, ir.Const) and e.b.value < 8:
            a = self.munchExpr(e.a)
            d = self.newTmp()
            self.emit('add %d0, %s0, {}'.format(e.b.value), dst=[d], src=[a])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '+':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.emit('add %d0, %s0, %s1', dst=[d], src=[a, b])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '-' and \
            isinstance(e.b, ir.Const) and e.b.value < 8:
            a = self.munchExpr(e.a)
            d = self.newTmp()
            self.emit('sub %d0, %s0, {}'.format(e.b.value), dst=[d], src=[a])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '-':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            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.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            self.emit('orr %s1, %s0', dst=[], src=[b, d])
            return d
        elif isinstance(e, ir.Binop) and e.operation == '<<':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            self.emit('lsl %s1, %s0', dst=[], src=[b, d]) # TODO: is d a source variable?
            return d
        elif isinstance(e, ir.Binop) and e.operation == '*':
            a = self.munchExpr(e.a)
            b = self.munchExpr(e.b)
            d = self.newTmp()
            self.move(d, a)
            # this mul instruction has operands swapped:
            self.emit('mul %s0, %d0', dst=[d], src=[b, d])
            return d
        elif isinstance(e, ir.Const) and e.value < 256:
            d = self.newTmp()
            self.emit('mov %d0, {}'.format(e.value), dst=[d])
            return d
        elif isinstance(e, ir.Const) and e.value < (2**31):
            d = self.newTmp()
            ln = self.frame.addConstant(e.value)
            self.emit('ldr %d0, {}'.format(ln), dst=[d])
            return d
        elif isinstance(e, ir.Mem) and isinstance(e.e, ir.Binop) and \
                e.e.operation == '+' and isinstance(e.e.b, ir.Const):
            base = self.munchExpr(e.e.a)
            d = self.newTmp()
            c = e.e.b.value
            self.emit('ldr %d0, [%s0 + {}]'.format(c), src=[base], dst=[d])
            return d
        elif isinstance(e, ir.Mem):
            # Load from memory
            base = self.munchExpr(e.e)
            d = self.newTmp()
            self.emit('ldr %d0, [%s0]', src=[base], dst=[d])
            return d
        elif isinstance(e, ir.Temp):
            return e
        elif isinstance(e, ir.Call):
            # Move arguments into proper locations:
            reguses = []
            for i, a in enumerate(e.arguments):
                loc = self.frame.argLoc(i)
                m = ir.Move(loc, a)
                self.munchStm(m)
                if isinstance(loc, ir.Temp):
                    reguses.append(loc)
            self.emit('bl {}'.format(e.f.name), src=reguses, dst=[self.frame.rv])
            d = self.newTmp()
            self.move(d, self.frame.rv)
            return d
        else:
            raise NotImplementedError('Expr --> {}'.format(e))

    def munchStm(self, s):
        if isinstance(s, ir.Terminator):
            pass
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Mem) and \
            isinstance(s.dst.e, ir.Binop) and s.dst.e.operation == '+' and \
            isinstance(s.dst.e.b, ir.Const):
            a = self.munchExpr(s.dst.e.a)
            val = self.munchExpr(s.src)
            c = s.dst.e.b.value
            self.emit('str %s1, [%s0 + {}]'.format(c), src=[a, val])
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Mem):
            memloc = self.munchExpr(s.dst.e)
            val = self.munchExpr(s.src)
            self.emit('str %s1, [%s0]', src=[memloc, val])
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Temp):
            val = self.munchExpr(s.src)
            dreg = s.dst
            self.emit('mov %d0, %s0', dst=[dreg], src=[val])
        elif isinstance(s, ir.Exp):
            # Generate expression code and discard the result.
            x = self.munchExpr(s.e)
            self.emit('nop', src=[x])
        elif isinstance(s, ir.Jump):
            tgt = self.targets[s.target]
            self.emit('b {}'.format(s.target.name), 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 = makeIns('b {}'.format(s.lab_no.name), jumps=[ntgt])
            # Explicitely add fallthrough:
            self.emit('beq {}'.format(s.lab_yes.name), jumps=[ytgt, jmp_ins])
            self.emit2(jmp_ins)
        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.ra = registerallocator.RegisterAllocator()
        self.outs = outs
        self.outs.getSection('code').address = 0x08000000
        self.outs.getSection('data').address = 0x20000000

    def generateFunc(self, irfunc):
        # Create a frame for this function:
        frame = ArmFrame(irfunc.name)

        # Canonicalize the intermediate language:
        canon.make(irfunc, frame)
        print('after canonicalize:')
        irfunc.dump()
        self.ins_sel.munchFunction(irfunc, frame)
        print('Selected instructions:')
        for i in frame.instructions:
            print(i)
        
        # Do register allocation:
        self.ra.allocFrame(frame)
        # TODO: Peep-hole here?

        # Add label and return and stack adjustment:
        frame.EntryExitGlue3()
        return frame

    def generate(self, ircode):
        # Munch program into a bunch of frames. One frame per function.
        # Each frame has a flat list of abstract instructions.
        # Generate code for all functions:
        self.frames = [self.generateFunc(func) for func in ircode.Functions]

        # Materialize assembly
        # Reparse the register allocated instructions into a stream of
        # real instructions.
        # TODO: this is ugly via string representations. This could be 
        # another interface?
        assembler = asm.Assembler(target=arm.armtarget, stream=self.outs)
        self.outs.selectSection('code')
        # assembly glue to make it work:
        self.outs.emit(arm.dcd_ins(Imm32(0x20000678))) # initial SP
        self.outs.emit(arm.dcd_ins(Imm32(0x08000009))) # reset vector
        self.outs.emit(arm.b_ins(LabelRef('main')))
        for frame in self.frames:
            for i in frame.instructions:
                assembler.assemble_line(str(i))

        # TODO: fixup references, do this in another way?
        self.outs.backpatch()
        self.outs.backpatch()
        return self.frames