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

# TODO: encode this in DSL (domain specific language)
# TBD: is this required?

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

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

armtarget = Target('arm')

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 armtarget.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


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 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 == '[]' 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]

def makeReg(cls, num, name):
    r = cls(num, name)
    armtarget.registers.append(r)
    return r

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

# Sanity checks:
assert isinstance(sp, ArmRegister)
assert isinstance(r3, ArmRegister)
assert ArmRegister.Create(ASymbol('r3')) is r3
assert ArmRegister.Create(ASymbol('sp')) is sp


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


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


@armtarget.instruction
class storeimm5_ins(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)


@armtarget.instruction
class loadimm5_ins(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


@armtarget.instruction
class ldr_pcrel(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)


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

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


@armtarget.instruction
class subregregimm3_ins(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)


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


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



@armtarget.instruction
class movregreg_ext_ins(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)
        

@armtarget.instruction
class mulregreg_ins(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)


@armtarget.instruction
class movregreg_ins(regreg_base):
    """ mov Rd, Rm (reg8 operands) """
    # TODO: match this:
    pattern = ir.Move(ir.Temp, ir.Temp)
    mnemonic = 'mov'
    opcode = 0


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


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


@armtarget.instruction
class bl_ins(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)


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


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


@armtarget.instruction
class blt_ins(cond_base_ins):
    mnemonic = 'blt'
    cond = 0b1011


@armtarget.instruction
class bgt_ins(cond_base_ins):
    mnemonic = 'bgt'
    cond = 0b1100


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

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

@armtarget.instruction
class addspsp_ins(addspsp_base):
    mnemonic = 'add'
    opcode = 0b101100000
    

@armtarget.instruction
class subspsp_ins(addspsp_base):
    mnemonic = 'sub'
    opcode = 0b101100001

armtarget.check()