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

# TODO: encode this in DSL (domain specific language)

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)

class Reg8Op:
    def __init__(self, num):
        assert num < 8
        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]
                if r.num < 8:
                    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

def isRegOffset(regname, x, y):
    if type(x) is ASymbol and type(y) is ANumber and x.name.upper() == regname:
        return y.number
    elif type(y) is ASymbol and type(x) is ANumber and y.name.upper() == regname:
        return x.number
    

class MemRegXRel:
    def __init__(self, offset):
        assert offset % 4 == 0
        self.offset = offset

    def __repr__(self):
        return '[{}, #{}]'.format(self.regname, self.offset)

    @classmethod
    def Create(cls, vop):
        if type(vop) is AUnop and vop.operation == '[]':
            vop = vop.arg # descent
            offset = isRegOffset(cls.regname, vop.arg1, vop.arg2)
            if type(offset) is int:
                if offset % 4 == 0:
                    offset = vop.arg2.number
                    return cls(offset)
            elif type(vop) is ASymbol and vop.name.upper() == self.regname:
                return cls(0)

class MemSpRel(MemRegXRel):
    regname = 'SP'

class MemPcRel(MemRegXRel):
    regname = 'PC'

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

    def __repr__(self):
        return '[{}, #{}]'.format(self.basereg, self.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
                    if offset > 120:
                        return
                    basereg = Reg8Op.Create(vop.arg1)
                    if not basereg:
                        return
                else:
                    return
            elif type(vop) is ASymbol:
                offset = 0
                basereg = Reg8Op.Create(vop)
                if not basereg:
                    return
            else:
                return
            return cls(getRegNum(basereg.num), offset)

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


@armtarget.instruction
class dcd_ins(ArmInstruction):
    mnemonic = 'dcd'
    operands = (Imm32,)
    def __init__(self, expr):
        if isinstance(expr, Imm32):
            self.expr = expr.imm
            self.label = None
        elif isinstance(expr, LabelRef):
            self.expr = 0
            self.label = expr
        else:
            raise NotImplementedError()

    def resolve(self, f):
        if self.label:
            self.expr = f(self.label.name)

    def encode(self):
        return u32(self.expr)

    def __repr__(self):
        return 'DCD 0x{0:X}'.format(self.expr)



# Memory related

class LS_imm5_base(ArmInstruction):
    """ ??? Rt, [Rn, imm5] """
    operands = (Reg8Op, MemR8Rel)
    def __init__(self, rt, memop):
        assert memop.offset % 4 == 0
        self.imm5 = memop.offset >> 2
        self.rn = memop.basereg.num
        self.rt = rt
        self.memloc = memop
        assert self.rn < 8
        assert self.rt.num < 8

    def encode(self):
        Rn = self.rn
        Rt = self.rt.num
        imm5 = self.imm5

        h = (self.opcode << 11) | (imm5 << 6) | (Rn << 3) | Rt
        return u16(h)
    def __repr__(self):
        return '{} {}, {}'.format(self.mnemonic, self.rt, self.memloc)

@armtarget.instruction
class storeimm5_ins(LS_imm5_base):
    mnemonic = 'STR'
    opcode = 0xC

@armtarget.instruction
class loadimm5_ins(LS_imm5_base):
    mnemonic = 'LDR'
    opcode = 0xD

class ls_sp_base_imm8(ArmInstruction):
    operands = (Reg8Op, MemSpRel)
    def __init__(self, rt, memop):
        self.rt = rt
        self.offset = memop.offset

    def encode(self):
        rt = self.rt.num
        assert rt < 8
        imm8 = self.offset >> 2
        assert imm8 < 256
        h = (self.opcode << 8) | (rt << 8) | imm8
        return u16(h)

    def __repr__(self):
        return '{} {}, [sp,#{}]'.format(self.mnemonic, self.rt, self.offset)

def align(x, m):
    while ((x % m) != 0):
        x = x + 1
    return x

@armtarget.instruction
class ldr_pcrel(ArmInstruction):
    """ ldr Rt, [PC, imm8], store value into memory """
    mnemonic = 'ldr'
    operands = (RegOp, LabelRef)
    def __init__(self, rt, label):
        assert isinstance(label, LabelRef)
        self.rt = rt
        self.label = label
        self.offset = 0

    def resolve(self, f):
        la = f(self.label.name)
        sa = align(self.address + 2, 4)
        self.offset = (la - sa)
        if self.offset < 0:
            self.offset = 0

    def encode(self):
        rt = self.rt.num
        assert rt < 8
        imm8 = self.offset >> 2
        assert imm8 < 256
        assert imm8 >= 0
        h = (0x9 << 11) | (rt << 8) | imm8
        return u16(h)

    def __repr__(self):
        return 'LDR {}, {}'.format(self.rt, self.label.name)

@armtarget.instruction
class ldr_sprel(ls_sp_base_imm8):
    """ ldr Rt, [SP, imm8] """
    mnemonic = 'LDR'
    opcode = 0x98

@armtarget.instruction
class str_sprel(ls_sp_base_imm8):
    """ str Rt, [SP, imm8] """
    mnemonic = 'STR'
    opcode = 0x90

@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

    def encode(self):
        rd = self.r
        opcode = self.opcode
        imm8 = self.imm
        h = (opcode << 11) | (rd << 8) | imm8
        return u16(h)
    def __repr__(self):
        return 'MOV {0}, xx?'.format(self.r)


@armtarget.instruction
class movregreg_ins(ArmInstruction):
    """ mov Rd, Rm """
    mnemonic = 'mov'
    operands = (Reg8Op, Reg8Op)
    def __init__(self, rd, rm):
        self.rd = rd
        self.rm = rm
    def encode(self):
        rd = self.rd.num
        rm = self.rm.num
        h = 0 | (rm << 3) | rd
        return u16(h)



# Arithmatics:

@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)

