import hashlib
import logging
import os
import random
import time
from typing import Optional, Tuple, List, Callable, Any, Union

from clarabot_crypto.base.cipher.aes import AESCipherSIVNonce96
from clarabot_crypto.base.key.kdf import HKDFKeyGeneratorV1, HKDFKeyGeneratorV2
from clarabot_crypto.context.account_signature import ACCOUNT_SIGNATURE_CRYPTO_CONTEXT, AccountSignatureConfigDumpArgs
from clarabot_crypto.context.nano import NANO_CRYPTO_CONTEXT, NanoConfigDumpArgs
from clarabot_crypto.context.nano_data import NANO_DATA_CRYPTO_CONTEXT, NanoDataConfigDumpArgs
from clarabot_crypto.context.nano_own import NANO_OWN_CRYPTO_CONTEXT, NanoOwnConfigDumpArgs

from clarabot_nano.common.config.constants import LocalSecrets
from clarabot_nano.common.config.extract.admin_password import AdminPasswordExtractedConfig
from clarabot_nano.common.config.extract.secrets import SecretsExtractedConfig
from clarabot_nano.common.error import AccessLoggableError
from clarabot_nano.common.protocol.errors import (
    NANO_REQUEST_ERROR_PART_RANGE,
    NANO_REQUEST_ERROR_CONFLICT,
    NANO_REQUEST_ERROR_ADMIN_PASSWORD,
    NANO_REQUEST_ERROR_DENIED,
    NANO_REQUEST_ERROR_ADMIN_PASSWORD_MISSING,
)
from clarabot_nano.common.protocol.request import NanoRequestContext, NanoBaseRequest
from clarabot_nano.common.protocol.response import NanoBaseResponse
from clarabot_nano.common.protocol.server import NanoServerResponse
from clarabot_nano.utility.handlers import AbstractHandler
from clarabot_nano.utility.serialize import msgpack_dumps


class NanoHandlerError(AccessLoggableError):
    @classmethod
    def format_message(cls, error_code: int, message=None, context=None):
        return super(NanoHandlerError, cls).format_message(
            error_code,
            message,
            account_id=context.get(NanoRequestContext.ACCOUNT_ID),
            ip_address=context.get(NanoRequestContext.REMOTE_IP),
        )

    def __init__(self, error_code: int, message: str, context: dict):
        # TODO: improve to include handler name for errors like:
        #   2021-06-03 17:11:18,538 244qm9Z822222: 4212 Admin-password required for request
        #  should be something like
        #   2021-06-03 17:11:18,538 244qm9Z822222: 4212 Admin-password required for request (Admin-list)
        #   2021-06-03 17:11:18,538 244qm9Z822222: Admin-list 4212 Admin-password required for request
        AccessLoggableError.__init__(self, error_code, message, context)


class MaybeConflictNanoHandlerError(NanoHandlerError):
    """
    This can be raised from processing to indicate some state-related inconsistency. The base framework will check
    if the data source has changed and the state error can be represented as a conflict error instead.
    """


class NanoHandleHelper(object):
    """
    This object shall be customized by the executing process for its handlers.
    """


