import struct, types
from target import Register, Instruction, Target, Imm8, Label, Imm3
from asmnodes import ASymbol, ANumber, AUnop, ABinop
from ppci import CompilerError
import ir

def u16(h):
    return struct.pack('<H', h)

def u32(x):
    return struct.pack('<I', x)

armtarget = Target('arm')

class ArmReg(Register):
    def __init__(self, num, name):
        super().__init__(name)
        self.num = num
    def __repr__(self):
        return self.name

class RegOp:
    def __init__(self, num):
        assert num < 16
        self.num = num

    @classmethod
    def Create(cls, vop):
        if type(vop) is ASymbol:
            name = vop.name
            regs = {}
            for r in armtarget.registers:
                regs[r.name] = r
            if name in regs:
                r = regs[name]
                return cls(r.num)

def getRegNum(n):
    for r in armtarget.registers:
        if r.num == n:
            return r

def getRegisterRange(n1, n2):
    regs = []
    if n1.num < n2.num:
        for n in range(n1.num, n2.num + 1):
            r = getRegNum(n)
            assert r
            regs.append(r)
    return regs

class MemoryOp:
    def __init__(self, basereg, offset):
        assert type(basereg) is ArmReg
        self.basereg = basereg
        self.offset = offset

    @classmethod
    def Create(cls, vop):
        if type(vop) is AUnop and vop.operation == '[]':
            vop = vop.arg # descent
            if type(vop) is ABinop:
                if vop.op == '+' and type(vop.arg1) is ASymbol and type(vop.arg2) is ANumber:
                    offset = vop.arg2.number
                    basereg = RegOp.Create(vop.arg1)
                    if not basereg:
                        return
                else:
                    return
            elif type(vop) is ASymbol:
                offset = 0
                basereg = RegOp.Create(vop)
                if not basereg:
                    return
            else:
                return
            return cls(getRegNum(basereg.num), offset)
        pass

class RegisterSet:
    def __init__(self, regs):
        assert type(regs) is set
        self.regs = regs
    def __repr__(self):
        return ','.join([str(r) for r in self.regs])
    @classmethod
    def Create(cls, vop):
        assert type(vop) is AUnop and vop.operation == '{}'
        assert type(vop.arg) is list
        regs = set()
        for arg in vop.arg:
            if type(arg) is ASymbol:
                reg = RegOp.Create(arg)
                if not reg:
                    return
                regs.add(reg)
            elif type(arg) is ABinop and arg.op == '-':
                reg1 = RegOp.Create(arg.arg1)
                reg2 = RegOp.Create(arg.arg2)
                if not reg1:
                    return
                if not reg2:
                    return
                for r in getRegisterRange(reg1, reg2):
                    regs.add(r)
            else:
                raise Exception('Cannot be')
        return cls(regs)

    def registerNumbers(self):
        return [r.num for r in self.regs]

# 8 bit registers:
r0 = ArmReg(0, 'r0')
armtarget.registers.append(r0)
r1 = ArmReg(1, 'r1')
armtarget.registers.append(r1)
r2 = ArmReg(2, 'r2')
armtarget.registers.append(r2)
r3 = ArmReg(3, 'r3')
armtarget.registers.append(r3)
r4 = ArmReg(4, 'r4')
armtarget.registers.append(r4)
r5 = ArmReg(5, 'r5')
armtarget.registers.append(r5)
r6 = ArmReg(6, 'r6')
armtarget.registers.append(r6)
r7 = ArmReg(7, 'r7')
armtarget.registers.append(r7)
# Other registers:
# TODO
sp = ArmReg(13, 'sp')
armtarget.registers.append(sp)
lr = ArmReg(14, 'lr')
armtarget.registers.append(lr)
pc = ArmReg(15, 'pc')
armtarget.registers.append(pc)

class ArmInstruction(Instruction):
    pass

class ldr_ins(ArmInstruction):
    mnemonic = 'ldr'
    opcode = 1337
    irpattern = 'todo'

class dcd_ins(ArmInstruction):
    mnemonic = 'dcd'
    def __init__(self, expr):
        self.expr = expr
    def encode(self):
        return u32(self.expr)

@armtarget.instruction
class storeimm5_ins(ArmInstruction):
    """ str Rt, [Rn, imm5], store value into memory """
    mnemonic = 'str'
    operands = (RegOp, MemoryOp)
    def __init__(self, rt, memop):
        assert memop.offset % 4 == 0
        self.imm5 = memop.offset >> 2
        self.rn = memop.basereg.num
        self.rt = rt.num

    def encode(self):
        Rn = self.rn
        Rt = self.rt
        imm5 = self.imm5
        h = (0xC << 11) | (imm5 << 6) | (Rn << 3) | Rt
        return u16(h)

