import asyncio
import hmac
import logging
import os
import time
from functools import partial
from typing import Optional, Dict, Set, List, Tuple, Union

import aiohttp
import safe_radix32
from clarabot_crypto.group_config.block import GroupConfigBlockMemberFlags
from clarabot_crypto.license.file import LicenseFile
from clarabot_nano_common.common.config.constants import LocalSecrets
from clarabot_nano_common.common.crypto.keyring import (
    NanoSelfAccountKeyring,
    NanoPeerAccountKeyring,
    NanoRoomKeyring,
    NanoAdhocKeyring,
)
from clarabot_nano_common.common.handler.base import NanoHandleHelper
from clarabot_nano_common.common.handler.registry import NanoHandlerClassRegistry
from clarabot_nano_common.common.protocol.adhoc import AdhocTokenFields
from clarabot_nano_common.common.protocol.request import NanoRequestContext
from clarabot_nano_common.gateway.state.gateway import AbstractGatewayState, NanoRejected
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads, hs256_sign, hs256_verify
from ttldict2 import TTLDict

from clarabot_nano.common.config.extract.admin_password import AdminPasswordExtractedConfig
from clarabot_nano.common.core.message import MSG_TYPE_REQUEST_STATE_UPDATE
from clarabot_nano.common.gateway.error import GatewayRequestProcessingError
from clarabot_nano.common.zmq.requester import AioClientRequester, AioWorkerPoolRequester
from clarabot_nano.gateway.adhocdb import AdhocKeyDatabase
from clarabot_nano.gateway.config import GatewayConfig
from clarabot_nano.gateway.logs import LogReaderManager
from clarabot_nano.gateway.request import get_room_access_info
from clarabot_nano.gateway.state.admin import GatewayAdminSessions
from clarabot_nano.gateway.state.config import RoomGroupConfigState
from clarabot_nano.gateway.state.keyring import KeyringCache
from clarabot_nano.inter.tmpfiledb import InterTemporaryFileDatabase

logger = logging.getLogger(__name__)


def _discard_task_from_map(m, k, t):
    m.pop(k, None)
    # t.exception()  # avoid logging?


class AdhocTokenData:
    __slots__ = ("room_id", "kr", "expires", "data", "handlers")

    def __init__(self, room_id: int, kr: NanoAdhocKeyring, expires: Optional[int], data: dict):
        self.room_id = room_id
        self.kr = kr
        self.expires = expires
        self.data = data
        if self.data.get(AdhocTokenFields.DATA_TYPE) == AdhocTokenFields.TYPE_FS and isinstance(
            self.data.get("path"), str
        ):
            self.handlers = AdhocTokenFields.HANDLERS[AdhocTokenFields.TYPE_FS]
        else:
            self.handlers = None


class NanoLicenseRejected(NanoRejected):
    def description(self) -> str:
        return "Server rejected the License"


class NanoDeviceRejected(NanoRejected):
    def description(self) -> str:
        return "Server rejected the device"


