import os

from cryptography.hazmat.primitives.ciphers import Cipher, CipherContext, algorithms, modes
from cryptography.hazmat.backends.openssl import backend as openssl_backend

from . import Writer


# ====================================================================================================================
# Use this class to create a PngBin image file with encrypted data.
#
# This class is derived from Writer class, So consult Its documentation for parameter details.
# The exception for the parameters are as below:
#   The multiple of `width` and `height` must be divisible by 4.
#   If `encryptor` is provided, It must be an AES cipher and CBC mode with key length of 256 bits and iv of 128 bits.
# ====================================================================================================================
class EncryptWriter(Writer):
    def __init__(self, width: int, height: int, fobj, auto_finish: bool = False, encryptor: CipherContext = None):
        """ Creates an encrypted PngBin writer instance.

        :param encryptor:
            a CipherContext type. if None, generates a new one with random key and random initialization vector (IV).
        """
        if width * height % 4 != 0:
            raise ValueError('The multiple of `width` and `height` must be divisible by 4.')
        if encryptor:
            if isinstance(encryptor, CipherContext):
                self._encryptor = encryptor
                self._key, self._iv = None, None

                # Access protected class members to gain key and iv info.
                # self._key = encryptor._ctx._cipher.key
                # self._iv = encryptor._ctx._mode.initialization_vector
            else:
                raise TypeError('"encryptor" must be a CipherContext type.')
        else:
            self._encryptor = self._gen_encryptor()

        super().__init__(width, height, fobj, auto_finish)
        self.__left = self._left

    @property
    def key(self) -> bytes:
        """ a cipher key that is generated by this class.
            returns None if you provided an encryptor when creating this class. """
        return self._key

    @property
    def iv(self) -> bytes:
        """ a cipher initialization vector (iv) that is generated by this class.
            returns None if you provided an encryptor when creating this class. """
        return self._iv

    @property
    def bytes_left(self) -> int:
        return self.__left

    def tell(self):
        return self._width * self._height * 4 - self.__left

    def write(self, data: (bytes, bytearray)) -> int:
        """ encrypts and writes input bytes-like object `data` to file.

        :param data: bytes-like object
        :return:
            length in bytes that has been written.
            This is guarantee to be the same as `len(data)` if underlying `fobj.write(data)` is guarantee.
            Unless `bytes_left` is 0 after the call.
        """
        # Developer note:
        #     Do not worry about data left inside `self._encryptor` for the last call,
        #     because when you call `finish()`, all data inside will be written with 0 paddings (if required).

        m = len(data)
        n = min(self.__left, m)
        self.__left -= n

        b = self._encryptor.update(data if m == n else data[:n])
        super().write(b)
        return n

    def _gen_encryptor(self) -> CipherContext:
        """ Generates an encryptor. """
        self._key, self._iv = os.urandom(32), os.urandom(16)
        cipher = Cipher(algorithms.AES(self._key), modes.CBC(self._iv), backend=openssl_backend)
        return cipher.encryptor()
