from ..basetarget import Instruction

from .token import ArmToken
from .registers import R0, SP, ArmRegister

# Helpers:

def rotate_right(v, n):
    """ bit-wise Rotate right n times """
    mask = (2**n) - 1
    mask_bits = v & mask
    return (v >> n) | (mask_bits << (32 - n))

def rotate_left(v, n):
    assert n >= 0
    assert n < 32
    return rotate_right(v, 32 - n)

masks = [rotate_right(0xFF, i * 2) for i in range(16)]
#0x000000FF,
#0xC000007F,
#0xF000000F,
#0xFC000003,
#0xFF000000, # 4

assert masks[4] == 0xFF000000, hex(masks[4])

def encode_imm32(v):
    """ Bundle 32 bit value into 4 bits rotation and 8 bits value
     """
    for i in range(0, 16):
        v2 = rotate_left(v, i*2)
        if (v2 & 0xFFFFFF00) == 0:
            rotation = i
            val = v2 & 0xFF
            x = (rotation << 8) | val
            return x
    raise Exception("Invalid value {}".format(v))

# Instructions:

class ArmInstruction(Instruction):
    def __init__(self):
        self.token = ArmToken()


class Mov(ArmInstruction):
    """ Mov Rd, imm16 """
    def __init__(self, reg, imm):
        super().__init__()
        self.reg = reg
        self.imm = imm

    def encode(self):
        self.token[0:12] = self.imm
        self.token.Rd = self.reg.num
        self.token[16:20] = 0
        self.token[20] = 0
        self.token[21:28] = 0b0011101
        self.token.cond = 0xE # Always!
        return self.token.encode()

    def relocations(self):
        return []

    def __repr__(self):
        return 'Mov {}, {}'.format(self.reg, self.imm)


def Add(*args):
    if len(args) == 3 and isinstance(args[0], ArmRegister) and \
            isinstance(args[1], ArmRegister):
        if isinstance(args[2], ArmRegister):
            return Add1(args[0], args[1], args[2])
        elif isinstance(args[2], int):
            return Add2(args[0], args[1], args[2])
    raise Exception()

def Sub(*args):
    if len(args) == 3 and isinstance(args[0], ArmRegister) and \
            isinstance(args[1], ArmRegister):
        if isinstance(args[2], ArmRegister):
            return Sub1(args[0], args[1], args[2])
        elif isinstance(args[2], int):
            return Sub2(args[0], args[1], args[2])
    raise Exception()

class OpRegRegReg(ArmInstruction):
    """ add rd, rn, rm """
    def __init__(self, rd, rn, rm, shift=0):
        super().__init__()
        self.rd = rd
        self.rn = rn
        self.rm = rm

    def encode(self):
        self.token[0:4] = self.rm.num
        self.token[4] = 0
        self.token[5:7] = 0
        self.token[7:12] = 0 # Shift
        self.token.Rd = self.rd.num
        self.token.Rn = self.rn.num
        self.token.S = 0 # Set flags
        self.token[21:28] = self.opcode
        self.token.cond = 0xE # Always!
        return self.token.encode()

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


class Add1(OpRegRegReg):
    opcode = 0b0000100


class Sub1(OpRegRegReg):
    opcode = 0b0000010


class Orr1(OpRegRegReg):
    opcode = 0b0001100


class OpRegRegImm(ArmInstruction):
    """ add rd, rn, imm12 """
    def __init__(self, rd, rn, imm):
        super().__init__()
        self.rd = rd
        self.rn = rn
        self.imm2 = encode_imm32(imm)
        self.imm = imm

    def encode(self):
        self.token[0:12] = self.imm2
        self.token.Rd = self.rd.num
        self.token.Rn = self.rn.num
        self.token.S = 0 # Set flags
        self.token[21:28] = self.opcode
        self.token.cond = 0xE # Always!
        return self.token.encode()

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


class Add2(OpRegRegImm):
    opcode = 0b0010100


class Sub2(OpRegRegImm):
    opcode = 0b0010010



# Branches:

class BranchBaseRoot(ArmInstruction):
    def __init__(self, target):
        super().__init__()
        self.target = target

    def encode(self):
        self.token.cond = self.cond
        self.token[24:28] = self.opcode
        return self.token.encode()

    def relocations(self):
        return [(self.target, 'b_imm24')]


class BranchBase(BranchBaseRoot):
    opcode = 0b1010

class BranchLinkBase(BranchBaseRoot):
    opcode = 0b1011

class Bl(BranchLinkBase):
    cond = 0xE

class B(BranchBase):
    cond = 0xE

class Beq(BranchBase):
    cond = 0x0

class Bgt(BranchBase):
    cond = 0xC

class Ble(BranchBase):
    cond = 0xD