import asyncio
import logging
import time
from functools import partial

import safe_radix32
from clarabot_crypto.context.base import UnknownCryptoConfigVersion
from clarabot_crypto.context.nano import NANO_CRYPTO_CONTEXT, NanoConfigLoadArgs
from clarabot_crypto.context.nano_own import NANO_OWN_CRYPTO_CONTEXT, NanoOwnConfigLoadArgs

from clarabot_nano.common.crypto.keyring import NanoPeerAccountKeyring
from clarabot_nano.common.error import AccessLoggableError
from clarabot_nano.common.handler.registry import NanoHandlerClassRegistry
from clarabot_nano.common.protocol.errors import (
    NANO_REQUEST_ERROR_PREPARSE,
    NANO_REQUEST_ERROR_PAYLOAD_PARSE,
    NANO_REQUEST_ERROR_DENIED,
    NANO_REQUEST_ERROR_INTERNAL_FAILURE,
    NANO_REQUEST_ERROR_TIMEOUT,
    NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE,
    NANO_REQUEST_ERROR,
    NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID,
    NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID,
    NANO_REQUEST_ERROR_MUTATION_ID,
    NANO_REQUEST_ERROR_PEER_NOT_TRUSTED,
    NANO_REQUEST_ERROR_TOO_MANY_REQUESTS,
)
from clarabot_nano.common.protocol.request import NanoBaseRequest, NanoRequestContext
from clarabot_nano.common.protocol.server import NanoServerRequest, NanoServerResponse
from clarabot_nano.common.request.constants import get_server_prefix_from_request_type, IDEMPOTENT_HANDLER_KEYS
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano.common.zmq.requester import NoWorkerError, AioWorkerPoolRequester
from clarabot_nano.gateway.errors import SERVER_ERROR_TOO_MANY_REQUESTS
from clarabot_nano.gateway.state.config import RoomGroupConfigNotReady
from clarabot_nano.gateway.state.gateway import GatewayState
from clarabot_nano.gateway.utility import GatewayServerError
from clarabot_nano.utility.serialize import msgpack_dumps, msgpack_loads

logger = logging.getLogger(__name__)
access_logger = logging.getLogger("nano/access-log")


class RequestProcessingError(AccessLoggableError):
    pass