class GatewayState(AbstractGatewayState, NanoHandleHelper):
    """
    This object holds pretty much all the state that the current connection is managing.
    """

    ADHOC_TOKEN_SALT_LENGTH = 16

    keyring_cache: KeyringCache

    def __init__(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        extracted_config: GatewayConfig,
        adhoc_db: AdhocKeyDatabase,
        self_kr: NanoSelfAccountKeyring,
        core_requester: AioClientRequester,
        worker_requester: AioWorkerPoolRequester,  # for invoking internal requests in gateway
        handler_registry: NanoHandlerClassRegistry,  # for invoking internal requests in gateway
        inter_tmp_db: InterTemporaryFileDatabase,
    ):
        AbstractGatewayState.__init__(self, ws, KeyringCache(self_kr))
        self.core_requester = core_requester
        self.worker_requester = worker_requester
        self.handler_registry = handler_registry
        self.inter_tmp_db = inter_tmp_db
        self.inter_job_queue = None
        # config stores and states
        self.extracted_config = extracted_config
        self.adhoc_db = adhoc_db
        self._adhoc_lock = asyncio.Lock()
        self.admin_sessions = GatewayAdminSessions(lambda: self.extracted_config)
        self.log_readers = LogReaderManager(lambda: self.part_size)
        self.room_config_store: Dict[int, RoomGroupConfigState] = dict()
        self.avatar_cache = TTLDict(300.0)
        self.adhoc_token_data_cache = TTLDict(600.0, 10000)
        self._kr_peer_load_tasks: Dict[int, asyncio.Future] = dict()
        self._kr_room_load_tasks: Dict[int, asyncio.Future] = dict()
        self._avatar_load_tasks: Dict[int, asyncio.Future] = dict()
        # job queues
        self._session_job_queue = asyncio.Queue()
        self._config_job_queue = asyncio.Queue()
        self._keyring_job_queue = asyncio.Queue()
        self._backup_job_queue = asyncio.Queue()
        self._miscellaneous_job_queue = asyncio.Queue()
        self._tmp_license = None  # save the used license for the init phase

    async def _customize_init_payload(self, init_payload: dict):
        if self.extracted_config.vendor:
            init_payload["vendor"] = self.extracted_config.vendor
        # Do not even present license files that would not be accepted by the client.
        if self.extracted_config.license and self.extracted_config.license_local_validity == "valid":
            init_payload["license"] = self._tmp_license = self.extracted_config.license
        else:
            self._tmp_license = None
        if self.extracted_config.device_id:
            init_payload["device_id"] = self.extracted_config.device_id

    async def _handle_init_close_code(self, code):
        if code == 4101:
            sent_license = self._tmp_license
            self._tmp_license = None
            await self.core_requester.request(
                [
                    MSG_TYPE_REQUEST_STATE_UPDATE,
                    msgpack_dumps(
                        {"_handler_key": "server_license_validity", "license": sent_license, "validity": "rejected"}
                    ),
                ]
            )
            raise NanoLicenseRejected()
        elif code == 4102:
            raise NanoDeviceRejected()
        await AbstractGatewayState._handle_init_close_code(self, code)

    async def _handle_customize_init_response_payload(self, payload: dict):
        sent_license = self._tmp_license
        self._tmp_license = None
        new_license = None
        if isinstance(payload.get("new_license"), str):
            try:
                if not LicenseFile.loads(payload["new_license"]).has_vendor(self.extracted_config.vendor):
                    raise ValueError("License has vendor mismatch")
            except Exception:
                logger.exception("Invalid new_license from server")
            else:
                hook_vendor_data = self.extracted_config.hook_vendor_data
                if not hook_vendor_data:
                    logger.error("Got new_license from server, but no vendor data hook is configured")
                else:
                    new_license = payload["new_license"]
        # NOTE: We allow running even if the license save/set fails because the server accepted this Nano.
        if new_license:
            reply = await self.core_requester.request(
                [
                    MSG_TYPE_REQUEST_STATE_UPDATE,
                    msgpack_dumps({"_handler_key": "new_license", "license": new_license}),
                ]
            )
            if len(reply) == 1 and bool(msgpack_loads(reply[0]).get("ok")):
                await asyncio.sleep(3.0)  # this will be interrupted by reconnect cycle
                return
        await self.core_requester.request(
            [
                MSG_TYPE_REQUEST_STATE_UPDATE,
                msgpack_dumps(
                    {
                        "_handler_key": "server_license_validity",
                        "license": sent_license,
                        "validity": "accepted" if sent_license or new_license else "unknown",
                    }
                ),
            ]
        )

    def _account_id_from_connection_ack(self, account_id: int):
        account_id_rad32 = safe_radix32.encode_fw(account_id)
        kr_account_id_rad32 = self.keyring_cache.self_kr.id.decode()
        if account_id_rad32 != kr_account_id_rad32:
            raise aiohttp.ClientError(
                f"inconsistent account-id from keyring and server: {account_id_rad32} {kr_account_id_rad32}"
            )
        self.account_id = account_id

    async def close(self):
        await AbstractGatewayState.close(self)
        await self.log_readers.close()

    def handle_new_extracted_config(self, extracted_config: GatewayConfig, old_extracted_config: GatewayConfig):
        """
        This helps to update the running state from config changes without requiring a reconnect to the server.
        """
        self.extracted_config = extracted_config
        if self.extracted_config.rooms != old_extracted_config.rooms:
            serves = self.extracted_config.get_serves()
            self._session_job_queue.put_nowait({"type": "serve-change", "serves": serves})
            self._loop.create_task(
                self.adhoc_db.drop_by_inverse_rooms([tmp["id"] for tmp in serves])
            ).add_done_callback(self._adhoc_db_serves_callback)
        # Notify the room-config manager task of the verified block changes.
        old_room_blocks = old_extracted_config.room_blocks
        room_blocks = extracted_config.room_blocks
        room_block_keys = set(old_extracted_config.room_blocks.keys())
        room_block_keys.update(extracted_config.room_blocks.keys())
        for room_id_rad32 in room_block_keys:
            if old_room_blocks.get(room_id_rad32) != room_blocks.get(room_id_rad32):
                self._config_job_queue.put_nowait(({"type": "verified", "room_id_rad32": room_id_rad32}))
        if not AdminPasswordExtractedConfig.self_matches(extracted_config, old_extracted_config):
            self.admin_sessions.reset()
        if (
            self.extracted_config.secrets[LocalSecrets.ADHOC_TOKEN_SIGN_SECRET]
            != old_extracted_config.secrets[LocalSecrets.ADHOC_TOKEN_SIGN_SECRET]
        ):
            self.adhoc_token_data_cache.clear()
        if self.extracted_config.get_drive_backups() != old_extracted_config.get_drive_backups():
            self._backup_job_queue.put_nowait({"type": "config-change"})

    def _adhoc_db_serves_callback(self, t):
        self.adhoc_token_data_cache.clear()
        try:
            t.result()
        except Exception:
            logger.exception("Error cleaning up not served adhoc-tokens")

    # keyring access api
    async def get_room_owner_keyring(
        self, room_id: int
    ) -> Optional[Tuple[Union[NanoPeerAccountKeyring, NanoSelfAccountKeyring], NanoRoomKeyring]]:
        kr = self.keyring_cache.get_own_room_kr(room_id)
        if kr:
            return self.keyring_cache.self_kr, kr
        p_rkr = self.keyring_cache.get_peer_room_kr(room_id)
        if p_rkr:
            pkr = await self.get_peer_account_keyring(p_rkr[0])
            return pkr, p_rkr[1]
        return await asyncio.shield(self._get_room_owner_keyring_task(room_id))

    async def get_owner_keyring_of_room(
        self, room_id: int
    ) -> Optional[Union[NanoPeerAccountKeyring, NanoSelfAccountKeyring]]:
        kr = self.keyring_cache.get_own_room_kr(room_id)
        if kr:
            return self.keyring_cache.self_kr
        p_rkr = self.keyring_cache.get_peer_room_kr(room_id)
        if p_rkr:
            owner_id = p_rkr[0]
        else:
            owner_id = self.keyring_cache.get_owner_of_peer_room(room_id)
        if owner_id:
            return await self.get_peer_account_keyring(owner_id)
        p_kr, _ = await asyncio.shield(self._get_room_owner_keyring_task(room_id))
        return p_kr

    def _get_room_owner_keyring_task(self, room_id: int) -> asyncio.Task:
        task = self._kr_room_load_tasks.get(room_id)
        if not task:
            self._kr_room_load_tasks[room_id] = task = asyncio.create_task(self._load_room_owner_keyring(room_id))
            task.add_done_callback(partial(_discard_task_from_map, self._kr_room_load_tasks, room_id))
        return task

    async def _load_room_owner_keyring(self, room_id: int):
        f = self._loop.create_future()
        self._keyring_job_queue.put_nowait({"type": "load-owner-room", "id": room_id, "future": f})
        return await f

    async def get_peer_account_keyring(self, account_id: int) -> NanoPeerAccountKeyring:
        kr = self.keyring_cache.get_peer_kr(account_id)
        if kr:
            return kr
        return await asyncio.shield(self._get_peer_account_keyring_task(account_id))

    def _get_peer_account_keyring_task(self, account_id: int) -> asyncio.Task:
        task = self._kr_peer_load_tasks.get(account_id)
        if not task:
            self._kr_peer_load_tasks[account_id] = task = asyncio.create_task(
                self._load_peer_account_keyring(account_id)
            )
            task.add_done_callback(partial(_discard_task_from_map, self._kr_peer_load_tasks, account_id))
        return task

    async def _load_peer_account_keyring(self, account_id: int):
        f = self._loop.create_future()
        self._keyring_job_queue.put_nowait({"type": "load-peer", "id": account_id, "future": f})
        return await f

    async def get_own_room_keyring(self, room_id: int) -> NanoRoomKeyring:
        kr = self.keyring_cache.get_own_room_kr(room_id)
        if kr:
            return kr
        return await asyncio.shield(self._get_room_keyring_task(room_id))

    def _get_room_keyring_task(self, room_id: int) -> asyncio.Task:
        task = self._kr_peer_load_tasks.get(room_id)
        if not task:
            self._kr_peer_load_tasks[room_id] = task = asyncio.create_task(self._load_room_keyring(room_id))
            task.add_done_callback(partial(_discard_task_from_map, self._kr_peer_load_tasks, room_id))
        return task

    async def _load_room_keyring(self, room_id: int):
        f = self._loop.create_future()
        self._keyring_job_queue.put_nowait({"type": "load-room", "id": room_id, "future": f})
        return await f

    # miscellaneous api
    async def get_account_avatar(self, account_id: int) -> Optional[Tuple[str, Optional[str]]]:
        if account_id == 0:
            return "Anonymous", None
        name_key = self.avatar_cache.get(account_id)
        if name_key is not None:
            return name_key
        return await asyncio.shield(self._get_account_avatar_task(account_id))

    def _get_account_avatar_task(self, account_id: int) -> asyncio.Task:
        task = self._avatar_load_tasks.get(account_id)
        if not task:
            self._avatar_load_tasks[account_id] = task = asyncio.create_task(self._load_account_avatar(account_id))
            task.add_done_callback(partial(_discard_task_from_map, self._avatar_load_tasks, account_id))
        return task

    async def _load_account_avatar(self, account_id: int):
        f = self._loop.create_future()
        self._miscellaneous_job_queue.put_nowait({"type": "load-avatar", "id": account_id, "future": f})
        return await f

    # room config utility api
    def set_served_rooms(self, served: Set[int]):
        """
        Updates the store of the group-configs of served rooms. Also produces events to load the new rooms.
        """
        old = []
        for room_id in self.room_config_store.keys():
            try:
                served.remove(room_id)
            except KeyError:
                old.append(room_id)
        for room_id in old:
            self.room_config_store.pop(room_id).close()  # cleanup waiters
        for room_id in served:
            self.room_config_store[room_id] = RoomGroupConfigState(room_id)
            self._config_job_queue.put_nowait({"type": "load-chain", "room_id": room_id})

    def get_verified_room_blocks(self, room_id: int) -> List[str]:
        verified_room_blocks = self.extracted_config.room_blocks.get(safe_radix32.encode_fw(room_id))
        if not verified_room_blocks:
            return []
        return sorted(verified_room_blocks)

    # adhoc room access utilities
    async def validate_adhoc_token(self, raw_token: str, target_room_id: int) -> Optional[AdhocTokenData]:
        """
        Get token data from cache or parse the token, then check its validity.
        """
        adhoc_token_data = self.adhoc_token_data_cache.get(raw_token)
        if not adhoc_token_data:
            token_data = hs256_verify(self.extracted_config.secrets[LocalSecrets.ADHOC_TOKEN_SIGN_SECRET], raw_token)
            if (
                not isinstance(token_data, dict)
                or not isinstance(token_id := token_data.get(AdhocTokenFields.ID), int)
                or not isinstance(token_salt := token_data.get(AdhocTokenFields.SALT), bytes)
                or len(token_salt) != self.ADHOC_TOKEN_SALT_LENGTH
            ):
                return
            adhoc_row = await self.adhoc_db.get(token_id)
            if not adhoc_row:
                return
            expires, room_id, secret, data, salt = adhoc_row
            if not hmac.compare_digest(salt, token_salt):
                return
            try:
                data = msgpack_loads(data)
            except Exception:
                return
            if not isinstance(data, dict):
                return
            try:
                kr = NanoAdhocKeyring.load(safe_radix32.encode_fw(room_id), {NanoAdhocKeyring.KEY1: secret})
            except Exception:
                return
            adhoc_token_data = AdhocTokenData(room_id, kr=kr, expires=expires, data=data)
            if not adhoc_token_data.handlers:
                return  # data not recognized
        else:
            token_data = None
        if not (await self._check_adhoc_token_data(adhoc_token_data, target_room_id)):
            return
        if token_data:  # newly parsed
            self.adhoc_token_data_cache[raw_token] = adhoc_token_data
        return adhoc_token_data

    async def _check_adhoc_token_data(self, adhoc_token_data: AdhocTokenData, target_room_id: int) -> Optional[bool]:
        """
        This is the primary access control behaviour of an (already validated) adhoc-token, it is performed for each
        request processing call.

        Adhoc-tokens are designed to provide access with an independent encryption key to a room from its static
        shared secret. Such is beneficial to avoid revealing unnecessary room data. Even though each adhoc-token has
        an independent secret for E2E key derivation, here their validity is tied to the room's permission config by
        requiring that the token-creator account is still a member. Essentially, this means that the adhoc tokens are
        tied to the "live" room permissions, they can only give access to the resources that the token creator has
        access to. Checking membership only is sufficient as long as all adhoc handlers are read-only.
        """
        if adhoc_token_data.room_id != target_room_id:
            return
        if adhoc_token_data.expires is not None and adhoc_token_data.expires < time.time():
            return
        account_id = adhoc_token_data.data.get(AdhocTokenFields.DATA_ACCOUNT_ID)
        if not isinstance(account_id, int):
            return
        try:
            config_member_flags = (await get_room_access_info(self, account_id, adhoc_token_data.room_id, None))[
                NanoRequestContext.ROOM_CONFIG_MEMBER
            ]
        except GatewayRequestProcessingError:
            return  # can't get permissions or access denied
        if not config_member_flags or not (config_member_flags & GroupConfigBlockMemberFlags.ACTIVE):
            return  # do not honor adhoc token in case its creator is no longer a member
        return True

    async def create_fs_adhoc_token(
        self, room_id: int, account_id: int, fs_path: str, ttl: Optional[int]
    ) -> Tuple[str, bytes]:
        return await self._create_adhoc_token(AdhocTokenFields.TYPE_FS, room_id, account_id, ttl, path=fs_path)

    async def _create_adhoc_token(
        self, adhoc_type: str, room_id: int, account_id: int, ttl: Optional[int] = None, **kwargs
    ) -> Tuple[str, bytes]:
        """
        The visible data in the token is purposefully kept minimal, so no information can be extracted from it.
        """
        data = {**kwargs, AdhocTokenFields.DATA_TYPE: adhoc_type, AdhocTokenFields.DATA_ACCOUNT_ID: account_id}
        adhoc_token_data = AdhocTokenData(
            room_id,
            kr=NanoAdhocKeyring.generate(safe_radix32.encode_fw(room_id)),
            expires=int(time.time()) + ttl if ttl else None,
            data=data,
        )
        adhoc_secret = adhoc_token_data.kr.get_key(NanoAdhocKeyring.KEY1)
        # NOTE: Since the token-ids may sometimes get re-used in sqlite table, the salt's random (not guessable and
        # practically non-repeating) property provides safety against adversaries harvesting adhoc tokens to exploit.
        # This will even prevent adversaries who may eventually be able to even guess the secret key for the HMAC
        # signature from guess-forging tokens.
        token_salt = os.urandom(self.ADHOC_TOKEN_SALT_LENGTH)
        token_id = await self.adhoc_db.create_token(
            adhoc_token_data.expires, room_id, adhoc_secret, msgpack_dumps(data), token_salt
        )
        adhoc_token = hs256_sign(
            self.extracted_config.secrets[LocalSecrets.ADHOC_TOKEN_SIGN_SECRET],
            {AdhocTokenFields.ID: token_id, AdhocTokenFields.SALT: token_salt},
        )
        self.adhoc_token_data_cache[adhoc_token] = adhoc_token_data
        return adhoc_token, adhoc_secret

    async def delete_adhoc_tokens_by_room(self, room_id: int):
        await self.adhoc_db.drop_by_room(room_id)
        room_tokens = [k for k, e in self.adhoc_token_data_cache.items() if e.value.room_id == room_id]
        for room_token in room_tokens:
            del self.adhoc_token_data_cache[room_token]

    # event handlers
    def handle_session_change(self, status: str):
        self._session_job_queue.put_nowait({"type": "session-change", "status": status})

    def handle_config_change(self, room_id: int, block: bytes, signature: bytes):
        self._config_job_queue.put_nowait(
            {"type": "add-block", "room_id": room_id, "block": block, "signature": signature}
        )

    def handle_serves_change(self, room_ids: Set[int]):
        self._config_job_queue.put_nowait({"type": "serves-change", "room_ids": set(room_ids)})

    def handle_trust_change(self, account_id: int, trust: bytes):
        self._keyring_job_queue.put_nowait({"type": "trust-peer", "id": account_id, "trust": trust})

    def handle_backup_request(self, room_id: str):
        self._backup_job_queue.put_nowait({"type": "backup-room", "target_room_id": room_id})
