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

import aiohttp
import safe_radix32
from aiohttp import WSMsgType
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.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.config.constants import CommonConfigConstants
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.crypto.keyring import NanoSelfAccountKeyring
from clarabot_nano.common.handler.registry import NanoHandlerClassRegistry
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano.common.upgrade import USE_PASSWORD_NORMALIZATION
from clarabot_nano.common.utility import IncreasingTimeout
from clarabot_nano.common.zmq.auth import setup_server_auth
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.config import GatewayConfig
from clarabot_nano.gateway.handler.create_mutation_id import CreateMutationIdHandler
from clarabot_nano.gateway.handler.create_room_block import CreateRoomConfigBlockHandler
from clarabot_nano.gateway.job.config import state_config_job
from clarabot_nano.gateway.job.keyring import state_keyring_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, NanoVersionRejected
from clarabot_nano.gateway.tofudb import FirstUseIdentityDatabase
from clarabot_nano.gateway.utility import check_single_response
from clarabot_nano.utility.network import create_session
from clarabot_nano.utility.serialize import msgpack_dumps, msgpack_loads

logger = logging.getLogger(__name__)


class GatewayProcess(AbstractAioConfiguredWorkerProcess):
    _extracted_config: GatewayConfig = None

    # message types between this and server
    REQUEST_CONNECTION_TERMINATE = "connection_terminate"
    REQUEST_START = "start"
    RESPONSE_DATA = "data"
    RESPONSE_ERROR = "error"
    RESPONSE_COMPLETE = "complete"
    RESPONSE_EVENT = "event"
    RELAYED_REQUEST = "request"  # request by a client
    RELAYED_ABORT = "abort"  # request by a client is no longer needed
    RELAYED_RESPONSE = "response"  # response to a client request

    VALID_RESPONSE_HANDLERS = {
        RESPONSE_DATA,
        RESPONSE_ERROR,
        RESPONSE_COMPLETE,
        RESPONSE_EVENT,
        RELAYED_REQUEST,
        RELAYED_ABORT,
    }

    INCREASING_TIMEOUT = IncreasingTimeout(4.0, 90.0, 0.5)

    def __init__(self):
        # 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._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
        # other configs and states
        self.__authenticator: ServerAuthenticator = None
        self.__active_state: GatewayState = None
        self.__disconnected = 0  # for reconnect delay adjustment
        self.__update_state = -1  # software update report/advice by server
        # message processing utilities
        self.__event_type_map = {
            "new-session": self._event_new_session,
            "lost-session": self._event_lost_session,
            "config-change": self._event_config_change,
            "trust-change": self._event_trust_change,
        }
        self._handler_registry = NanoHandlerClassRegistry(
            [
                CreateMutationIdHandler,
                CreateRoomConfigBlockHandler,
            ]
        )

    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)
                try:
                    await AbstractAioConfiguredWorkerProcess._main_wrapper(self)
                finally:
                    self._worker_server.stop()
                    self._worker_server_client.close()

    def _create_main_tasks(self) -> Sequence[Awaitable]:
        """
        Our additional main task is for processing the messages of the workers (when we are the server).
        """
        return [self._process_server_messages(), self._request_stats_reporter()]

    async def _on_new_extracted_config(self, 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

    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)

    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(),
                    }
                ),
            ]
        )
        with FirstUseIdentityDatabase.create(
            asyncio.get_event_loop(), str(self._extracted_config.data_dir / "tofu.db"), sqlite_timeout=15.0
        ) as tofu_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 create_session(**CommonConfigConstants.AIOHTTP_CLIENT_SESSION_OPTS) as session:
                self.__authenticator = ServerAuthenticator(
                    session,
                    CommonConfigConstants.SERVER_DOMAIN,
                    CommonConfigConstants.SERVER_PATH,
                    self._extracted_config.data_dir,
                )
                await AbstractAioConfiguredWorkerProcess._main(self, session, tofu_db)

    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")
        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 _work(self, session: aiohttp.ClientSession, tofu_db: FirstUseIdentityDatabase):
        """
        Wait for login credentials to be available or try and use them and establish connection. If successful enter
        main websocket loop to serve remote requests.
        """
        if self.__authenticator.could_try_resume():
            await self._set_status("connecting")
            try:
                self_kr = await self.__authenticator.resume()
            except Exception:
                # session error, local storage error, decryption error or connection error
                logger.exception("Failed to resume session")
            else:
                try:
                    async with self._setup_websocket(session, self_kr) as state:
                        del self_kr
                        self.__disconnected = 0
                        await self._set_status("connected")
                        # websocket is alive, launch main processing tasks and bind their life to the state
                        done, _ = await asyncio.wait(
                            [
                                state.ensure_bound_task(self._websocket_work_receive()),
                                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, tofu_db)),
                            ],
                            return_when=asyncio.FIRST_COMPLETED,
                        )
                        for t in done:
                            if t.done() and not t.cancelled():
                                t.result()  # may raise exception
                except aiohttp.ClientResponseError:
                    # NOTE: Some errors may indicate bad configuration, TBD?
                    logger.info("Low-level client response error", exc_info=True)
                except aiohttp.ClientConnectionError:
                    # Connection errors are expected and should be ignored most of the time.
                    logger.debug("Low-level client connection error", exc_info=True)
                except aiohttp.ClientError:
                    logger.exception("Error in connection (client)")
                except NanoVersionRejected:
                    self.__update_state = 4  # No longer supported/update available.
                    self.__disconnected += 100  # no reason to retry more frequently
                    logger.exception("This Nano version is deprecated, update to a recent version to fix it.")
                except Exception:
                    logger.exception("Unexpected error")
                    raise
            self.__disconnected += 1
            await self._set_status("disconnected")
        else:
            self.__disconnected = 0
            await self._set_status("waiting-config")
        # backoff-delay when exiting in case of an error (and not cancellation)
        await asyncio.sleep(self.INCREASING_TIMEOUT.get(self.__disconnected))

    @asynccontextmanager
    async def _setup_websocket(
        self, session: aiohttp.ClientSession, self_kr: NanoSelfAccountKeyring
    ) -> AsyncContextManager[GatewayState]:
        """
        Create, connect, authenticate, initialise and yield a new websocket connection.
        """
        async with session.ws_connect(
            f"https://{CommonConfigConstants.SERVER_DOMAIN}{CommonConfigConstants.SERVER_PATH}wsn"
        ) as ws:
            state = GatewayState(ws, self._extracted_config, self_kr, self._core_requester)
            del self_kr
            try:
                self.__update_state = await state.initialise(
                    await self.__authenticator.csrf_from_session(), self.__authenticator.get_continue_token()
                )
                self.__active_state = state
                yield state
            finally:
                await state.close()  # if this raises the whole process shall exit
                self.__active_state = None

    async def _websocket_work_receive(self):
        """
        Process messages from the server. They can be either:
            - an event of some state change pushed by the server
            - a client request relayed by the server
            - a response to a request submitted by this gateway connection
        """
        while True:
            msg = await self.__active_state.ws.receive()
            if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
                break
            if msg.type != WSMsgType.BINARY:  # pragma: no cover
                raise aiohttp.ClientError("text transport not supported")
            try:
                message = msgpack_loads(msg.data)
            except Exception:  # pragma: no cover
                raise aiohttp.ClientError(f"invalid message format: {msg.type} {msg.data}")
            if not isinstance(message, dict):  # pragma: no cover
                raise aiohttp.ClientError(f"invalid message: {message}")
            if (t := message.get("type")) not in self.VALID_RESPONSE_HANDLERS:
                logger.error("Unknown message type from server: %s\n%s", t, message)
                continue
            if t == self.RESPONSE_EVENT:
                op_id = None  # na for events
            elif not isinstance(op_id := message.get("id"), int):
                logger.error("Invalid message id from server: %s\n%s", op_id, message)
                continue
            getattr(self, t)(op_id, message.get("payload"))

    # server request handler bootstraps
    def data(self, op_id, payload):
        """
        Data frame to a request made by this gateway.
        """
        self.__active_state.collect_response_data(op_id, payload)

    def error(self, op_id, payload):
        """
        Error frame to a request made by this gateway.
        """
        self.__active_state.collect_response_data(op_id, payload)
        self.__active_state.complete_response_data(op_id)

    def complete(self, op_id, payload):
        """
        End marker to a request made by this gateway.
        """
        self.__active_state.complete_response_data(op_id)

    def event(self, op_id, payload):
        """
        Some state (session, connection, license) or database entity has changed.
        """
        self.__event_type_map.get(payload.get("type"), self._event_unknown)(payload)

    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 abort(self, op_id, payload):
        """
        Cancellation of a relayed request from a client.
        """
        self.__active_state.cancel_task(op_id)

    # specific event handlers
    def _event_unknown(self, payload):
        logger.warning("Unknown event-type received from server: %s\n%s", payload.get("type"), payload)

    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)
