import logging
import ir
from target import Label, Comment, Alignment, LabelRef, Imm32, DebugInfo, Nop
from target import Imm3
import cortexm3 as arm
from ppci import CompilerError
import registerallocator
from instructionselector import InstructionSelector
import irmach
from irmach import AbstractInstruction as makeIns
import canon
import transform
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 = [arm.r0, arm.r1, arm.r2, arm.r3, arm.r4, arm.r5, arm.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] = arm.r0
        self.tempMap[self.p1] = arm.r1
        self.tempMap[self.p2] = arm.r2
        self.tempMap[self.p3] = arm.r3
        self.tempMap[self.p4] = arm.r4
        self.tempMap[self.fp] = arm.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(arm.Label(self.name)))
        self.instructions.insert(1, makeIns(arm.push_ins(arm.RegisterSet({arm.lr, arm.r7}))))
        # Reserve stack space for locals:
        self.instructions.insert(2, makeIns(arm.subspsp_ins(arm.sp, arm.sp, arm.Imm7(self.stacksize))))
        # Setup frame pointer:
        self.instructions.insert(3, makeIns(arm.movregreg_ext_ins(arm.r7, arm.sp)))
        # Stack grows downwards
        self.instructions.append(makeIns(arm.addspsp_ins(arm.sp, arm.sp, arm.Imm7(self.stacksize))))
        self.instructions.append(makeIns(arm.pop_ins(arm.RegisterSet({arm.pc, arm.r7}))))
        # Add constant literals:
        self.instructions.append(makeIns(Alignment(4))) # Align at 4 bytes
        for ln, v in self.constants:
            self.instructions.append(makeIns(arm.Label(ln)))
            self.instructions.append(makeIns(arm.dcd_ins(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()
            c = Imm3(e.b.value)
            self.emit(arm.addregregimm3_ins, others=[c], 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(arm.addregs_ins, 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()
            c = Imm3(e.b.value)
            self.emit(arm.subregregimm3_ins, others=[c], 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(arm.subregs_ins, 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(arm.orrregs_ins, 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(arm.lslregs_ins, 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(arm.mulregreg_ins, dst=[d], src=[b, d])
            return d
        elif isinstance(e, ir.Const) and e.value < 256:
            d = self.newTmp()
            self.emit(arm.mov_imm8_ins, others=[arm.Imm8(e.value)], dst=[d])
            return d
        elif isinstance(e, ir.Const) and e.value < (2**31):
            d = self.newTmp()
            ln = LabelRef(self.frame.addConstant(e.value))
            self.emit(arm.ldr_pcrel, others=[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(arm.loadimm5_ins, others=[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(arm.loadimm5_ins, others=[0], 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(arm.bl_ins(LabelRef(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(arm.storeimm5_ins, others=[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(arm.storeimm5_ins, others=[0], src=[memloc, val])
        elif isinstance(s, ir.Move) and isinstance(s.dst, ir.Temp):
            val = self.munchExpr(s.src)
            dreg = s.dst
            self.move(dreg, 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(arm.b_ins(LabelRef(s.target.name)), jumps=[tgt])
        elif isinstance(s, ir.CJump):
            a = self.munchExpr(s.a)
            b = self.munchExpr(s.b)
            self.emit(arm.cmp_ins, src=[a, b])
            ntgt = self.targets[s.lab_no]
            ytgt = self.targets[s.lab_yes]
            jmp_ins = makeIns(arm.b_ins(LabelRef(s.lab_no.name)), jumps=[ntgt])
            opnames = {'<': arm.blt_ins, '>':arm.bgt_ins, '==':arm.beq_ins}
            op = opnames[s.cond](LabelRef(s.lab_yes.name))
            self.emit(op, jumps=[ytgt, jmp_ins]) # Explicitely add fallthrough
            self.emit2(jmp_ins)
        else:
            raise NotImplementedError('Stmt --> {}'.format(s))

    def move(self, dst, src):
        self.emit(arm.movregreg_ext_ins, src=[src], dst=[dst])


# 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):
        """ Generate code for one function into a frame """
        # Cleanup function:
        transform.removeEmptyBlocks(irfunc)

        # Create a frame for this function:
        frame = ArmFrame(irfunc.name)

        # Canonicalize the intermediate language:
        canon.make(irfunc, frame)
        self.ins_sel.munchFunction(irfunc, frame)
        
        # Do register allocation:
        self.ra.allocFrame(frame)
        # TODO: Peep-hole here?

        # Can we materialize here??

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

        # Materialize assembly
        # Materialize the register allocated instructions into a stream of
        # real instructions.
        frame.lower_to(self.outs)
        return frame

    def generate(self, ircode):
        self.outs.selectSection('code')
        # assembly glue to make it work:
        # TODO: this must be in source code, not in compiler
        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')))

        # 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]

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