import struct
from asmnodes import ASymbol, AInstruction, ANumber, AUnop, ABinop
from .basetarget import Register, Instruction, Target, Label, LabelRef
from .basetarget import Imm32, Imm8, Imm7, Imm3


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

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

# Operands:

class ArmRegister(Register):
    def __init__(self, num, name):
        super().__init__(name)
        self.num = num

    def __repr__(self):
        return self.name

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


class Reg8Op(ArmRegister):
    pass


class Reg16Op(ArmRegister):
    pass


R0 = Reg8Op(0, 'r0')
R1 = Reg8Op(1, 'r1')
R2 = Reg8Op(2, 'r2')
R3 = Reg8Op(3, 'r3')
R4 = Reg8Op(4, 'r4')
R5 = Reg8Op(5, 'r5')
R6 = Reg8Op(6, 'r6')
R7 = Reg8Op(7, 'r7')
# Other registers:
# TODO
SP = ArmRegister(13, 'sp')
LR = ArmRegister(14, 'lr')
PC = ArmRegister(15, 'pc')

registers = [R0, R1, R2, R3, R4, R5, R6, R7, SP, LR, PC]


class RegSpOp:
    @classmethod
    def Create(cls, vop):
        if type(vop) is ASymbol:
            if vop.name.lower() == 'sp':
                return cls()


def getRegNum(n):
    for r in 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 == '[]' and type(vop.arg) is ABinop and vop.arg.op == '+':
            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 MemR8Rel:
    def __init__(self, basereg, offset):
        assert type(basereg) is Reg8Op
        assert type(offset) is int
        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 = ArmRegister.Create(arg)
                if not reg:
                    return
                regs.add(reg)
            elif type(arg) is ABinop and arg.op == '-':
                reg1 = ArmRegister.Create(arg.arg1)
                reg2 = ArmRegister.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]



# Instructions:

class ArmInstruction(Instruction):
    pass


allins = []


def instruction(i):
    allins.append(i)
    return i


@instruction
class Dcd(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
        elif isinstance(expr, int):
            self.expr = expr
            self.label = None
        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)


@instruction
class nop_ins(ArmInstruction):
    mnemonic = 'nop'
    operands = tuple()

    def encode(self):
        return bytes()

    def __repr__(self):
        return 'NOP'


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


@instruction
class Str2(LS_imm5_base):
    mnemonic = 'STR'
    opcode = 0xC

    @classmethod
    def fromim(cls, im):
        mem = MemR8Rel(im.src[0], im.others[0])
        return cls(im.src[1], mem)


@instruction
class Ldr2(LS_imm5_base):
    mnemonic = 'LDR'
    opcode = 0xD

    @classmethod
    def fromim(cls, im):
        mem = MemR8Rel(im.src[0], im.others[0])
        return cls(im.dst[0], mem)

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


@instruction
class Ldr3(ArmInstruction):
    """ ldr Rt, LABEL, load value from pc relative position """
    mnemonic = 'ldr'
    operands = (Reg8Op, LabelRef)
    def __init__(self, rt, label):
        assert isinstance(label, LabelRef)
        self.rt = rt
        self.label = label
        self.offset = 0

    @classmethod
    def fromim(cls, im):
        return cls(im.dst[0], im.others[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
        assert self.offset % 4 == 0
        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)


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


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


@instruction
class Mov3(ArmInstruction):
    """ mov Rd, imm8, move immediate value into register """
    mnemonic = 'mov'
    opcode = 4 # 00100 Rd(3) imm8
    operands = (Reg8Op, Imm8)
    def __init__(self, rd, imm):
        if type(imm) is int:
            imm = Imm8(imm)
        assert type(imm) is Imm8
        self.imm = imm.imm
        assert type(rd) is Reg8Op, str(type(rd))
        self.rd = rd

    @classmethod
    def fromim(cls, im):
        return cls(im.dst[0], im.others[0])

    def encode(self):
        rd = self.rd.num
        opcode = self.opcode
        imm8 = self.imm
        h = (opcode << 11) | (rd << 8) | imm8
        return u16(h)

    def __repr__(self):
        return 'MOV {}, {}'.format(self.rd, self.imm)



# Arithmatics:



class regregimm3_base(ArmInstruction):
    operands = (Reg8Op, Reg8Op, Imm3)
    def __init__(self, rd, rn, imm3):
        self.rd = rd
        self.rn = rn
        assert type(imm3) is Imm3
        self.imm3 = imm3

    @classmethod
    def fromim(cls, im):
        return cls(im.dst[0], im.src[0], im.others[0])

    def encode(self):
        rd = self.rd.num
        rn = self.rn.num
        imm3 = self.imm3.imm
        opcode = self.opcode
        h = (self.opcode << 9) | (imm3 << 6) | (rn << 3) | rd
        return u16(h)

    def __repr__(self):
        return '{} {}, {}, {}'.format(self.mnemonic, self.rd, self.rn, self.imm3.imm)

@instruction
class Add2(regregimm3_base):
    """ add Rd, Rn, imm3 """
    mnemonic = 'add'
    opcode = 0b0001110


@instruction
class Sub2(regregimm3_base):
    """ sub Rd, Rn, imm3 """
    mnemonic = 'sub'
    opcode = 0b0001111


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

    @classmethod
    def fromim(cls, im):
        return cls(im.dst[0], im.src[0], im.src[1])

    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)