async def state_request_job(
    gws: GatewayState,
    worker_pool: AioWorkerPoolRequester,
    gateway_handler_registry: NanoHandlerClassRegistry,
    op_id: int,
    server_request: bytes,
):
    """
    All requests from the clients are relayed by the server arrive here and get their own execution task.
    Preliminary authentication will check if the requesting party is allowed to make a request or not.
    """
    # TODO: timeout on the whole processing?
    account_id = ip_address = None  # to help logging unexpected errors
    process_start_t = time.perf_counter()
    try:
        try:
            nano_request = msgpack_loads(server_request)
        except Exception:
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Request is malformed")
        # primitive pre-validations
        if not isinstance(payload := nano_request.get(NanoServerRequest.PAYLOAD), bytes):
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Request payload must be bytearray")
        if not isinstance(cache_key := nano_request.get(NanoServerRequest.CACHE_KEY), bytes) and cache_key is not None:
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Request cache_key must be bytearray or null")
        if (
            not isinstance(transfer_key := nano_request.get(NanoServerRequest.TRANSFER_KEY), bytes)
            and transfer_key is not None
        ):
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Request transfer_key must be bytearray or null")
        if not isinstance(part := nano_request.get(NanoServerRequest.PART), int) and part is not None:
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Request part must be int or null")
        if cache_key and transfer_key:
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PREPARSE,
                "Request cache_key and transfer_key must not be present at the same time",
            )
        if part is None:
            pass  # NOTE: The cache or transfer key may be present with missing part when resuming a transfer.
        elif part <= 0:
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Explicit request for part 0 is invalid")
        # The server related request is deemed valid and we move on to validate the E2E encrypted user request.
        ip_address = nano_request.get(NanoServerRequest.REMOTE_IP)
        # The keys to decrypt the encrypted request (payload) are looked up by metadata (account_id and room_id).
        account_id = nano_request.get(NanoServerRequest.ACCOUNT_ID)
        meta_room_id = nano_request.get(NanoServerRequest.ROOM_ID)
        if isinstance(meta_room_id, int):
            # NOTE: Exiting here for not served rooms spares us the work of decrypting the payload. There is no useful
            # information an adversary may gain from polling the served rooms by this. The server already knows which
            # rooms this Nano serves. Users other than the owner have no ability to send requests to the Nano directly.
            # Since they will request message routing from the server by the room-id there is no way for them to
            # influence this field, as it is validated and filled in by the server. If they provide a manipulated
            # payload with correctly encrypted but maliciously different embedded room-id the inconsistency will stop
            # the processing and will not reveal served-room ids either.
            if not gws.room_config_store.get(meta_room_id):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_DENIED,
                    "Room is not served",
                    account_id=account_id,
                    maybe_account=True,
                    ip_address=ip_address,
                )
            room_kr = await gws.get_own_room_keyring(meta_room_id)
        elif meta_room_id is None:
            room_kr = None
        else:
            # NOTE: ip-address is not logged, this is still server-related field mostly
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Requesting room-id is invalid")
        if not isinstance(account_id, int):
            # NOTE: ip-address is not logged, this is still server-related field mostly
            raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Requesting account-id is missing")
        elif account_id == gws.account_id:  # requesting user is owner supposedly
            peer_kr = None
            payload_decrypter = partial(
                NANO_OWN_CRYPTO_CONTEXT.load, NanoOwnConfigLoadArgs(payload, gws.keyring_cache.self_kr, time.time())
            )
            # NOTE: Room data is not used, its authenticity is checked later from the room-id of the payload solely.
        else:
            if not room_kr:
                # NOTE: ip-address is not logged, this is still server-related field mostly
                raise RequestProcessingError(NANO_REQUEST_ERROR_PREPARSE, "Room-id is required for requesting identity")
            if account_id == 0:  # requesting user is anonymous
                if gws.extracted_config.deny_anonymous:
                    # Enforcing this policy override here is an optimization, by exiting before payload decryption.
                    # It is checked later against the requesting account's access level, instead of requester identity.
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_DENIED, "Anonymous request denied", ip_address=ip_address
                    )
                eph_pk = nano_request.get(NanoServerRequest.EPHEMERAL_PUBLIC_KEY)
                if not isinstance(eph_pk, bytes):
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_PREPARSE,
                        "Request of anonymous is not valid without an ephemeral pk",
                        ip_address=ip_address,
                    )
                peer_kr = NanoPeerAccountKeyring.load_anonymous(eph_pk)
                # NOTE: Anonymous has no identity to check for trust.
            else:  # requesting user is registered
                peer_kr = await gws.get_peer_account_keyring(account_id)
                # Enforce TOFU or explicit trust fingerprint identity validation of the peer.
                if not peer_kr.trusted:
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_PEER_NOT_TRUSTED,
                        "Requesting account's identity has changed and verification is required",
                        account_id=account_id,
                        maybe_account=True,
                        ip_address=ip_address,
                    )
            payload_decrypter = partial(
                NANO_CRYPTO_CONTEXT.load,
                NanoConfigLoadArgs(payload, gws.keyring_cache.self_kr, peer_kr, room_kr, time.time()),
            )
        # decrypt the payload
        try:
            decrypted_payload = payload_decrypter()
        except UnknownCryptoConfigVersion:
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_PARSE,  # TODO: specific error?
                "Request payload crypto version is not supported",
                account_id=account_id,
                maybe_account=True,
                ip_address=ip_address,
            )
        except ValueError:
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_PARSE,
                "Request payload was not cryptographically authentic",
                account_id=account_id,
                maybe_account=True,
                ip_address=ip_address,
            )
        del payload_decrypter  # free
        # The account_id provided by the server is guaranteed to be authentic from this point onwards, because the
        # associated keys successfully decrypted the payload from an authenticated ciphertext. The meta-room-id may
        # be authentic as well, but it must be checked to be consistent with the one specified in the payload.
        try:
            loaded_payload = msgpack_loads(decrypted_payload)  # deserialize decrypted payload
        except Exception:
            loaded_payload = None
        if not isinstance(loaded_payload, dict):
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_PARSE,
                "Bad payload format",
                account_id=account_id,
                ip_address=ip_address,
            )
        # preload type-key from payload; this is used to lookup the server_prefix (pool) to forward the request to
        handler_key = loaded_payload.get(NanoBaseRequest.TYPE)
        if not isinstance(handler_key, str):
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID,
                "Type is invalid",
                account_id=account_id,
                ip_address=ip_address,
            )
        server_prefix = get_server_prefix_from_request_type(handler_key)
        # parse room_id from payload if any
        room_id = loaded_payload.get(NanoBaseRequest.ROOM_ID)
        if isinstance(room_id, int):
            pass
        elif isinstance(room_id, str):
            try:
                room_id = safe_radix32.decode(room_id)
            except (UnicodeEncodeError, OverflowError):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID,
                    "Bad room-id format",
                    account_id=account_id,
                    ip_address=ip_address,
                )
        elif room_id is not None:
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID,
                "Room-id is invalid",
                account_id=account_id,
                ip_address=ip_address,
            )
        if room_id != meta_room_id:
            # NOTE: This prevents malicious or nosy users to poll the Nano for served room-ids, because the server
            # will inform us which room they are sending the request from.
            # NOTE: This prevents the owner from issuing room-related requests directly to the Nano. This can
            # be allowed, but for now we see no benefit of doing so.
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID,
                "Room-id is inconsistent",
                account_id=account_id,
                ip_address=ip_address,
            )
        # The meta_room_id provided by the server is guaranteed to be authentic from this point onwards, because the
        # same room_id was specified by the requesting party inside a successfully decrypted payload from an
        # authenticated ciphertext.
        # validate request by room_id and account_id
        if not room_id:
            if account_id != gws.account_id:  # pragma: no cover
                # NOTE: This cannot happen because aside from the owner, anyone else is required to encrypt the
                # payload using the room-secret. This means the room id must be specified in metadata by the server
                # for the decryption to succeed. The embedded room id must also match the metadata room id or the
                # processing stops before from inconsistency error.
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_DENIED,
                    "Room-designation is missing from request",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            # NOTE: room_id == to meta_room_id so room_kr is None here
            room = None  # some nano-admin command by the owner
        elif (room_apis := gws.extracted_config.rooms.get(room_id)) is None or (
            room_config_state := gws.room_config_store.get(room_id)
        ) is None:
            raise RequestProcessingError(
                NANO_REQUEST_ERROR_DENIED,
                "Room is not served or has no API set up",
                account_id=account_id,
                ip_address=ip_address,
            )
        else:
            try:
                access_flags = await room_config_state.get_flags(account_id)
            except RoomGroupConfigNotReady:
                if room_config_state.state_is_final():
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_DENIED,
                        "Room config is revoked or failed",
                        account_id=account_id,
                        ip_address=ip_address,
                    )
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE,
                    "Room config is not ready",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            # NOTE: Access flags are resolved for account_id: room allows access if access_flags is not None.
            # NOTE: The "deny_anonymous" policy must be checked here (as well) to enforce it against the requesting
            # account's level of access. Registered users may submit requests using their registered identity while
            # only having the anonymous access level. In this case they have the room's pinned key from an
            # anonymous-share URL presumably.
            if not access_flags or (access_flags[0] == 0 and gws.extracted_config.deny_anonymous):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_DENIED, "Access denied", account_id=account_id, ip_address=ip_address
                )
            room = {
                NanoRequestContext.ROOM_ID: room_id,  # parsed id for the server
                NanoRequestContext.ROOM_KR: room_kr if server_prefix == ShortNames.GATEWAY else room_kr.dump_internal(),
                NanoRequestContext.ROOM_APIS: room_apis,  # served API flags
                NanoRequestContext.ROOM_CONFIG_MEMBER: access_flags[0],
                NanoRequestContext.ROOM_CONFIG_GROUP: access_flags[1],
            }
        # Check if the request needs and has proper mutation-session values.
        # This is done last in the preparation phase to avoid de-syncing the sequence counter at errors.
        if handler_key not in IDEMPOTENT_HANDLER_KEYS:
            if account_id == 0:
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID,
                    f"Mutation type {handler_key} is forbidden",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            mutation_key = loaded_payload.get(NanoBaseRequest.MUTATION_KEY)
            if not isinstance(mutation_key, str):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_MUTATION_ID,
                    "Mutation-key is invalid",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            mutation_sequence = loaded_payload.get(NanoBaseRequest.MUTATION_SEQUENCE)
            if not isinstance(mutation_sequence, int):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_MUTATION_ID,
                    "Mutation-sequence is invalid",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            if not gws.mutate_tracker.check(account_id, mutation_key, mutation_sequence):
                raise RequestProcessingError(
                    NANO_REQUEST_ERROR_MUTATION_ID,
                    "Mutation-id is invalid",
                    account_id=account_id,
                    ip_address=ip_address,
                )
    except RequestProcessingError as e:
        access_logger.error(e.log)
        response = {NanoServerResponse.ERROR: e.code}
    except Exception as e:
        if isinstance(e, GatewayServerError) and e.error == SERVER_ERROR_TOO_MANY_REQUESTS:
            error_code = NANO_REQUEST_ERROR_TOO_MANY_REQUESTS
            error_msg = "Server throttling while preprocessing request"
        else:
            error_code = NANO_REQUEST_ERROR
            error_msg = "Unexpected error preprocessing request"
        log_msg = RequestProcessingError.format_message(
            error_code, error_msg, account_id=account_id, ip_address=ip_address
        )
        if error_code == NANO_REQUEST_ERROR_TOO_MANY_REQUESTS:
            access_logger.warning(log_msg)
        else:
            access_logger.error(log_msg)
            logger.exception(log_msg)
        # TODO: count unexpected errors and if it hits a rate quit (core will start a new gateway)
        response = {NanoServerResponse.ERROR: error_code}
    else:
        # Request processing preparation was successful. Delegate the execution to a server-worker.
        request_context = {
            # server config
            NanoRequestContext.PART_SIZE: gws.part_size,
            NanoRequestContext.PART_APPROX: gws.part_approx,
            NanoRequestContext.S_HMAC_PARTS: gws.s_hmac_parts,
            # request args
            NanoRequestContext.OWNER_ACCOUNT_ID: gws.account_id,
            NanoRequestContext.ACCOUNT_ID: account_id,
            NanoRequestContext.REMOTE_IP: ip_address,
            NanoRequestContext.CACHE_KEY: cache_key,
            NanoRequestContext.TRANSFER_KEY: transfer_key,
            NanoRequestContext.PART: part,
            NanoRequestContext.ROOM: room,
            # Initial request processing will create a fingerprint that binds the cache or transfer key to the
            # request for authenticity. To avoid processing large requests, but keep the validity guarantee only
            # the head of the original request will be used including the nonce and the auth-tag.
            NanoRequestContext.REQUEST_HEAD: payload[: NanoRequestContext.REQUEST_HEAD_LENGTH],
        }
        try:
            if server_prefix == ShortNames.GATEWAY:
                request_context[NanoRequestContext.OWNER_KR] = gws.keyring_cache.self_kr
                request_context[NanoRequestContext.PEER_KR] = peer_kr
                # There are a few special handlers that run in the gateway for optimization.
                try:
                    response = await gateway_handler_registry.invoke(
                        gws, gws.extracted_config, request_context, loaded_payload
                    )
                except GatewayServerError as e:
                    if e.error == SERVER_ERROR_TOO_MANY_REQUESTS:
                        raise RequestProcessingError(NANO_REQUEST_ERROR_TOO_MANY_REQUESTS) from e
                    raise e
            else:
                request_context[NanoRequestContext.OWNER_KR] = gws.keyring_cache.self_kr.dump_internal()
                request_context[NanoRequestContext.PEER_KR] = peer_kr.dump_internal() if peer_kr else None
                response_frames = await worker_pool.request(
                    [msgpack_dumps(request_context), decrypted_payload], server_prefix=server_prefix
                )
                if not (1 <= (frame_number := len(response_frames)) <= 2):
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_INTERNAL_FAILURE,
                        "Server response framing is invalid",
                        account_id=account_id,
                        ip_address=ip_address,
                    )
                try:
                    response = msgpack_loads(response_frames[0])
                except Exception:
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_INTERNAL_FAILURE,
                        "Server response format is invalid",
                        account_id=account_id,
                        ip_address=ip_address,
                    )
                if not isinstance(response, dict):
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_INTERNAL_FAILURE,
                        "Server response type is invalid",
                        account_id=account_id,
                        ip_address=ip_address,
                    )
                if frame_number == 2:
                    response[NanoServerResponse.PAYLOAD] = response_frames[1]
                elif not response:
                    # Preprocessing has failed in the server.
                    raise RequestProcessingError(
                        NANO_REQUEST_ERROR_INTERNAL_FAILURE,
                        "Server could not parse request",
                        account_id=account_id,
                        ip_address=ip_address,
                    )
            # After processing from gateway or server handler.
            response[NanoServerResponse.LATENCY] = time.perf_counter() - process_start_t
        except asyncio.TimeoutError:
            access_logger.error(
                RequestProcessingError.format_message(
                    NANO_REQUEST_ERROR_TIMEOUT,
                    "Server timed out for request (may be dupe)",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            )
            response = {NanoServerResponse.ERROR: NANO_REQUEST_ERROR_TIMEOUT}
        except NoWorkerError:
            access_logger.error(
                RequestProcessingError.format_message(
                    NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE,
                    "No appropriate server available",
                    account_id=account_id,
                    ip_address=ip_address,
                )
            )
            response = {NanoServerResponse.ERROR: NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE}
        except RequestProcessingError as e:  # pre-processing in server or post-processing in gateway errors
            # The server should have already logged the request in access, but since its response is broken we
            # will log again with a note about the situation.
            access_logger.error(f"{e.log} (may be dupe)")
            logger.exception(e.log)
            response = {NanoServerResponse.ERROR: e.code}
    # Send response.
    await gws.send_response(op_id, response)
