import math
import multiprocessing
import random
import struct
from enum import IntEnum

from ..errors import TokenError
from .aes_gcm import encrypt, decrypt
from .reed_solomon import pad, unpad
from .seeded_shuffler import shuffle_bytes, unshuffle_bytes

"""
Every parameter but the key has a constant length due to its' algorithms restraints:
- algorithm: 1 byte
- ecc_block_size: 1 byte (0 <= ecc_block_size <= 255).
- ecc_symbol_num: 1 byte (0 < ecc_symbol_num < ecc_block_size).
- verification_tag: 16 bytes (hardcoded into aes encryption algorithm).
- nonce: 16 bytes (hardcoded into aes encryption algorithm).
- update_header: 8 bytes (hardcoded into aes encryption algorithm).
- shuffling seed: 64 bytes (hardcoded into shuffling algorithm).
- key: this will be concatenated to the end of the token so it will be the remaining bytes.

Algorithm specific parameters:
- min_alpha: 8 bytes (per pythons size for floating-point numbers).
- num_of_sacrificed_bits: 1 byte (ranges from 0-8).

From this we can know that the length of any token generated by this algorithm, will be at least the minimum length (
each case has a different length), if its less, then it's not a valid token.
"""


class Algorithms(IntEnum):
    LSB = 1
    BPCS = 2


def generic_wrap(plaintext: bytes, key: bytes, ecc_block_size: int,
                 ecc_symbol_num: int) -> tuple[bytes, bytes, bytes, bytes, bytes]:
    update_logger = multiprocessing.get_logger()

    update_logger.info("Encrypting data...")
    encrypted, verification_tag, nonce, update_header, key = encrypt(plaintext, key)

    update_logger.info("Shuffling data...")
    random.seed()
    seed = random.randbytes(64)
    random.seed(seed)
    shuffled = shuffle_bytes(encrypted, seed)

    update_logger.info("Padding data with error correction...")
    wrapped = pad(shuffled, ecc_block_size, ecc_symbol_num)

    return wrapped, verification_tag, nonce, update_header, seed


def wrap_bpcs(plaintext: bytes, key: bytes, ecc_block_size: int, ecc_symbol_num: int, min_alpha: float):
    """
    Wraps contents. First with encryption, then with error correction.
    :param plaintext: the plaintext to be wrapped
    :param key: the encryption key
    :param ecc_block_size: the error correction block size
    :param ecc_symbol_num: the bytes per block that we want to dedicate for error correction
    :param min_alpha: the minimum complexity coefficient
    :return: the wrapped plaintext, and the matching token
    :raises TokenError: if the created token is invalid
    """
    update_logger = multiprocessing.get_logger()
    update_logger.info("Wrapping data to be embedded with BPCS...")

    algorithm = Algorithms.BPCS.to_bytes(1)
    algorithm = struct.pack("c", algorithm)

    wrapped, verification_tag, nonce, update_header, seed = generic_wrap(plaintext, key, ecc_block_size, ecc_symbol_num)

    update_logger.info("Generating token...")
    token = (algorithm + ecc_block_size.to_bytes(1) + ecc_symbol_num.to_bytes(1) +
             verification_tag + nonce + update_header + seed + struct.pack("d", min_alpha) + key)

    if not len(token) == (115 + len(key)):
        raise TokenError("Token creating parameters are invalid.")

    update_logger.info("Wrapping completed!")
    return bytes(wrapped), token


def wrap_lsb(plaintext: bytes, key: bytes, ecc_block_size: int, ecc_symbol_num: int, num_of_sacrificed_bits: int):
    update_logger = multiprocessing.get_logger()
    update_logger.info("Wrapping data to be embedded with LSB...")

    algorithm = Algorithms.LSB.to_bytes(1)
    algorithm = struct.pack("c", algorithm)

    wrapped, verification_tag, nonce, update_header, seed = generic_wrap(plaintext, key, ecc_block_size, ecc_symbol_num)

    update_logger.info("Generating token...")
    token = (algorithm + ecc_block_size.to_bytes(1) + ecc_symbol_num.to_bytes(1) +
             verification_tag + nonce + update_header + seed +
             struct.pack("c", num_of_sacrificed_bits.to_bytes(1)) + key)

    if not len(token) == (108 + len(key)):
        raise TokenError("Token creating parameters are invalid.")

    update_logger.info("Wrapping completed!")
    return bytes(wrapped), token