@instruction
class Add(regregreg_base):
    mnemonic = 'ADD'
    opcode = 0b0001100


@instruction
class Sub(regregreg_base):
    mnemonic = 'SUB'
    opcode = 0b0001101


@instruction
class Mov2(ArmInstruction):
    """ mov rd, rm """
    operands = (ArmRegister, ArmRegister)
    mnemonic = 'MOV'
    def __init__(self, rd, rm):
        self.rd = rd
        self.rm = rm

    @classmethod
    def fromim(cls, im):
        return cls(im.dst[0], im.src[0])

    def encode(self):
        Rd = self.rd.num & 0x7
        D = (self.rd.num >> 3) & 0x1
        Rm = self.rm.num
        opcode = 0b01000110
        return u16((opcode << 8) | (D << 7) |(Rm << 3) | Rd)

    def __repr__(self):
        return '{} {}, {}'.format(self.mnemonic, self.rd, self.rm)
        

@instruction
class Mul(ArmInstruction):
    """ mul Rn, Rdm """
    operands = (Reg8Op, Reg8Op)
    mnemonic = 'MUL'
    def __init__(self, rn, rdm):
        self.rn = rn
        self.rdm = rdm

    @classmethod
    def fromim(cls, im):
        assert im.src[1] is im.dst[0]
        return cls(im.src[0], im.dst[0])

    def encode(self):
        rn = self.rn.num
        rdm = self.rdm.num
        opcode = 0b0100001101
        h = (opcode << 6) | (rn << 3) | rdm
        return u16(h)

    def __repr__(self):
        return '{} {}, {}'.format(self.mnemonic, self.rn, self.rdm)


class regreg_base(ArmInstruction):
    """ ??? Rdn, Rm """
    operands = (Reg8Op, Reg8Op)
    # TODO: integrate with the code gen interface:
    src = (0, 1)
    dst = (0,)
    def __init__(self, rdn, rm):
        self.rdn = rdn
        self.rm = rm

    @classmethod
    def fromim(cls, im):
        return cls(im.src[0], im.src[1])

    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)


@instruction
class movregreg_ins(regreg_base):
    """ mov Rd, Rm (reg8 operands) """
    mnemonic = 'mov'
    opcode = 0


@instruction
class And(regreg_base):
    mnemonic = 'AND'
    opcode = 0b0100000000


@instruction
class Orr(regreg_base):
    mnemonic = 'ORR'
    opcode = 0b0100001100


@instruction
class Cmp(regreg_base):
    mnemonic = 'CMP'
    opcode = 0b0100001010


@instruction
class Lsl(regreg_base):
    mnemonic = 'LSL'
    opcode = 0b0100000010


@instruction
class cmpregimm8_ins(ArmInstruction):
    """ cmp Rn, imm8 """
    mnemonic = 'cmp'
    opcode = 5 # 00101
    operands = (Reg8Op, 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)

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


@instruction
class B(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)


@instruction
class Bl(jumpBase_ins):
    mnemonic = 'BL'
    def encode(self):
        imm32 = wrap_negative(self.offset >> 1, 32)
        imm11 = imm32 & 0x7FF
        imm10 = (imm32 >> 11) & 0x3FF
        j1 = 1 # TODO: what do these mean?
        j2 = 1
        s = (imm32 >> 24) & 0x1
        h1 = (0b11110 << 11) | (s << 10) | imm10 
        h2 = (0b1101 << 12) | (j1 << 13) | (j2 << 11) | imm11
        return u16(h1) + u16(h2)


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)


@instruction
class Beq(cond_base_ins):
    mnemonic = 'beq'
    cond = 0


@instruction
class Bne(cond_base_ins):
    mnemonic = 'bne'
    cond = 1


@instruction
class Blt(cond_base_ins):
    mnemonic = 'blt'
    cond = 0b1011


@instruction
class Bgt(cond_base_ins):
    mnemonic = 'bgt'
    cond = 0b1100


@instruction
class Push(ArmInstruction):
    operands = (RegisterSet,)
    mnemonic = 'push'

    def __init__(self, regs):
        if type(regs) is set:
            regs = RegisterSet(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)


@instruction
class Pop(ArmInstruction):
    operands = (RegisterSet,)
    mnemonic = 'pop'

    def __init__(self, regs):
        if type(regs) is set:
            regs = RegisterSet(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
        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)


@instruction
class Yield(ArmInstruction):
    operands = ()
    mnemonic = 'yield'

    def encode(self):
        return u16(0xbf10)

# misc:

# add/sub SP:
class addspsp_base(ArmInstruction):
    operands = (RegSpOp, RegSpOp, Imm7)
    def __init__(self, _sp, _sp2, imm7):
        self.imm7 = imm7.imm
        assert self.imm7 % 4 == 0
        self.imm7 >>= 2

    def encode(self):
        return u16((self.opcode << 7) |self.imm7)

    def __repr__(self):
        return '{} sp, sp, {}'.format(self.mnemonic, self.imm7 << 2)

@instruction
class AddSp(addspsp_base):
    mnemonic = 'add'
    opcode = 0b101100000


@instruction
class SubSp(addspsp_base):
    mnemonic = 'sub'
    opcode = 0b101100001