import asyncio
import logging
from concurrent.futures.thread import ThreadPoolExecutor
from contextlib import contextmanager, asynccontextmanager
from functools import partial
from typing import Sequence, Type, Optional, List, AsyncContextManager

import zmq

from clarabot_nano.common.config.client import AbstractAioConfiguredWorkerProcess
from clarabot_nano.common.core.message import MSG_TYPE_GREETING, MSG_TYPE_SHUTDOWN
from clarabot_nano.common.crypto.keyring import NanoSelfAccountKeyring, NanoPeerAccountKeyring, NanoRoomKeyring
from clarabot_nano.common.handler.registry import NanoHandlerClassRegistry
from clarabot_nano.common.protocol.errors import NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE, NANO_REQUEST_ERROR
from clarabot_nano.common.protocol.request import NanoRequestContext
from clarabot_nano.common.protocol.server import NanoServerResponse
from clarabot_nano.common.zmq.constants import REQUEST_MESSAGE, REPLY_MESSAGE
from clarabot_nano.common.zmq.dealer_client import AioZmqDealerClientThread
from clarabot_nano.common.zmq.inproc_client import AioInprocClient
from clarabot_nano.common.zmq_utils import polling_send_multipart
from clarabot_nano.server.base.config import ServerConfig
from clarabot_nano.server.base.handler.base import AbstractNanoServerHandler, ServerHandleHelper
from clarabot_nano.utility.aio import run_tasks, force_wait_predicate
from clarabot_nano.utility.serialize import msgpack_loads, msgpack_dumps

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


class GatewayZmqClient(AioZmqDealerClientThread):
    def __init__(self, loop, auth_data, port, inproc_bind_name, server_prefix: str):
        AioZmqDealerClientThread.__init__(
            self, loop, port, auth_data, inproc_bind_name, (server_prefix.encode("utf-8"),)
        )
        self.connected = False

    def _thread_init_client_socket(self, client_socket):
        polling_send_multipart(client_socket, [MSG_TYPE_GREETING])
        poll = zmq.Poller()
        poll.register(client_socket, zmq.POLLIN)
        sockets = dict(poll.poll(int(self._keepalive_failure * 1000)))
        if client_socket not in sockets:
            raise Exception("Could not connect to Gateway")
        message = client_socket.recv_multipart()
        if len(message) != 1 or message[0] != MSG_TYPE_GREETING:
            raise Exception("Invalid connect response from Gateway")
        self.connected = True

    def _thread_socket_loop(self, client_socket, inproc_socket):
        try:
            AioZmqDealerClientThread._thread_socket_loop(self, client_socket, inproc_socket)
        finally:
            try:
                client_socket.send_multipart([MSG_TYPE_SHUTDOWN], flags=zmq.NOBLOCK)
            except Exception:  # pragma: no cover
                pass  # continue exit if graceful shutdown cannot be sent to gateway


class GatewayConnection(GatewayZmqClient):
    def __init__(self, loop, *args, thread_name_prefix=None):
        self.exit_event = asyncio.Event()
        self.client = None
        GatewayZmqClient.__init__(self, loop, *args)
        self.start(
            callback=partial(loop.call_soon_threadsafe, self.exit_event.set),
            name_prefix=thread_name_prefix,
        )

    async def wait_valid(self, key, port) -> bool:
        if self._auth_data != key or self._port != port:
            return False
        # wait until gateway connection is established or fails
        while not self.connected:
            if not self.is_alive():
                await asyncio.sleep(1)
                return False  # error reason should be logged from thread, just exit and retry
            await asyncio.sleep(0.1)
        return not self.exit_event.is_set() and self.is_alive()

    def create_inproc_client(self):
        if not self.client:
            self.client = GatewayZmqClient.create_inproc_client(self)
        return self.client