class regregreg_base(ArmInstruction):
    """ ??? Rd, Rn, Rm """
    operands = (Reg8Op, Reg8Op, Reg8Op)
    def __init__(self, rd, rn, rm):
        self.rd = rd
        self.rn = rn
        self.rm = rm
    def encode(self):
        rd = self.rd.num
        rn = self.rn.num
        rm = self.rm.num
        h = (self.opcode << 9) | (rm << 6) | (rn << 3) | rd
        return u16(h)
    def __repr__(self):
        return '{} {}, {}, {}'.format(self.mnemonic, self.rd, self.rn, self.rm)

@armtarget.instruction
class addregs_ins(regregreg_base):
    mnemonic = 'ADD'
    opcode = 0b0001100

@armtarget.instruction
class subregs_ins(regregreg_base):
    mnemonic = 'SUB'
    opcode = 0b0001101

class regreg_base(ArmInstruction):
    """ ??? Rdn, Rm """
    operands = (Reg8Op, Reg8Op)
    def __init__(self, rdn, rm):
        self.rdn = rdn
        self.rm = rm
    def encode(self):
        rdn = self.rdn.num
        rm = self.rm.num
        h = (self.opcode << 6) | (rm << 3) | rdn
        return u16(h)
    def __repr__(self):
        return '{} {}, {}'.format(self.mnemonic, self.rdn, self.rm)

@armtarget.instruction
class andregs_ins(regreg_base):
    mnemonic = 'AND'
    opcode = 0b0100000000

@armtarget.instruction
class orrregs_ins(regreg_base):
    mnemonic = 'ORR'
    opcode = 0b0100001100

@armtarget.instruction
class cmp_ins(regreg_base):
    mnemonic = 'CMP'
    opcode = 0b0100001010

@armtarget.instruction
class lslregs_ins(regreg_base):
    mnemonic = 'LSL'
    opcode = 0b0100000010

@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)

# Jumping:

def wrap_negative(x, bits):
    b = struct.unpack('<I', struct.pack('<i', x))[0]
    mask = (1 << bits) - 1
    return b & mask

class jumpBase_ins(ArmInstruction):
    operands = (LabelRef,)
    def __init__(self, target_label):
        assert type(target_label) is LabelRef
        self.target = target_label
        self.offset = 0

    def resolve(self, f):
        la = f(self.target.name)
        sa = self.address + 4
        self.offset = (la - sa)
        #if self.offset < 0:
        #    # TODO: handle negative jump
        #    self.offset = 0

    def __repr__(self):
        return '{} {}'.format(self.mnemonic, self.target.name)

@armtarget.instruction
class b_ins(jumpBase_ins):
    mnemonic = 'B'
    def encode(self):
        imm11 = wrap_negative(self.offset >> 1, 11)
        h = (0b11100 << 11) | imm11 # | 1 # 1 to enable thumb mode
        return u16(h)

class cond_base_ins(jumpBase_ins):
    def encode(self):
        imm8 = wrap_negative(self.offset >> 1, 8)
        h = (0b1101 << 12) | (self.cond << 8) | imm8
        return u16(h)

@armtarget.instruction
class beq_ins(cond_base_ins):
    mnemonic = 'beq'
    cond = 0

@armtarget.instruction
class beq_ins(cond_base_ins):
    mnemonic = 'bne'
    cond = 1

@armtarget.instruction
class push_ins(ArmInstruction):
    operands = (RegisterSet,)
    mnemonic = 'push'
    def __init__(self, regs):
        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)

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

armtarget.check()