from constants import BLOCK_BITS, BLOCK_BYTES
from lilliput_tbc import LilliputTBCEnc


def ArrayToBlockbytesMatrix(array):
    vector = list(array)

    blocks_nb = len(vector)//BLOCK_BYTES

    block_starts = (
        i*BLOCK_BYTES for i in range(blocks_nb)
    )

    matrix = [
        vector[start:start+BLOCK_BYTES] for start in block_starts
    ]

    padding_len = len(vector)%BLOCK_BYTES

    if padding_len > 0:
        padding = vector[-padding_len:]
        matrix.append(padding)

    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


class TagValidationError(Exception):
    def __init__(self, announced, computed):
        msg = '\n'.join((
            'Invalid tag:',
            announced.hex().upper()+' (announced)',
            computed.hex().upper()+' (computed)'
        ))

        super().__init__(msg)
        self._announced = announced
        self._computed = computed