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

import lilliput_tbc as ltbc
from helpers import ArrayToBlockbytesMatrix, BlockbytesMatrixToBytes


BLOCK_BITS = 128
KEY_BITS = 128
TWEAK_BITS = 128
TWEAKEY_BITS = KEY_BITS + TWEAK_BITS
LANE_BITS = 64
LANES = int((TWEAKEY_BITS) / LANE_BITS)
ROUNDS = 32

BLOCK_BYTES = int(BLOCK_BITS / 8)
KEY_BYTES = int(KEY_BITS / 8)
TWEAK_BYTES = int(TWEAK_BITS / 8)
TWEAKEY_BYTES = int(TWEAKEY_BITS / 8)

A_BITS = BLOCK_BITS
M_BITS = BLOCK_BITS
N_BITS = 120
N_BYTES = int(N_BITS / 8)


def InitParameters(key_bits = 128, tweak_bits = 128, rounds = 32) :
    global KEY_BITS
    global KEY_BYTES
    global TWEAK_BITS
    global TWEAK_BYTES
    global TWEAKEY_BITS
    global TWEAKEY_BYTES
    global LANES
    global ROUNDS

    KEY_BITS = key_bits
    TWEAK_BITS = tweak_bits
    TWEAKEY_BITS = KEY_BITS + TWEAK_BITS
    LANES = int((TWEAKEY_BITS) / LANE_BITS)
    ROUNDS = rounds
    KEY_BYTES = int(KEY_BITS / 8)
    TWEAK_BYTES = int(TWEAK_BITS / 8)
    TWEAKEY_BYTES = int(TWEAKEY_BITS / 8)


###############################################################################

def XorState(state1, state2) :
    state_output = [state1[byte] ^ state2[byte] for byte in range(0, len(state1))]
    return state_output

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 LowPart(array, number_bits) :
    shifted = 0
    for byte in range(0, len(array)) :
        shifted |= (array[byte] << (8 * byte))

    mask = 0
    for bit in range(0, number_bits) :
        mask |= (0x1 << bit)

    lower_part = shifted & mask

    will_padd = 0
    if (number_bits % 8) != 0 :
        will_padd = 1

    lower_part_byte = [0 for byte in range(0, int(number_bits / 8)  + will_padd)]
    for byte in range(0, int(number_bits / 8)  + will_padd) :
        lower_part_byte[byte] = lower_part & 0xff
        lower_part = lower_part >> 8

    return lower_part_byte

###############################################################################

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

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

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


def BuildAuth(A, key) :
    Auth = [0 for byte in range(0, BLOCK_BYTES)]
    l_a = int(A_BITS / BLOCK_BITS)
    if int(A_BITS % BLOCK_BITS) > 0 :
        will_padd = 1
    else :
        will_padd = 0

    for i in range(0, l_a) :
        tweak = TweakAssociatedData(i, padded = 0)
        enc = ltbc.LilliputTBCEnc(tweak, key, A[i])
        Auth = XorState(Auth, enc)

    if (A_BITS % BLOCK_BITS) == 0 :
        return Auth

    tweak = TweakAssociatedData(l_a, padded = 1)
    ad_padded = Padding10LSB(A[l_a], (A_BITS % BLOCK_BITS))
    enc = ltbc.LilliputTBCEnc(tweak, key, ad_padded)
    Auth = XorState(Auth, enc)

    return Auth

################################################################################

def TweakTag(j, padded = 0) :
        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 == 1 :
            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 = int(M_BITS / BLOCK_BITS)
    if int(M_BITS % BLOCK_BITS) > 0 :
        will_padd = 1
    else :
        will_padd = 0

    tag = list(Auth)

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

    if will_padd == 1 :
        tweak = TweakTag(l, padded = 1)
        m_padded = Padding10LSB(M[l], M_BITS % BLOCK_BITS)
        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 = int(M_BITS / BLOCK_BITS)
    if int(M_BITS % BLOCK_BITS) > 0 :
        will_padd = 1
    else :
        will_padd = 0

    C = [[0 for byte in range(0, 16)] for j in range(0, l + will_padd)]

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

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

    return C

################################################################################
def SCT2Enc(A, M, N, key, tweak_bits, rounds) :
    InitParameters(len(key)*8, tweak_bits, rounds)

    global A_BITS
    global M_BITS

    A_BITS = len(A)*8
    M_BITS = len(M)*8

    A = ArrayToBlockbytesMatrix(A)
    M = ArrayToBlockbytesMatrix(M)

    ltbc.KEY_BITS = KEY_BITS
    ltbc.ROUNDS = ROUNDS
    ltbc.TWEAK_BITS = TWEAK_BITS
    ltbc.LANES = LANES

    ltbc.TWEAKEY_BITS = TWEAKEY_BITS
    ltbc.KEY_BYTES = KEY_BYTES
    ltbc.TWEAK_BYTES = TWEAK_BYTES
    ltbc.TWEAKEY_BYTES = TWEAKEY_BYTES

    ltbc.TKs = [[0 for byte in range(0, TWEAKEY_BYTES)] for round in range(0, ROUNDS)]
    ltbc.RTKs = [[0 for byte in range(0, 8)] for round in range(0, ROUNDS)]
    ltbc.States = [[0 for byte in range(0, BLOCK_BYTES)] for round in range(0, ROUNDS)]


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

    return BlockbytesMatrixToBytes(C), bytes(tag)


def SCT2Dec(A, C, N, tag, key, tweak_bits, rounds) :
    InitParameters(len(key)*8, tweak_bits, rounds)

    global A_BITS
    global M_BITS

    A_BITS = len(A)*8
    M_BITS = len(C)*8

    A = ArrayToBlockbytesMatrix(A)
    C = ArrayToBlockbytesMatrix(C)

    ltbc.KEY_BITS = KEY_BITS
    ltbc.ROUNDS = ROUNDS
    ltbc.TWEAK_BITS = TWEAK_BITS
    ltbc.LANES = LANES

    ltbc.TWEAKEY_BITS = TWEAKEY_BITS
    ltbc.KEY_BYTES = KEY_BYTES
    ltbc.TWEAK_BYTES = TWEAK_BYTES
    ltbc.TWEAKEY_BYTES = TWEAKEY_BYTES

    ltbc.TKs = [[0 for byte in range(0, TWEAKEY_BYTES)] for round in range(0, ROUNDS)]
    ltbc.RTKs = [[0 for byte in range(0, 8)] for round in range(0, ROUNDS)]
    ltbc.States = [[0 for byte in range(0, BLOCK_BYTES)] for round in range(0, ROUNDS)]


    M = MessageEncryption(C, N, tag, key)
    Auth = BuildAuth(A, key)
    tag2 = MesssageAuthTag(M, N, Auth, key)

    if(tag == tag2) :
        return BlockbytesMatrixToBytes(M)