"""
    SCT 2 for lilliput ae 2
"""

import lilliput_tbc as ltbc
from constants import BLOCK_BITS, BLOCK_BYTES
from helpers import (
    ArrayToBlockbytesMatrix,
    BlockbytesMatrixToBytes,
    BuildAuth,
    Padding10LSB,
    XorState
)


TWEAK_BITS = 128
TWEAK_BYTES = TWEAK_BITS//8


def TweakTag(j, padded):
        tweak = [0 for byte in range(0,  TWEAK_BYTES)]

        tweak[TWEAK_BYTES - 1] |= ((j >> 120) & 0xf)
        for byte in range(TWEAK_BYTES - 2, -1, -1) :
            tweak[byte] = (j >> (8 * byte)) & 0xff

        if padded:
            tweak[TWEAK_BYTES - 1] |= 0x40

        return tweak

def TweakTagEnd(N) :
    tweak = [0 for byte in range(0,  TWEAK_BYTES)]

    for byte in range(0, TWEAK_BYTES - 1) :
        tweak[byte] = N[byte]
    tweak[TWEAK_BYTES - 1] = 0x10

    return tweak


def AddTagJ(tag, j) :
    array_j = [0 for byte in range(0, TWEAK_BYTES)]
    for byte in range(0, TWEAK_BYTES) :
        array_j[byte] = (j >> (byte * 8))

    xorr = XorState(tag, array_j)

    xorr[TWEAK_BYTES - 1] |= 0x80

    return xorr


def MesssageAuthTag(M, N, Auth, key) :
    l = len(M)//BLOCK_BYTES
    padding_bytes = len(M)%BLOCK_BYTES

    tag = list(Auth)
    M = ArrayToBlockbytesMatrix(M)

    for j in range(0, l) :
        tweak = TweakTag(j, False)
        encryption = ltbc.LilliputTBCEnc(tweak, key, M[j])
        tag = XorState(tag, encryption)

    if padding_bytes > 0 :
        tweak = TweakTag(l, True)
        m_padded = Padding10LSB(M[l], 8*padding_bytes)
        encryption = ltbc.LilliputTBCEnc(tweak, key, m_padded)
        tag = XorState(tag, encryption)

    tweak = TweakTagEnd(N)
    encryption = ltbc.LilliputTBCEnc(tweak, key, tag)
    tag = encryption

    return tag


def MessageEncryption(M, N, tag, key) :
    l = len(M)//BLOCK_BYTES
    padding_bytes = len(M)%BLOCK_BYTES

    M = ArrayToBlockbytesMatrix(M)
    C = []

    for j in range(0, l) :
        tweak = AddTagJ(tag, j)
        padded_nounce = list(N) + [0x00]
        encryption = ltbc.LilliputTBCEnc(tweak, key, padded_nounce)
        C.append(XorState(M[j], encryption))

    if padding_bytes > 0:
        tweak = AddTagJ(tag, l)
        padded_nounce = list(N) + [0x00]
        encryption = ltbc.LilliputTBCEnc(tweak, key, padded_nounce)
        C.append(XorState(M[l], encryption))

    return C

################################################################################
def encrypt(A, M, N, key) :
    K = list(key)

    Auth = BuildAuth(TWEAK_BITS, A, K)
    tag = MesssageAuthTag(M, N, Auth, K)
    C = MessageEncryption(M, N, tag, K)

    return BlockbytesMatrixToBytes(C), bytes(tag)


def decrypt(A, C, N, tag, key) :
    K = list(key)
    tag = list(tag)

    M = BlockbytesMatrixToBytes(
        MessageEncryption(C, N, tag, K)
    )
    Auth = BuildAuth(TWEAK_BITS, A, K)
    tag2 = MesssageAuthTag(M, N, Auth, K)

    if tag == tag2:
        return M