class AbstractNanoHandler(AbstractHandler):
    # Basic authentication policy flags that can be specialized for the handler.
    # NOTE: The identity of the requester that is present in the context is checked by the gateway to be authentic.
    REQUIRES_OWNER_IDENTITY = False
    REQUIRES_ADMIN_PASSWORD = False  # REQUIRES_OWNER_IDENTITY must be True for this to be effective
    REJECTS_ANONYMOUS_IDENTITY = True  # REQUIRES_OWNER_IDENTITY is stronger if both are set

    # Handlers that are technical or not important to be logged (idempotent) may set this to a lower (DEBUG) value.
    SUCCESS_LOG_LEVEL = logging.INFO

    def __init__(
        self,
        helper: NanoHandleHelper,
        extracted_config: Union[SecretsExtractedConfig, AdminPasswordExtractedConfig],
        context: dict,
        payload: dict,
    ):
        self.helper = helper
        self.extracted_config = extracted_config
        self.context = context
        self.payload = payload
        # these fields are populated during execution
        self.cacheable = None
        self.request_key = None
        self.cache_key = None
        self.transfer_key = None
        self.count = None

    async def execute(self) -> dict:
        """
        Handle request caching and ensure part consistency (as much as possible) for the response.
        """
        # Enforce owner-admin policy.
        if self.REQUIRES_OWNER_IDENTITY:
            if self.context[NanoRequestContext.OWNER_ACCOUNT_ID] == self.context[NanoRequestContext.ACCOUNT_ID]:
                if self.REQUIRES_ADMIN_PASSWORD and not self.extracted_config.check_admin_password(
                    self.payload.get(NanoBaseRequest.ADMIN_PASSWORD)
                ):
                    if self.extracted_config.can_setup_password():
                        error_code = NANO_REQUEST_ERROR_ADMIN_PASSWORD_MISSING
                    else:
                        error_code = NANO_REQUEST_ERROR_ADMIN_PASSWORD
                    raise NanoHandlerError(error_code, "Admin-password required for request", self.context)
            else:
                raise NanoHandlerError(NANO_REQUEST_ERROR_DENIED, "Owner identity required for request", self.context)
        elif self.REJECTS_ANONYMOUS_IDENTITY and self.context[NanoRequestContext.ACCOUNT_ID] == 0:
            raise NanoHandlerError(
                NANO_REQUEST_ERROR_DENIED, "Anonymous identity is not valid for request", self.context
            )
        # Get request-key before processing.
        self.cacheable, self.request_key, self.count = await self._create_request_key()
        # assert if the response is cacheable the count must be at least 2 and the request_key must be something
        assert not self.cacheable or (
            self.count > 1 and isinstance(self.request_key, list) and self.request_key
        ), f"{self.cacheable} {self.count} {self.request_key} "
        # assert if the response is not cacheable the count must be at least 1
        assert self.cacheable or (self.count > 0 and (self.request_key is not None or self.count == 1))
        self.cache_key, self.transfer_key = self._secure_request_key(self.cacheable, self.request_key)
        try:
            if (part := self.context[NanoRequestContext.PART]) is None:
                # Check data source consistency for continuation if either cache or transfer keys are available.
                if self.context[NanoRequestContext.CACHE_KEY] or self.context[NanoRequestContext.TRANSFER_KEY]:
                    if self.request_key is None:
                        raise NanoHandlerError(
                            NANO_REQUEST_ERROR_PART_RANGE,
                            "Handler does not support part requests (continue)",
                            self.context,
                        )
                    elif (
                        self.cache_key != self.context[NanoRequestContext.CACHE_KEY]
                        or self.transfer_key != self.context[NanoRequestContext.TRANSFER_KEY]
                    ):
                        raise NanoHandlerError(
                            NANO_REQUEST_ERROR_CONFLICT, "Datasource state conflict (continue)", self.context
                        )
                response = {
                    NanoServerResponse.CACHE_KEY: self.cache_key,
                    NanoServerResponse.TRANSFER_KEY: self.transfer_key,
                    NanoServerResponse.COUNT: self.count,
                    NanoServerResponse.PAYLOAD: await self._execute_initial(),
                }
                part_key = None
            elif self.request_key is None:
                raise NanoHandlerError(
                    NANO_REQUEST_ERROR_PART_RANGE, "Handler does not support part requests", self.context
                )
            else:  # part > 0
                # Assert that the current request-key matches the one submitted with the request.
                if (
                    self.cache_key != self.context[NanoRequestContext.CACHE_KEY]
                    or self.transfer_key != self.context[NanoRequestContext.TRANSFER_KEY]
                ):
                    raise NanoHandlerError(NANO_REQUEST_ERROR_CONFLICT, "Datasource state conflict", self.context)
                # Handle out-of-range part requests.
                if part <= 0 or part >= self.count:
                    raise NanoHandlerError(NANO_REQUEST_ERROR_PART_RANGE, "Part request out of range", self.context)
                # Calculate S-HMAC part beginning if any should be produced for the current response.
                # To make distributed & stateless file (source data) chunking possible we cannot build on the already
                # authenticated-encrypted payloads. Arbitrary parts may be re-requested by the clients and it is
                # unacceptable that an arbitrary request to any part would require the full processing of the source
                # data in question. As a sufficiently secure, but a little less elegant workaround to not having
                # monolithic chunking procedures, the source data will be fingerprinted with a random salt and that
                # will be provided with a signature of the owner. This will be done on deterministically random ranges
                # of the data, while fully covering it whole.
                response = {
                    NanoServerResponse.PAYLOAD: await self._execute_part(part, self._get_s_hmac_part_beginning(part))
                }
                part_key = self._create_part_encryption_key()
        except MaybeConflictNanoHandlerError:
            interrupt_cacheable, interrupt_request_key, _ = await self._create_request_key()
            if self.cacheable != interrupt_cacheable or self.request_key != interrupt_request_key:
                raise NanoHandlerError(
                    NANO_REQUEST_ERROR_CONFLICT, "Datasource state conflict during-run", self.context
                )
            raise
        # Get request-key after processing and assert that it matches the one from before the processing. This
        # is supposed to ensure the state of the data source was consistent during the processing. Unfortunately
        # this is not absolutely fool-proof, but more strict solutions would only be possible by platform and
        # data-source specific implementations.
        # TODO: This check could be optional for some handlers as optimization.
        post_cacheable, post_request_key, _ = await self._create_request_key()
        if self.cacheable != post_cacheable or self.request_key != post_request_key:
            raise NanoHandlerError(NANO_REQUEST_ERROR_CONFLICT, "Datasource state conflict post-run", self.context)
        # Serialize and encrypt the payload.
        payload = msgpack_dumps(response[NanoServerResponse.PAYLOAD])
        if part_key is None:
            # First part is P2P encrypted all the time.
            if (peer_kr := self.context[NanoRequestContext.PEER_KR]) is None:
                payload = NANO_OWN_CRYPTO_CONTEXT.dump(
                    NanoOwnConfigDumpArgs(payload, self.context[NanoRequestContext.OWNER_KR], time.time())
                )
            else:
                payload = NANO_CRYPTO_CONTEXT.dump(
                    NanoConfigDumpArgs(
                        payload,
                        self.context[NanoRequestContext.OWNER_KR],
                        peer_kr,
                        self.context[NanoRequestContext.ROOM][NanoRequestContext.ROOM_KR],
                        time.time(),
                    )
                )
        else:
            # Data parts are always encrypted with a key that is unique for the data and is shared with the
            # requester with P2P setup.
            payload = NANO_DATA_CRYPTO_CONTEXT.dump(NanoDataConfigDumpArgs(payload, part_key))
        response[NanoServerResponse.PAYLOAD] = payload
        return response

    async def _create_request_key(self) -> Tuple[bool, Optional[list], int]:
        """
        Request key construction options:
            - With the first component of the returned tuple, the handler can decide if the response shall be
              cacheable or not.
            - The second returned value is the raw request_key. It's construction must be implemented so the value
              and its serialization will be deterministic. Ideally the key should be a MAC of the source data that
              is O(1) to validate. Since that is usually not possible the implementation may take viable tradeoffs.
              For example a file may be represented by its path, modification timestamp and size. In a well behaving
              system any change to any of those would indicate that the data source has changed. The key must be as
              concise as possible. If it exceeds a certain length the tail of it will be shortened by transforming
              it to a secure hash. The values that provide the most uniqueness shall be in the front, so they do not
              get hashed, just encrypted. Including the room-id that is being used for the query could fragment the
              response-cache without much benefit, so including it is not recommended.
            - The third returned value must be the number of parts the response will be split to. It must be included
              in the request_key or derived from its components if exact consistency is required. An example where
              that is not desired is a directory listing. Consistency is is lazy there and slices (parts) are merged
              together by using an overlap.

        If the request is neither multipart nor has state property, the request_key shall be None. A short random
        byte-array will be generated to provide some uniqueness, but otherwise the handler implementation must refuse
        utilizing the continue features.

        For multipart responses the parts count will change if the server adjusts the part_size for the cache. This
        must invalidate cached entries implicitly (halving the size will make the first half of the 2nd part the new
        4th part). This is handled by the server, it is guaranteed to avoid conflicts between the different size
        entries. The old size entries will be rotated out as they will no longer be hit at all.

        Some handlers may need to do major work to determine the number of parts for the complete response. In those
        cases the first call to this method could have a significant part of the execution.
        If request is impossible to be executed then raise an error.
        """
        raise NotImplementedError()

    def _secure_request_key(
        self, cacheable: bool, request_key: Optional[list]
    ) -> Tuple[Optional[bytes], Optional[bytes]]:
        """
        The output of this method will be used as a look-up key for the payload. It must be deterministic, which
        means that the serialization and encryption of the request_key must be deterministic.

        On the other hand the output must change in response to a lot of inputs in a way to minimize (ideally
        completely neutralize) chances of collisions. Part encryption keys (PKs) must similarly change while being
        practically independent of each other. These changing will auto-invalidate caches or response streams where
        the data source has changed. The following inputs are used in creating a cache or transfer key:
            - serialized_key (from request_key): Represents the state of the data source. A change in this does not
              prohibit any previous users from accessing it, but the caching or streaming will need to be refreshed.
            - local_secret_key: An exclusive locally generated and never transmitted secret key will be used for the
              encryption of the serialized_key. This protects the information in the request_key of the subject of the
              data.

        Additionally for requests where the response is not cacheable the request_key will include the request_head in
        the salt. This makes all such requests have exclusive and unique transfer_keys and thus part-keys as well.
        Because those part keys are only ever disclosed between the Nano and the requesting party all other parts
        can be considered authentic and need no further digital signature. This makes it possible to respond with a
        volatile but still multi-part response like a directory listing.
        Using this construct also hides the fact that the same user or different users have been requesting the same
        (undisclosed) resource from the server.

        The cached responses where the part_key is shared with multiple users the response is protected by additional
        digital signatures of the Nano's owner.

        All responses make use of the request_head to calculate the request-authenticity-hash which prevents a
        malicious server from replaying responses.

        TODO: Nano should notify or force the server to clear the caches when the secret key changes for it locally.
        Maximum length of the returned value is 240.
        """
        if request_key is None:
            # Single part, not cached response. (these are asserted in execute() before calling this)
            # The transfer_key is required to be present, but it is not utilized,
            encrypted_key = os.urandom(12)
        else:
            serialized_key = msgpack_dumps(request_key)  # NOTE: determinism depends on the input
            # Limit the length of the key. Because of the independent nature of the Nano, the server cannot detect or
            # handle cache_key collisions in any way. For this reason we opt to not simply create it by hashing, but
            # having it encrypted securely rather. AES is a bijection which means it will not produce any collisions
            # for any input and nonce as long as the secret key is the same. The request key shall place the data-
            # source properties that have the most uniqueness in the front as that part will not get hashed ever.
            # If the serialized key is too long, its tail will be hashed. Realistically even simply hashing the key
            # would be secure, so at this point we accept the virtual possibility of a collision.
            if len(serialized_key) >= 212:  # equals excludes collision between hashed and not hashed keys
                serialized_key = serialized_key[:180] + hashlib.sha3_256(serialized_key[180:]).digest()  # Keccak!
            info = hashlib.sha256(serialized_key).digest()
            if not cacheable:
                info += self.context[NanoRequestContext.REQUEST_HEAD]
            # We use SIV for encrypting the cache_key or transfer_key, so coincidental nonce reuse is not an issue.
            # Even though this is never decrypted the nonce and auth-tag are kept to provide collision-free
            # values reliably.
            # The key of the encryption can't be changing or derived frequently. It must be the same for all responses
            # so the bijection property of AES will guarantee us no collisions.
            encrypted_key = AESCipherSIVNonce96.encrypt(
                serialized_key,
                self.extracted_config.secrets[LocalSecrets.SERVER_CACHE_KEY_ENCRYPTION]
                if cacheable
                else self.extracted_config.secrets[LocalSecrets.SERVER_TRANSFER_KEY_ENCRYPTION],
                # NOTE: No salt necessary because the IKM is a very good quality secure local secret.
                HKDFKeyGeneratorV1.new_key(
                    12, self.extracted_config.secrets[LocalSecrets.SERVER_REQUEST_KEY_IKM], None, info
                ),
            )
        return (encrypted_key, None) if cacheable else (None, encrypted_key)

    async def _execute_initial(self) -> dict:
        """
        The first response's primary role is to setup the authenticity of the response. To provide proof that the
        response is for the current request the cache_key or transfer_key (depending on caching) is hashed with the
        request's head. This will prevent response replays by malicious servers.

        The initial response is always encoded in the same P2P encryption the requester used. This cryptographically
        binds together the unique request and the cache_key or transfer_key for the request in a way the client can
        authenticate easily.

        For multipart responses each part has to somehow be bound to this authenticated initial response or the
        original request directly. The original request is unique each time it is submitted by a client which means
        cacheable responses cannot be bound to it directly. So for simplicity and compatibility all parts shall get
        bound to the initial response. Using a secret key - only provided to the client in the initial authenticated
        response - for all part responses will prevent the server from replaying a different response-sequence
        masquerading under a different cache or transfer key.

        Overriding this method is most meaningful when no caching will be used and the actual response can be
        embedded in the initial one.
        """
        p = {
            NanoBaseResponse.REQUEST_AUTHENTICITY_HASH: hashlib.sha256(
                (self.cache_key if self.cacheable else self.transfer_key)
                + self.context[NanoRequestContext.REQUEST_HEAD]
            ).digest()[:16],
            NanoBaseResponse.COUNT: self.count,  # Count of all response parts for client, including the initial one.
        }
        if self.count > 1:  # further parts will be bound to the initial response by this key
            p[NanoBaseResponse.PART_KEY] = self._create_part_encryption_key()
        return p

    def _create_part_encryption_key(self) -> bytes:
        """
        A key is derived from a local secret using the secured transfer or cache key. No room-specific key
        is used, which allows caching of a resource among multiple rooms for the same Nano. There would be no benefit
        to the contrary option.

        Using unique keys for cached objects prevents passive access from cache by using a compromised account.
        All accounts need to get the specific key for the specific version of a specific file even if it is
        available from cache. This also means the Nano can at any point change the KDF secrets and thus permanently
        render all cache objects unusable.

        Accounts that has received a key for a cached object can continue to retrieve it provided all parts are
        available from the cache. If even a single is missing the Nano will not be able to continue that version of
        the file because the secured request_key will mismatch.
        """
        if self.cacheable:
            return HKDFKeyGeneratorV2.new_key(
                64, self.extracted_config.secrets[LocalSecrets.SERVER_CACHE_KEY_IKM], None, self.cache_key
            )
        return HKDFKeyGeneratorV2.new_key(
            64, self.extracted_config.secrets[LocalSecrets.SERVER_TRANSFER_KEY_IKM], None, self.transfer_key
        )

    def _get_s_hmac_part_beginning(self, part: int) -> Optional[int]:
        """
        Returns the beginning of the range (start, part] for which the current part should create an S-HMAC.
        The deterministic random nature of the mac part indexes makes DoS attacks harder to execute, as the most
        costly parts are not always on the same indexes.
        """
        r = random.Random(self.cache_key if self.cacheable else self.transfer_key)
        start_part = 0
        while True:
            next_part = start_part + r.randint(*self.context[NanoRequestContext.S_HMAC_PARTS])
            if next_part == part:
                return start_part
            elif next_part < part:
                start_part = next_part
            elif part == self.count - 1:
                return start_part
            else:
                break

    async def _execute_part(self, part: int, s_hmac_start_part: Optional[int]) -> dict:
        """
        Create the response payload of part.
        If S-HMAC start is present the implementation should produce the signed fingerprint
        for the range of (s_hmac_start_part, part].
        If the payload will be encrypted p2p then MAC creation is unnecessary.
        """
        return {
            NanoBaseResponse.PART: part,  # Part for client
        }

    def _add_s_hmac(self, payload: dict, bytes_count: int, fingerprint: bytes, salt: bytes):
        """
        Create and sign the fingerprint data, then add them to the payload.

        This is necessary to make it impossible for a malicious server and member to work together and create a
        forged multi-part response for an other account.
        """
        payload[NanoBaseResponse.MAC] = msgpack_dumps(
            {
                "c": bytes_count,  # number of bytes included in the fingerprint ending with the current payload
                "f": fingerprint,  # sha512 truncated for example
                "s": salt,  # from CSPRNG
            }
        )
        payload[NanoBaseResponse.MAC_SIGNATURE] = ACCOUNT_SIGNATURE_CRYPTO_CONTEXT.dump(
            AccountSignatureConfigDumpArgs(payload[NanoBaseResponse.MAC], self.context[NanoRequestContext.OWNER_KR])
        )

    def success_log_message(self) -> str:
        return NanoHandlerError.format_message(0, self._log_message(), self.context)

    def _log_message(self) -> str:
        raise NotImplementedError()

    def is_request_by_owner(self) -> bool:
        return self.context[NanoRequestContext.OWNER_ACCOUNT_ID] == self.context[NanoRequestContext.ACCOUNT_ID]

    def _approximate_part_slicing(self, items: list, size: Callable[[Any], int], overlap: int = 0) -> List[list]:
        """
        This will work well provided the list has items of relatively small size compared to the PART_SIZE.
        Large items are also accepted as long as they are each smaller than PART_SIZE. However a positive overlap
        could easily lead to impossible slice sizes, in which case OverflowError will be thrown.
        The size argument must be a callable that precisely approximates the size of an item from the list. It will
        be used with each item exactly once. The simplest general implementation could serialize the item to know
        its exact size. For example: `lambda item: len(msgpack_dumps(item))`
        This implementation creates slices deterministically.
        The overlap shall be used by the client to merge the approximately sliced data. In case two neighbour parts
        have no overlap the slicing has changed too much and the request must be restarted. This requires some form
        of identification data to be available in the entries.
        """
        slice_size = 0
        slices = [coll := []]
        for item in items:
            item_size = size(item)
            slice_size += item_size
            if slice_size >= self.context[NanoRequestContext.PART_APPROX][0]:
                if item_size > self.context[NanoRequestContext.PART_APPROX][1]:
                    # push current part to next slice because it is huge
                    slices.append(coll := coll[len(coll) - overlap :])
                    coll.append(item)
                else:
                    coll.append(item)
                    slices.append(coll := coll[len(coll) - overlap :])
                slice_size = sum(size(i) for i in coll)
                if slice_size >= self.context[NanoRequestContext.PART_APPROX][0]:
                    raise OverflowError("Overlap ate all the space")
            else:
                coll.append(item)
        if not slices[-1] and slices[0]:
            slices.pop(-1)
        return slices