def get_base_data_from_token(token: bytes):
    algorithm, token = ord(struct.unpack("c", token[:1])[0]), token[1:]
    ecc_block_size, token = ord(token[:1]), token[1:]
    ecc_symbol_num, token = ord(token[:1]), token[1:]
    verification_tag, token = token[:16], token[16:]
    nonce, token = token[:16], token[16:]
    update_header, token = token[:8], token[8:]
    seed, token = token[:64], token[64:]
    return (ecc_block_size, ecc_symbol_num), (verification_tag, nonce, update_header), seed, token


def get_bpcs_token_info(token: bytes):
    """
    Gets the wrapping and bpcs info from a given token.
    :param token: the token to infer data from
    :return: the token's info in the form:
    (ecc_block_size, ecc_symbol_num), (verification_tag, nonce, update_header, key), min_alpha
    :raises TokenError: if the token is invalid
    """
    update_logger = multiprocessing.get_logger()
    update_logger.info("Extracting data from BPCS token...")
    if not len(token) >= 115:
        raise TokenError("Token length is invalid. Token length must be equal to or greater than 42 bytes.")

    ((ecc_block_size, ecc_symbol_num),
     (verification_tag, nonce, update_header), seed, token) = get_base_data_from_token(token)

    min_alpha, token = struct.unpack("d", token[:8])[0], token[8:]
    key = token
    update_logger.info("Data extraction completed!")
    return (ecc_block_size, ecc_symbol_num), (verification_tag, nonce, update_header, key), seed, min_alpha


def get_lsb_token_info(token: bytes):
    update_logger = multiprocessing.get_logger()
    update_logger.info("Extracting data from LSB token...")

    if not len(token) >= 108:
        raise TokenError("Token length is invalid. Token length must be equal to or greater than 43 bytes.")

    ((ecc_block_size, ecc_symbol_num),
     (verification_tag, nonce, update_header), seed, token) = get_base_data_from_token(token)
    num_of_sacrificed_bits, token = int.from_bytes(struct.unpack("c", token[:1])[0]), token[1:]

    key = token
    update_logger.info("Data extraction completed!")
    return (ecc_block_size, ecc_symbol_num), (verification_tag, nonce, update_header, key), seed, num_of_sacrificed_bits


def unwrap(wrapped: bytes, ecc_block_size: int, ecc_symbol_num: int, verification_tag: bytes, nonce: bytes,
           update_header: bytes, seed: bytes, key: bytes):
    """
    Unwraps content, using the following wrapping parameters.
    :param wrapped: the wrapped content
    :param ecc_block_size: the rs block size
    :param ecc_symbol_num: the number of bytes per block that we want to dedicate for error correction
    :param verification_tag: the GCM tag
    :param nonce: the GCM nonce
    :param update_header: the GCM header
    :param seed: the shuffling seed
    :param key: the encryption key
    :return: the unwrapped content
    :raises ReedSolomonError: if the function was fed invalid parameters, or the data had too many errors for the rs
    algorithm to fix
    :raises ContentWrapperError: if the decryption process fails
    """
    update_logger = multiprocessing.get_logger()
    update_logger.info("Unwrapping extracted data...")

    update_logger.info("Unpadding data...")
    shuffled = unpad(wrapped, ecc_block_size, ecc_symbol_num)

    update_logger.info("Unshuffling data...")
    ciphertext = unshuffle_bytes(shuffled, seed)

    update_logger.info("Decrypting data...")
    dec = decrypt(ciphertext, key, verification_tag, nonce, update_header)

    update_logger.info("Unwrapping completed!")
    return dec


def get_max_unwapped_length(max_bit_embbeding_input_length: int, ecc_block_size: int, ecc_symbol_num: int) -> int:
    return math.floor(
        ((max_bit_embbeding_input_length * (ecc_block_size - ecc_symbol_num))
         /
         (8 * ecc_block_size))
        - 16
    )
