import struct
from random import randint

from ..helpers.checksumHelper import CheckSum
from ..network.enums import types as TYPE

TEA_ROUNDS      = 32
TEA_XOR_MASK    = 0x5E94A3CF
KEY             = []

def EncryptBlock(v0, v1):
    sum = 0
    
    for i in range(TEA_ROUNDS):
        v0 += ((v1 << 4 ^ v1 >> 5) + v1) ^ (sum + key[sum & 3])

def Encrypt(var_input): # bytes
    var_input_len = len(var_input) + TYPE.SIZEOF_CHAR * 2 # (checksum, encodedPad)
    paddingBytes = 0
    
    if var_input_len % 8 != 0:
        paddingBytes = 8 - (var_input_len % 8)
    
    finalLen = var_input_len + paddingBytes
    
    # Randomize pad size for some reason
    encodedPad = randint(0, TYPE.MAX_UCHAR)
    encodedPad <<= 4
    encodedPad &= 0xF0 # Remove bits lost from UCHAR size (C) [BYTE MAGIC]
    encodedPad |= paddingBytes
    
    var_output = [0] * finalLen #bytes(finalLen)
    var_output[2 + paddingBytes:] = [b for b in var_input]
    
    # Write the pad size variable
    var_output[1] = encodedPad
    
    # Write the padding
    var_output[2:2+len(paddingBytes)] = bytes([randint(0, TYPE.MAX_UCHAR) for _ in range(len(paddingBytes))])
    
    # Calculate checksum
    checksum = CheckSum()
    checksum.add(var_output[1:var_input_len + 1 + paddingBytes])
    
    var_checksum = checksum.get()
    var_checksum = (var_checksum << 4) ^ var_checksum
    var_checksum &= TYPE.MAX_UCHAR # [BYTE MAGIC]
    
    var_output[0] = var_checksum
    
    # Encryption
    for i in range(0, finalLen, 8):
        v0 = var_output[i]
        v1 = var_output[i + TYPE.SIZEOF_INT]
    
    return bytes(var_output)