from .constants import BLOCK_BITS, BLOCK_BYTES
from . import tbc


def bytes_to_block_matrix(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 block_matrix_to_bytes(matrix):
    return bytes(byte for block in matrix for byte in block)


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


def pad10(X):
    zeroes = [0] * (BLOCK_BYTES-len(X)-1)
    return zeroes + [0b10000000] + X


def _tweak_associated_data(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 build_auth(t, A, key):
    Auth = [0 for byte in range(0, BLOCK_BYTES)]
    l_a = len(A)//BLOCK_BYTES
    need_padding = len(A)%BLOCK_BYTES > 0

    A = bytes_to_block_matrix(A)

    for i in range(0, l_a):
        tweak = _tweak_associated_data(t, i, padded=False)
        enc = tbc.encrypt(tweak, key, A[i])
        Auth = xor_state(Auth, enc)

    if not need_padding:
        return Auth

    tweak = _tweak_associated_data(t, l_a, padded=True)
    ad_padded = pad10(A[l_a])
    enc = tbc.encrypt(tweak, key, ad_padded)
    Auth = xor_state(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