import asyncio
import logging
import time
from contextlib import asynccontextmanager
from typing import Sequence, Awaitable, List

import aiohttp
import safe_radix32
from clarabot_crypto.auth.common import CLIENT_SALT_INPUT_LENGTH
from clarabot_crypto.auth.password import account_password_change, create_account_password_change_request
from clarabot_crypto.keyring.private import DEFAULT_PRIVATE_KEYRING_CLASS
from clarabot_nano_common.common.config.constants import CommonConfigConstants
from clarabot_nano_common.common.crypto.keyring import NanoSelfAccountKeyring
from clarabot_nano_common.common.protocol.errors import NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE
from clarabot_nano_common.common.protocol.server import NanoServerResponse
from clarabot_nano_common.gateway.handler.create_mutation_id import CreateMutationIdHandler
from clarabot_nano_common.gateway.process import AbstractGatewayProcess
from clarabot_nano_common.gateway.tofudb import FirstUseIdentityDatabase
from clarabot_nano_common.gateway.utility import check_single_response
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads

from clarabot_nano import __version__
from clarabot_nano.common.auth.handler import ServerAuthenticator, LoginFailed
from clarabot_nano.common.auth.two_factor import gateway_to_control, TwoFactorError
from clarabot_nano.common.config.client import AbstractAioConfiguredWorkerProcess
from clarabot_nano.common.core.message import (
    MSG_TYPE_REQUEST_STATE_UPDATE,
    MSG_TYPE_GATEWAY_LOGIN,
    MSG_TYPE_GATEWAY_LOGOUT,
    MSG_TYPE_GATEWAY_CHANGE_PASSWORD,
    MSG_TYPE_CONTROL_REQ_STATS,
)
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano.common.upgrade import USE_PASSWORD_NORMALIZATION
from clarabot_nano.common.zmq.auth import setup_server_auth
from clarabot_nano.common.zmq.constants import INTERNAL_REQUEST_MESSAGE, REPLY_MESSAGE
from clarabot_nano.common.zmq.handler import AioZmqMessageHandler
from clarabot_nano.common.zmq.inproc_client import AioInprocClient
from clarabot_nano.common.zmq.requester import AioWorkerPoolRequester, AioClientRequester
from clarabot_nano.common.zmq.router_server import AioZmqRouterServerThread
from clarabot_nano.common.zmq_utils import zmq_authenticator
from clarabot_nano.gateway.adhocdb import AdhocKeyDatabase
from clarabot_nano.gateway.config import GatewayConfig
from clarabot_nano.gateway.handler.admin.create_session import AdminCreateSessionHandler
from clarabot_nano.gateway.handler.admin.log.read import AdminLogReadHandler
from clarabot_nano.gateway.handler.admin.log.stream import AdminLogStreamHandler
from clarabot_nano.gateway.handler.count_adhoc_tokens import CountAdhocTokensHandler
from clarabot_nano.gateway.handler.create_adhoc_token import CreateAdhocTokenHandler
from clarabot_nano.gateway.handler.create_room_block import CreateRoomConfigBlockHandler
from clarabot_nano.gateway.handler.inter_nano_job.add import AddInterNanoJobHandler
from clarabot_nano.gateway.handler.inter_nano_job.delete import DeleteSpecificInterNanoJobHandler
from clarabot_nano.gateway.handler.inter_nano_job.describe import DescribeSpecificInterNanoJobHandler
from clarabot_nano.gateway.handler.inter_nano_job.list import ListInterNanoJobsHandler
from clarabot_nano.gateway.handler.reset_adhoc_tokens import ResetAdhocTokensHandler
from clarabot_nano.gateway.job.backup import state_backup_job
from clarabot_nano.gateway.job.config import state_config_job
from clarabot_nano.gateway.job.inter import GatewayInterNanoJobQueue
from clarabot_nano.gateway.job.internal import state_internal_request_job
from clarabot_nano.gateway.job.keyring import state_keyring_job
from clarabot_nano.gateway.job.miscellaneous import state_miscellaneous_job
from clarabot_nano.gateway.job.request import state_request_job
from clarabot_nano.gateway.job.session import state_session_job
from clarabot_nano.gateway.state.gateway import GatewayState
from clarabot_nano.inter.tmpfiledb import InterTemporaryFileDatabase