@armtarget.instruction
class loadimm5_ins(ArmInstruction):
    """ str Rt, [Rn, imm5], store value into memory """
    mnemonic = 'ldr'
    operands = (RegOp, MemoryOp)
    def __init__(self, rt, memop):
        assert memop.offset % 4 == 0
        self.imm5 = memop.offset >> 2
        self.rn = memop.basereg.num
        self.rt = rt.num

    def encode(self):
        Rn = self.rn
        Rt = self.rt
        imm5 = self.imm5
        h = (0xD << 11) | (imm5 << 6) | (Rn << 3) | Rt
        return u16(h)

@armtarget.instruction
class mov_ins(ArmInstruction):
    """ mov Rd, imm8, move immediate value into register """
    mnemonic = 'mov'
    opcode = 4 # 00100 Rd(3) imm8
    operands = (RegOp, Imm8)
    irpattern = ir.ImmLoad
    def __init__(self, rd, imm):
        self.imm = imm.imm
        self.r = rd.num

    @classmethod
    def FromIr(cls, ir_ins):
        pass

    def encode(self):
        rd = self.r
        opcode = self.opcode
        imm8 = self.imm
        h = (opcode << 11) | (rd << 8) | imm8
        return u16(h)
        
@armtarget.instruction
class movregreg_ins(ArmInstruction):
    """ mov Rd, Rm """
    mnemonic = 'mov'
    opcode = 8 # 01000 Rd(3) imm8
    operands = (RegOp, RegOp)
    def __init__(self, rd, rm):
        self.rd = rd
        self.rm = rm
    def encode(self):
        rd = self.rd.num
        D = (rd & 0x8) >> 3
        assert D < 2
        rd = rd & 0x7
        rm = self.rm.num
        assert rm < 16
        opcode = self.opcode
        h = (opcode << 11) | (3 << 9) | (D << 7) | (rm << 3) | rd
        return u16(h)

@armtarget.instruction
class addregregimm3_ins(ArmInstruction):
    """ add Rd, Rn, imm3 """
    mnemonic = 'add'
    opcode = 3 # 00011
    operands = (RegOp, RegOp, Imm3)
    irpattern = 3
    def __init__(self, rd, rn, imm3):
        self.rd = rd
        self.rn = rn
        self.imm3 = imm3
    def encode(self):
        rd = self.rd.num
        rn = self.rn.num
        imm3 = self.imm3.imm
        opcode = self.opcode
        h = (opcode << 11) | (1 << 10) | (imm3 << 6) | (rn << 3) | rd
        return u16(h)

@armtarget.instruction
class cmpregimm8_ins(ArmInstruction):
    """ cmp Rn, imm8 """
    mnemonic = 'cmp'
    opcode = 5 # 00101
    operands = (RegOp, Imm8)
    def __init__(self, rn, imm):
        self.rn = rn
        self.imm = imm
    def encode(self):
        rn = self.rn.num
        imm = self.imm.imm
        opcode = self.opcode
        h = (opcode << 11) | (rn << 8) | imm
        return u16(h)

@armtarget.instruction
class jmp_ins(ArmInstruction):
    operands = (Label,)
    mnemonic = 'jmp'
    def __init__(self, target_label):
        self.target = target_label
    def fixUp(self):
        pass
    def encode(self):
        h = 1337 # TODO
        return u16(h)

@armtarget.instruction
class push_ins(ArmInstruction):
    operands = (RegisterSet,)
    mnemonic = 'push'
    def __init__(self, regs):
        print(self.operands)
        assert (type(regs),) == self.operands, (type(regs),)
        self.regs = regs
    def __repr__(self):
        return '{0} {{{1}}}'.format(self.mnemonic, self.regs)
    def encode(self):
        reg_list = 0
        M = 0
        for n in self.regs.registerNumbers():
            if n < 8:
                reg_list |= (1 << n)
            elif n == 14:
                M = 1
            else:
                raise NotImplementedError('not implemented for this register')
        h = (0x5a << 9) | (M << 8) | reg_list
        return u16(h)

@armtarget.instruction
class pop_ins(ArmInstruction):
    operands = (RegisterSet,)
    mnemonic = 'pop'
    def __init__(self, regs):
        self.regs = regs
    def __repr__(self):
        return '{0} {{{1}}}'.format(self.mnemonic, self.regs)
    def encode(self):
        reg_list = 0
        P = 0
        for n in self.regs.registerNumbers():
            if n < 8:
                reg_list |= (1 << n)
            elif n == 15:
                P = 1
            else:
                raise NotImplementedError('not implemented for this register')
        h = (0x5E << 9) | (P << 8) | reg_list
        return u16(h)
        return u16(0)

@armtarget.instruction
class yield_ins(ArmInstruction):
    operands = ()
    mnemonic = 'yield'
    def encode(self):
        return u16(0xbf10)

armtarget.check()