#  Part of the Concrete Compiler Project, under the BSD3 License with Zama Exceptions.
#  See https://github.com/zama-ai/concrete/blob/main/LICENSE.txt for license information.

"""Client support."""
from typing import List, Optional, Union, Dict
import numpy as np

# pylint: disable=no-name-in-module,import-error
from mlir._mlir_libs._concretelang._compiler import ClientSupport as _ClientSupport

# pylint: enable=no-name-in-module,import-error

from .public_result import PublicResult
from .key_set import KeySet
from .key_set_cache import KeySetCache
from .client_parameters import ClientParameters
from .public_arguments import PublicArguments
from .lambda_argument import LambdaArgument
from .lwe_secret_key import LweSecretKey
from .wrapper import WrapperCpp
from .utils import ACCEPTED_INTS, ACCEPTED_NUMPY_UINTS, ACCEPTED_TYPES


class ClientSupport(WrapperCpp):
    """Client interface for doing key generation and encryption.

    It provides features that are needed on the client side:
    - Generation of public and private keys required for the encrypted computation
    - Encryption and preparation of public arguments, used later as input to the computation
    - Decryption of public result returned after execution
    """

    def __init__(self, client_support: _ClientSupport):
        """Wrap the native Cpp object.

        Args:
            client_support (_ClientSupport): object to wrap

        Raises:
            TypeError: if client_support is not of type _ClientSupport
        """
        if not isinstance(client_support, _ClientSupport):
            raise TypeError(
                f"client_support must be of type _ClientSupport, not {type(client_support)}"
            )
        super().__init__(client_support)

    # pylint: disable=arguments-differ
    @staticmethod
    def new() -> "ClientSupport":
        """Build a ClientSupport.

        Returns:
            ClientSupport
        """
        return ClientSupport.wrap(_ClientSupport())

    # pylint: enable=arguments-differ

    @staticmethod
    def key_set(
        client_parameters: ClientParameters,
        keyset_cache: Optional[KeySetCache] = None,
        secret_seed: Optional[int] = None,
        encryption_seed: Optional[int] = None,
        initial_lwe_secret_keys: Optional[Dict[int, LweSecretKey]] = None,
    ) -> KeySet:
        """Generate a key set according to the client parameters.

        If the cache is set, and include equivalent keys as specified by the client parameters,
        the keyset is loaded, otherwise, a new keyset is generated and saved in the cache.
        If keygen is required, it will first initialize the secret keys provided, if any.

        Args:
            client_parameters (ClientParameters): client parameters specification
            keyset_cache (Optional[KeySetCache], optional): keyset cache. Defaults to None.
            secret_seed (Optional[int]): secret seed, must be a positive 128 bits integer
            encryption_seed (Optional[int]): encryption seed, must be a positive 128 bits integer
            initial_lwe_secret_keys (Optional[Dict[int, LweSecretKey]]): keys to init the keyset
                with before keygen. It maps keyid to secret key

        Raises:
            TypeError: if client_parameters is not of type ClientParameters
            TypeError: if keyset_cache is not of type KeySetCache
            AssertionError: if seed components is not uint64

        Returns:
            KeySet: generated or loaded keyset
        """
        secret_seed = 0 if secret_seed is None else secret_seed
        encryption_seed = 0 if encryption_seed is None else encryption_seed

        if secret_seed < 0 or secret_seed >= 2**128:
            raise ValueError("secret_seed must be a positive 128 bits integer")
        if encryption_seed < 0 or encryption_seed >= 2**128:
            raise ValueError("encryption_seed must be a positive 128 bits integer")
        secret_seed_msb = (secret_seed >> 64) & 0xFFFFFFFFFFFFFFFF
        secret_seed_lsb = (secret_seed) & 0xFFFFFFFFFFFFFFFF
        encryption_seed_msb = (encryption_seed >> 64) & 0xFFFFFFFFFFFFFFFF
        encryption_seed_lsb = (encryption_seed) & 0xFFFFFFFFFFFFFFFF

        if keyset_cache is not None and not isinstance(keyset_cache, KeySetCache):
            raise TypeError(
                f"keyset_cache must be None or of type KeySetCache, not {type(keyset_cache)}"
            )

        if initial_lwe_secret_keys is not None:
            intial_sks = {i: sk.cpp() for i, sk in initial_lwe_secret_keys.items()}
        else:
            intial_sks = {}

        cpp_cache = None if keyset_cache is None else keyset_cache.cpp()
        return KeySet.wrap(
            _ClientSupport.key_set(
                client_parameters.cpp(),
                cpp_cache,
                secret_seed_msb,
                secret_seed_lsb,
                encryption_seed_msb,
                encryption_seed_lsb,
                intial_sks,
            ),
        )

    @staticmethod
    def encrypt_arguments(
        client_parameters: ClientParameters,
        keyset: KeySet,
        args: List[Union[int, np.ndarray]],
        circuit_name: str,
    ) -> PublicArguments:
        """Prepare arguments for encrypted computation.

        Pack public arguments by encrypting the ones that requires encryption, and leaving the rest as plain.
        It also pack public materials (public keys) that are required during the computation.

        Args:
            client_parameters (ClientParameters): client parameters specification
            keyset (KeySet): keyset used to encrypt arguments that require encryption
            args (List[Union[int, np.ndarray]]): list of scalar or tensor arguments
            circuit_name(str): the name of the circuit for which to encrypt

        Raises:
            TypeError: if client_parameters is not of type ClientParameters
            TypeError: if keyset is not of type KeySet
            TypeError: if circuit_name is not of type str

        Returns:
            PublicArguments: public arguments for execution
        """
        if not isinstance(client_parameters, ClientParameters):
            raise TypeError(
                f"client_parameters must be of type ClientParameters, not {type(client_parameters)}"
            )
        if not isinstance(keyset, KeySet):
            raise TypeError(f"keyset must be of type KeySet, not {type(keyset)}")
        if not isinstance(circuit_name, str):
            raise TypeError(
                f"circuit_name must be of type str, not {type(circuit_name)}"
            )

        signs = client_parameters.input_signs()
        if len(signs) != len(args):
            raise RuntimeError(
                f"function has arity {len(signs)} but is applied to too many arguments"
            )

        lambda_arguments = [
            ClientSupport._create_lambda_argument(arg, signed)
            for arg, signed in zip(args, signs)
        ]
        return PublicArguments.wrap(
            _ClientSupport.encrypt_arguments(
                client_parameters.cpp(),
                keyset.cpp(),
                [arg.cpp() for arg in lambda_arguments],
                circuit_name,
            )
        )

    @staticmethod
    def decrypt_result(
        client_parameters: ClientParameters,
        keyset: KeySet,
        public_result: PublicResult,
        circuit_name: str,
    ) -> Union[int, np.ndarray]:
        """Decrypt a public result using the keyset.

        Args:
            client_parameters (ClientParameters): client parameters for decryption
            keyset (KeySet): keyset used for decryption
            public_result (PublicResult): public result to decrypt
            circuit_name (str): name of the circuit for which to decrypt

        Raises:
            TypeError: if keyset is not of type KeySet
            TypeError: if public_result is not of type PublicResult
            TypeError: if circuit_name is not of type str
            RuntimeError: if the result is of an unknown type

        Returns:
            Union[int, np.ndarray]: plain result
        """
        if not isinstance(keyset, KeySet):
            raise TypeError(f"keyset must be of type KeySet, not {type(keyset)}")
        if not isinstance(public_result, PublicResult):
            raise TypeError(
                f"public_result must be of type PublicResult, not {type(public_result)}"
            )
        if not isinstance(circuit_name, str):
            raise TypeError(
                f"circuit_name must be of type str, not {type(circuit_name)}"
            )
        results = _ClientSupport.decrypt_result(
            client_parameters.cpp(), keyset.cpp(), public_result.cpp(), circuit_name
        )

        def process_result(result):
            lambda_arg = LambdaArgument.wrap(result)
            is_signed = lambda_arg.is_signed()
            if lambda_arg.is_scalar():
                return (
                    lambda_arg.get_signed_scalar()
                    if is_signed
                    else lambda_arg.get_scalar()
                )

            if lambda_arg.is_tensor():
                return np.array(
                    (
                        lambda_arg.get_signed_tensor_data()
                        if is_signed
                        else lambda_arg.get_tensor_data()
                    ),
                    dtype=(np.int64 if is_signed else np.uint64),
                ).reshape(lambda_arg.get_tensor_shape())

            raise RuntimeError("unknown return type")

        processed_results = tuple(map(process_result, results))
        if len(processed_results) == 1:
            return processed_results[0]
        return processed_results

    @staticmethod
    def _create_lambda_argument(
        value: Union[int, np.ndarray], signed: bool
    ) -> LambdaArgument:
        """Create a lambda argument holding either an int or tensor value.

        Args:
            value (Union[int, numpy.array]): value of the argument, either an int, or a numpy array
            signed (bool): whether the value is signed

        Raises:
            TypeError: if the values aren't in the expected range, or using a wrong type

        Returns:
            LambdaArgument: lambda argument holding the appropriate value
        """

        # pylint: disable=too-many-return-statements,too-many-branches
        if not isinstance(value, ACCEPTED_TYPES):
            raise TypeError(
                "value of lambda argument must be either int, numpy.array or numpy.(u)int{8,16,32,64}"
            )
        if isinstance(value, ACCEPTED_INTS):
            if (
                isinstance(value, int)
                and not np.iinfo(np.int64).min <= value < np.iinfo(np.uint64).max
            ):
                raise TypeError(
                    "single integer must be in the range [-2**63, 2**64 - 1]"
                )
            if signed:
                return LambdaArgument.from_signed_scalar(value)
            return LambdaArgument.from_scalar(value)
        assert isinstance(value, np.ndarray)
        if value.dtype not in ACCEPTED_NUMPY_UINTS:
            raise TypeError("numpy.array must be of dtype (u)int{8,16,32,64}")
        if value.shape == ():
            if isinstance(value, np.ndarray):
                # extract the single element
                value = value.max()
            # should be a single uint here
            if signed:
                return LambdaArgument.from_signed_scalar(value)
            return LambdaArgument.from_scalar(value)
        if value.dtype == np.uint8:
            return LambdaArgument.from_tensor_u8(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.uint16:
            return LambdaArgument.from_tensor_u16(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.uint32:
            return LambdaArgument.from_tensor_u32(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.uint64:
            return LambdaArgument.from_tensor_u64(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.int8:
            return LambdaArgument.from_tensor_i8(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.int16:
            return LambdaArgument.from_tensor_i16(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.int32:
            return LambdaArgument.from_tensor_i32(
                value.flatten().tolist(), list(value.shape)
            )
        if value.dtype == np.int64:
            return LambdaArgument.from_tensor_i64(
                value.flatten().tolist(), list(value.shape)
            )
        raise TypeError("numpy.array must be of dtype (u)int{8,16,32,64}")