class ServerProcess(AbstractAioConfiguredWorkerProcess):
    _extracted_config: ServerConfig = None  # to fix type hinting

    @classmethod
    @contextmanager
    def create(cls, server_prefix: str):
        with ThreadPoolExecutor(thread_name_prefix="TP-") as pool:
            yield cls(pool, server_prefix)

    def __init__(self, pool: ThreadPoolExecutor, server_prefix: str, log_with_pid: bool = True):
        self._pool = pool
        self._server_prefix = server_prefix
        self._handler_registry = NanoHandlerClassRegistry(self._init_get_handler_classes())
        extracted_config_cls = self._get_extracted_config_cls()
        if not issubclass(extracted_config_cls, ServerConfig):  # pragma: no cover
            raise Exception("Invalid extracted config cls for server")
        self._gateway_connection: GatewayConnection = None
        AbstractAioConfiguredWorkerProcess.__init__(
            self, server_prefix, log_with_pid, extracted_config_cls=extracted_config_cls
        )

    async def _main_wrapper(self):
        # The requester is used for error reporting only by some of the servers.
        self._core_requester = self.create_core_requester()
        with self._core_requester:
            await AbstractAioConfiguredWorkerProcess._main_wrapper(self)

    def _get_extracted_config_cls(self):
        return ServerConfig

    def _init_get_handler_classes(self) -> Sequence[Type[AbstractNanoServerHandler]]:
        raise NotImplementedError()

    async def _on_new_extracted_config(self, old_extracted_config: Optional[ServerConfig]):
        ec = self._extracted_config  # TODO: check & document why this optimization is possible and correct
        if (
            old_extracted_config is None
            or ec.gateway_port != old_extracted_config.gateway_port
            or ec.gateway_key != old_extracted_config.gateway_key
        ):
            self.queue_main_event({"type": "new-configs"})

    async def _process_main_event(self, event):
        """
        Work task is not running, execute command requests or events made from config changes.
        """
        t = event.get("type")
        if t == "new-configs":
            pass  # new config has the updated access data, that will be used for new work task
        else:  # pragma: no cover
            raise Exception(f"unknown server main event {t}")

    async def _main(self, *args):
        try:
            await AbstractAioConfiguredWorkerProcess._main(self, *args)
        finally:
            await self._cleanup_gateway_connection(final=True)

    async def _work(self):
        ec = self._extracted_config
        gateway = await self._setup_gateway_connection(ec.gateway_port, ec.gateway_key)
        if not gateway:
            return
        async with self._create_handle_helper(ec) as helper:
            del ec  # free
            work_tasks = self._create_work_tasks(gateway.client, helper)
            work_tasks.append(asyncio.ensure_future(self._process_gateway_messages(gateway.client, helper)))
            work_tasks.append(asyncio.ensure_future(gateway.exit_event.wait()))
            await run_tasks(work_tasks)

    async def _setup_gateway_connection(self, gateway_port, gateway_key):
        # Cleanup previous connection if it is not valid or alive.
        if self._gateway_connection:
            if not await self._gateway_connection.wait_valid(gateway_key, gateway_port):
                await self._cleanup_gateway_connection()
        # Wait if we have no gateway available.
        if not gateway_port or not gateway_key:
            while True:
                # This will block if no gateway config is available when starting.
                # When a gateway starts up this work task will end and a new one will start with the new config.
                await asyncio.sleep(2.5)
        # Create gateway connection and clients.
        if not self._gateway_connection:
            self._gateway_connection = GatewayConnection(
                self._loop,
                gateway_key,
                gateway_port,
                "gwwk",
                self._server_prefix,
                thread_name_prefix=self.__class__.__name__,
            )
        # Wait until the new client is ready.
        if not await self._gateway_connection.wait_valid(gateway_key, gateway_port):
            await self._cleanup_gateway_connection()
        else:
            self._gateway_connection.create_inproc_client()
            return self._gateway_connection

    async def _cleanup_gateway_connection(self, final=False):
        if self._gateway_connection:
            if self._gateway_connection.client:
                self._gateway_connection.client.close()
            # asyncio-compatible thread join
            self._gateway_connection.stop(0)
            try:
                await force_wait_predicate(self._gateway_connection.is_finished, timeout=7.5 if final else 15.0)
            except asyncio.TimeoutError:  # fatal error; cannot stop the gateway zmq client thread
                self._core_exit_callback()
                raise
            finally:  # Prevent double-cleanup of the GatewayConnection object, it would not work anyways.
                self._gateway_connection = None

    @asynccontextmanager
    async def _create_handle_helper(self, extracted_config) -> AsyncContextManager[ServerHandleHelper]:
        """
        Used from the work task so this can use the extracted config safely.
        """
        yield ServerHandleHelper(self._pool)

    def _create_work_tasks(self, gateway_client: AioInprocClient, helper: ServerHandleHelper) -> List[asyncio.Future]:
        return []

    async def _process_gateway_messages(self, gateway_client, helper):
        request_tasks = set()

        def cleanup_task_cb(f):
            request_tasks.remove(f)

        try:
            while True:
                messages = await gateway_client.recv_multipart_messages()
                for message in messages:
                    if message[0] == REQUEST_MESSAGE:
                        request_tasks.add(
                            task := self._loop.create_task(
                                self._process_gateway_request(gateway_client, helper, message)
                            )
                        )
                        task.add_done_callback(cleanup_task_cb)
                    else:  # pragma: no cover
                        raise Exception("Unknown message from Gateway")
        finally:
            for t in request_tasks:
                t.cancel()
            await asyncio.gather(*request_tasks, return_exceptions=True)

    async def _process_gateway_request(self, gateway_client, helper, message):
        """
        Execution if already in a task specific to the request.
        Message frames:
            0: REQUEST-MARKER
            1: REQUEST-ID for tracking by the gateway
            2: context
            3: payload
        """
        try:
            try:
                context = msgpack_loads(message[2])
                payload = msgpack_loads(message[3])
                # basic pre-checking
                if not isinstance(context, dict):
                    raise Exception(f"Context type is {type(context)}")
                elif not isinstance(payload, dict):
                    raise Exception(f"Payload type is {type(payload)}")
            except Exception:
                # empty response means pre-processing has failed
                gateway_client.send_multipart([REPLY_MESSAGE, message[1], msgpack_dumps({})])
                return
            try:
                context[NanoRequestContext.OWNER_KR] = NanoSelfAccountKeyring.load_internal(
                    context[NanoRequestContext.OWNER_KR]
                )
                if raw_peer_kr := context.get(NanoRequestContext.PEER_KR):
                    context[NanoRequestContext.PEER_KR] = NanoPeerAccountKeyring.load_internal(raw_peer_kr)
                if room := context.get(NanoRequestContext.ROOM):
                    room[NanoRequestContext.ROOM_KR] = NanoRoomKeyring.load_internal(room[NanoRequestContext.ROOM_KR])
            except Exception:
                logger.exception("Unexpected error while preloading context")
                response = {NanoServerResponse.ERROR: NANO_REQUEST_ERROR}
            else:
                response = await self._handler_registry.invoke(helper, self._extracted_config, context, payload)
            # Frame the payload outside of the custom data if any. (small serialization optimization)
            payload = response.pop(NanoServerResponse.PAYLOAD, None)
            response_frames = [REPLY_MESSAGE, message[1], msgpack_dumps(response)]
            if payload:
                response_frames.append(payload)
            gateway_client.send_multipart(response_frames)
        except asyncio.CancelledError:
            # Send a temporary error if possible when the request-task is cancelled.
            # This could work instead of the client waiting for a longer timeout.
            try:
                gateway_client.send_multipart(
                    [
                        REPLY_MESSAGE,
                        message[1],
                        msgpack_dumps({NanoServerResponse.ERROR: NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE}),
                    ]
                )
            except Exception:  # pragma: no cover
                pass  # We could already be shutting down, not much we can do here.
            raise
