from constants import BLOCK_BITS, BLOCK_BYTES
from lilliput_tbc import LilliputTBCEnc


def ArrayToBlockbytesMatrix(array) :
    length = len(array)
    pad = 0
    if(length % BLOCK_BYTES == 0) :
        number_blocks = int(length / BLOCK_BYTES)
    else :
        number_blocks = int((length + (BLOCK_BYTES - (length % BLOCK_BYTES))) / BLOCK_BYTES)
        pad = 1

    matrix = [[0] * BLOCK_BYTES for block in range(0, number_blocks - pad)]
    if(pad == 1) :
        matrix.append([0] * (length % BLOCK_BYTES))

    for byte in range(0, length) :
        matrix[int(byte / BLOCK_BYTES)][byte % BLOCK_BYTES] = array[byte]

    return matrix


def BlockbytesMatrixToBytes(matrix):
    return bytes(byte for block in matrix for byte in block)


def XorState(state1, state2):
    return [s1^s2 for (s1, s2) in zip(state1, state2)]


def Padding10LSB(array, number_bits):
    shifted = 0
    for byte in range(0, len(array)):
        shifted |= (array[byte] << (8 * byte))
    shifted = (shifted << (BLOCK_BITS - number_bits)) & 0xffffffffffffffffffffffffffffffff

    padded = shifted | (0x1 << (BLOCK_BITS - number_bits - 1))

    array_padded = [0 for byte in range(0, BLOCK_BYTES)]
    for byte in range(0, BLOCK_BYTES):
        array_padded[byte] = (padded & (0xff << (8 * byte))) >> (8 * byte)

    return array_padded


def _tweakAssociatedData(t, i, padded):
    t_bytes = t//8
    tweak = [0]*(t_bytes)

    mask = 0xff
    for byte in range(t_bytes-1):
        tweak[byte] = (i & mask) >> (byte * 8)
        mask = mask << 8

    mask = (0xf << (8 * t_bytes-1))
    tweak[-1] = (i & mask) >> ((t_bytes-1)*8)
    if not padded:
        tweak[-1] |= 0x20
    else:
        tweak[-1] |= 0x60

    return tweak


def BuildAuth(t, A, key):
    Auth = [0 for byte in range(0, BLOCK_BYTES)]
    l_a = len(A)//BLOCK_BYTES

    padding_bytes = len(A)%BLOCK_BYTES

    A = ArrayToBlockbytesMatrix(A)

    for i in range(0, l_a):
        tweak = _tweakAssociatedData(t, i, padded=False)
        enc = LilliputTBCEnc(tweak, key, A[i])
        Auth = XorState(Auth, enc)

    if padding_bytes == 0:
        return Auth

    tweak = _tweakAssociatedData(t, l_a, padded=True)
    ad_padded = Padding10LSB(A[l_a], padding_bytes*8)
    enc = LilliputTBCEnc(tweak, key, ad_padded)
    Auth = XorState(Auth, enc)

    return Auth