logger = logging.getLogger(__name__)


class GatewayProcess(AbstractGatewayProcess, AbstractAioConfiguredWorkerProcess):
    _extracted_config: GatewayConfig = None
    _authenticator: ServerAuthenticator
    _active_state: GatewayState

    def __init__(self):
        AbstractGatewayProcess.__init__(
            self,
            __version__,
            {
                "new-session": self._event_new_session,
                "lost-session": self._event_lost_session,
                "config-change": self._event_config_change,
                "trust-change": self._event_trust_change,
                "privy-event": self._event_privy,
            },
            [
                AdminCreateSessionHandler,
                CreateAdhocTokenHandler,
                CreateMutationIdHandler,
                CreateRoomConfigBlockHandler,
                ResetAdhocTokensHandler,
                CountAdhocTokensHandler,
                AdminLogReadHandler,
                AdminLogStreamHandler,
                AddInterNanoJobHandler,
                DeleteSpecificInterNanoJobHandler,
                DescribeSpecificInterNanoJobHandler,
                ListInterNanoJobsHandler,
            ],
        )
        # super class init will start core client and load initial config
        AbstractAioConfiguredWorkerProcess.__init__(self, ShortNames.GATEWAY, extracted_config_cls=GatewayConfig)
        # zmq facilities
        self._worker_server: AioZmqRouterServerThread = None
        self._worker_server_client: AioInprocClient = None
        self._worker_server_handler = AioZmqMessageHandler(type_index=1)
        self._worker_requester: AioWorkerPoolRequester = None
        self._core_requester: AioClientRequester = None
        self._tofu_db: FirstUseIdentityDatabase = None
        self._adhoc_db: AdhocKeyDatabase = None
        self._temp_db: InterTemporaryFileDatabase = None
        self._core_handler.add_handler(MSG_TYPE_GATEWAY_LOGIN, self._handle_core_login)
        self._core_handler.add_handler(MSG_TYPE_GATEWAY_LOGOUT, self._handle_core_logout)
        self._core_handler.add_handler(MSG_TYPE_GATEWAY_CHANGE_PASSWORD, self._handle_core_change_password)
        # stats
        self._request_stats_count = 0

    async def _main_wrapper(self):
        """
        Extend the configured client process with a server component for use with workers.
        """
        self._core_requester = self.create_core_requester()
        with self._core_requester:
            with zmq_authenticator() as auth:
                # setup server facility then proceed with main asyncio bootstrap
                self._worker_server = AioZmqRouterServerThread(
                    self._loop, "server", self.get_config()["local"]["localhost_network_security"]
                )
                setup_server_auth(auth, self._worker_server)
                self._worker_server.start(name_prefix=self.__class__.__name__)  # bind port
                self._worker_server_client = self._worker_server.create_inproc_client()
                # create router-dealer requester with round-robin identity pool and keepalive management
                self._worker_requester = AioWorkerPoolRequester(self._worker_server_client, self._worker_server_handler)
                self._worker_server_handler.add_handler(INTERNAL_REQUEST_MESSAGE, self._worker_request_internal)
                try:
                    await AbstractAioConfiguredWorkerProcess._main_wrapper(self)
                finally:
                    self._worker_server.stop()
                    self._worker_server_client.close()

    @asynccontextmanager
    async def _create_main_tasks(self) -> Sequence[Awaitable]:
        """
        Our additional main task is for processing the messages of the workers (when we are the server).
        The TOFU identity database is also started here, because this is the point that is before the main-tasks
        nursery and we have the extracted config for use as well.
        """
        async with FirstUseIdentityDatabase.create(
            str(self._extracted_config.data_dir / "tofu.db"), sqlite_timeout=15.0
        ) as tofu_db:
            self._tofu_db = tofu_db
            async with AdhocKeyDatabase.create(
                str(self._extracted_config.data_dir / "adhoc.db"), sqlite_timeout=15.0
            ) as adhoc_db:
                self._adhoc_db = adhoc_db
                yield [
                    self._process_server_messages(),
                    self._request_stats_reporter(),
                    tofu_db.closed.wait(),
                    adhoc_db.closed.wait(),
                ]

    async def _on_new_extracted_config(self, old_extracted_config):
        if old_extracted_config:
            if self._active_state:
                self._active_state.handle_new_extracted_config(self._extracted_config, old_extracted_config)
            # NOTE: if some config requires complete reconnect then put event by .queue_main_event() instead
            if (
                self._extracted_config.vendor != old_extracted_config.vendor
                or self._extracted_config.license != old_extracted_config.license
                or self._extracted_config.device_id != old_extracted_config.device_id
            ):
                self.queue_main_event({"type": "reconnect"})

    async def _process_server_messages(self):
        """
        Handle messages from workers. These are mostly keepalive and reply messages.
        """
        while True:
            messages = await self._worker_server_client.recv_multipart_messages()
            for message in messages:
                await self._worker_server_handler.handle(message)

    def _worker_request_internal(self, message: Sequence[bytes]):
        if self._active_state:
            self._active_state.ensure_bound_task(
                asyncio.create_task(state_internal_request_job(self._active_state, self._worker_server_client, message))
            )
        else:
            # TODO: maybe we want to support internal requests without active gateway?
            #   currently this may be only relevant to the document server: if an editing is in-progress at logout
            #   the changed document may not be able to get saved, because the context (from here) won't be available
            self._worker_server_client.send_multipart(
                [
                    message[0],
                    REPLY_MESSAGE,
                    message[2],
                    msgpack_dumps({NanoServerResponse.ERROR: NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE}),
                ]
            )

    async def _request_stats_reporter(self):
        """
        Submit request statistics periodically for summary to the control process.
        """
        since, self._request_stats_count = time.perf_counter(), 0
        while True:
            await asyncio.sleep(2.5)
            t, c = since, self._request_stats_count
            since, self._request_stats_count = time.perf_counter(), 0
            c /= since - t  # normalize stat to request/second average
            self._core_requester.notify([MSG_TYPE_CONTROL_REQ_STATS, msgpack_dumps({"reqs": c})])

    async def _main(self):
        await self._set_status("starting")
        # initialise state by publishing server access to workers through the core
        await self._core_requester.request(
            [
                MSG_TYPE_REQUEST_STATE_UPDATE,
                msgpack_dumps(
                    {
                        "_handler_key": "set_gateway_access",
                        "port": self._worker_server.port,
                        "key": self._worker_server.get_client_auth_info(),
                    }
                ),
            ]
        )
        async with InterTemporaryFileDatabase.create(
            str(self._extracted_config.data_dir / "gwtmp.db"), sqlite_timeout=10.0
        ) as temp_file_db:
            self._temp_db = temp_file_db
            # main websocket connection management loop
            # The two things that are managed here have an effect on each other. When an event (or command request) is
            # to be processed the work coroutine must be stopped and restarted after to ensure consistency. This is not
            # a problem since such events actually relate to the connection state and should not be frequent.
            async with self._work_setup() as session:
                await AbstractAioConfiguredWorkerProcess._main(self, session)

    async def _create_authenticator(self, session: aiohttp.ClientSession) -> ServerAuthenticator:
        return ServerAuthenticator(
            session,
            CommonConfigConstants.SERVER_DOMAIN,
            CommonConfigConstants.SERVER_PATH,
            self._extracted_config.data_dir,
        )

    async def _process_main_event(self, event):
        """
        Work coroutine is not running, execute command requests or events made from config changes.
        """
        t = event.get("type")
        if t == "login":
            try:
                await self._authenticator.login(
                    event["email"],
                    event["password"],
                    server_session_key=event.get("server_session_key"),
                    otp_salt=event.get("otpSalt"),
                    tfa_code=event.get("2faCode"),
                )
            except TwoFactorError as e:
                await self._set_command_result(e.message)
            except Exception as e:
                if not isinstance(e, LoginFailed):
                    logger.exception("Login failed")
                await self._set_command_result("login failed")
            else:
                await self._set_command_result("login successful")
                self._disconnected = 0  # reset delay for new attempt
        elif t == "logout":
            try:
                success = await self._authenticator.logout()
            except Exception:
                logger.exception("Logout failed")
                success = False
            if success:
                await self._set_command_result("logout successful")
            else:
                await self._set_command_result("logout failed")
        elif t == "reconnect":
            pass
        else:  # pragma: no cover
            raise Exception(f"unknown gateway main event {t}")

    async def _set_status(self, status: str):
        await self._core_requester.request(
            [
                MSG_TYPE_REQUEST_STATE_UPDATE,
                msgpack_dumps({"_handler_key": "set_gateway_status", "status": status, "update": self._update_state}),
            ]
        )

    async def _set_command_result(self, result: str):
        await self._core_requester.request(
            [
                MSG_TYPE_REQUEST_STATE_UPDATE,
                msgpack_dumps({"_handler_key": "set_gateway_command_result", "command_result": result}),
            ]
        )

    async def _handle_core_login(self, message: Sequence[bytes]):
        if len(message) != 2:
            raise Exception("Invalid login command")
        data = msgpack_loads(message[1])
        if not isinstance(email := data.get("email"), str) or not email:
            await self._set_command_result("invalid or missing email")
            return
        if not isinstance(password := data.get("password"), str) or not password:
            await self._set_command_result("invalid or missing password")
            return
        # try a login (or change account)
        event = {"type": "login", "email": email, "password": password}
        for k in ("server_session_key", "otpSalt", "2faCode"):
            event[k] = data.get(k)
        self.queue_main_event(event)

    async def _handle_core_logout(self, message: Sequence[bytes]):
        if len(message) != 2:  # NOTE: args are not used
            raise Exception("Invalid logout command")
        self.queue_main_event({"type": "logout"})

    async def _handle_core_change_password(self, message: Sequence[bytes]):
        if len(message) != 2:
            raise Exception("Invalid change-password command")
        data = msgpack_loads(message[1])
        if not isinstance(password := data.get("password"), str) or not password:
            await self._set_command_result("invalid or missing password")
            return
        if (
            not isinstance(client_salt_input := data.get("client_salt_input"), bytes)
            or len(client_salt_input) != CLIENT_SALT_INPUT_LENGTH
        ):
            await self._set_command_result("invalid or missing client_salt_input")
            return
        if not self._active_state:
            await self._set_command_result("must be connected to change password")
            return
        req = create_account_password_change_request(
            self._active_state.keyring_cache.self_kr,
            *account_password_change(
                DEFAULT_PRIVATE_KEYRING_CLASS,
                self._active_state.keyring_cache.self_kr,
                DEFAULT_PRIVATE_KEYRING_CLASS.normalize_password(password) if USE_PASSWORD_NORMALIZATION else password,
                client_salt_input,
                enforce_normalized_password=USE_PASSWORD_NORMALIZATION,
            ),
        )
        for k in ("otpSalt", "2faCode"):
            req[k] = data.get(k)
        result = "password change successful"
        try:
            async with self._active_state.server_request({"type": "edit-password", "variables": req}) as future:
                resp = check_single_response(await future, dict, "edit account password")
            if isinstance(r := resp["data"]["editPassword"], dict) and (two_factor := r.get("2FA")):
                result = gateway_to_control(two_factor)
            elif resp["data"]["editPassword"] is not True:
                raise Exception("Unexpected response value for edit-password")
        except Exception:
            result = "password change failed"
        await self._set_command_result(result)

    async def _create_state_bound_tasks(self, state) -> List[asyncio.Future]:
        return [
            state.ensure_bound_task(state_session_job(state, state._session_job_queue, self._set_status)),
            state.ensure_bound_task(state_config_job(state, state._config_job_queue)),
            state.ensure_bound_task(state_keyring_job(state, state._keyring_job_queue, self._tofu_db)),
            state.ensure_bound_task(state_miscellaneous_job(state, state._miscellaneous_job_queue)),
            state.ensure_bound_task(state_backup_job(state, state._backup_job_queue, self._temp_db)),
            state.ensure_bound_task(state.inter_job_queue.run()),
            state.ensure_bound_task(
                state.adhoc_db.cleanup_not_served_rooms_expired_tokens(state.extracted_config.get_serves())
            ),
            state.ensure_bound_task(self._temp_db.cleanup_old_temp_files()),
        ]

    async def _create_state(self, ws: aiohttp.ClientWebSocketResponse, self_kr: NanoSelfAccountKeyring) -> GatewayState:
        state = GatewayState(
            ws,
            self._extracted_config,
            self._adhoc_db,
            self_kr,
            self._core_requester,
            self._worker_requester,
            self._handler_registry,
            self._temp_db,
        )
        state.inter_job_queue = GatewayInterNanoJobQueue(state)
        return state

    def request(self, op_id, payload):
        """
        Relayed request from a client.
        """
        self._request_stats_count += 1
        self._active_state.bind_client_task(
            asyncio.create_task(
                state_request_job(self._active_state, self._worker_requester, self._handler_registry, op_id, payload)
            ),
            op_id,
        )

    def _event_new_session(self, payload):
        """
        The nano got a license binding and a session. The token shall be saved and used for rapid reconnect support.
        """
        if not isinstance(token := payload.get("token"), str):
            raise Exception(f"Invalid new session event: {payload}")
        self._authenticator.save_continue_token(token)
        self._active_state.handle_session_change("new")

    def _event_lost_session(self, payload):
        """
        The continue token should be kept. This does not mean that the continue token is invalidated.
        """
        self._active_state.handle_session_change("lost")

    def _event_config_change(self, payload):
        """
        Config update events will bring the Nano new config blocks. These are queued and will be applied by the
        group-config managing task.
        """
        if isinstance(event := payload.get("event"), dict):
            if (
                isinstance(room_id := event.get("roomId"), str)
                and isinstance(block := event.get("block"), bytes)
                and isinstance(signature := event.get("signature"), bytes)
            ):
                try:
                    room_id = safe_radix32.decode(room_id)
                except (UnicodeEncodeError, OverflowError):
                    pass
                else:
                    self._active_state.handle_config_change(room_id, block, signature)
                    return
        logger.warning("Invalid config-change event %s", payload)

    def _event_trust_change(self, payload):
        """
        Trust update events will bring the Nano new peer-account trust fingerprints. These are queued and will
        be applied by the keyring managing task.
        """
        if isinstance(event := payload.get("event"), dict):
            if isinstance(peer_id := event.get("peerId"), str) and isinstance(trust := event.get("trust"), bytes):
                try:
                    peer_id = safe_radix32.decode(peer_id)
                except (UnicodeEncodeError, OverflowError):
                    pass
                else:
                    self._active_state.handle_trust_change(peer_id, trust)
                    return
        logger.warning("Invalid trust-change event %s", payload)

    def _event_privy(self, payload):
        """
        Request or info sent by a Privy Nano (Backup agent).
        """
        if isinstance(event := payload.get("event"), dict):
            if isinstance(room_id := event.get("privyRoomId"), str) and isinstance(
                request := event.get("privyRequest"), str
            ):
                try:
                    safe_radix32.decode(room_id)
                except (UnicodeEncodeError, OverflowError):
                    pass
                else:
                    if request == "backup-now":
                        self._active_state.handle_backup_request(room_id)
                        return
        logger.warning("Invalid privy event %s", payload)
