from __future__ import annotations

import asyncio
import logging
import time
from typing import Sequence, List

import wait_for2

from clarabot_nano.common.core.message import (
    MSG_TYPE_KEEPALIVE,
    MSG_TYPE_GREETING,
    MSG_TYPE_SHUTDOWN,
    MSG_TYPE_CONTROL_SERVICE_ERROR,
)
from clarabot_nano.common.log import CustomCallbackHandler
from clarabot_nano.common.zmq.constants import REQUEST_MESSAGE, REPLY_MESSAGE, KEEPALIVE_FAILURE, NOTIFY_ID
from clarabot_nano.common.zmq.handler import AioZmqMessageHandler
from clarabot_nano.common.zmq.inproc_client import AioInprocClient
from clarabot_nano.utility.serialize import uint2bytes, msgpack_dumps


class ControlNotifierHandler(CustomCallbackHandler):
    """
    A handler class that sends an error notification to the control process.
    """

    def __init__(self, aio_core_requester: AioClientRequester):
        CustomCallbackHandler.__init__(self, self._notify_control)
        self._aio_core_requester = aio_core_requester

    def _notify_control(self, msg):
        self._aio_core_requester.notify([MSG_TYPE_CONTROL_SERVICE_ERROR, msgpack_dumps({"msg": msg})])


class AioClientRequester(object):
    """
    Wrap an AioInprocClient along a AioZmqMessageHandler to create an asynchronous request-replay pattern API.
    """

    _ID_WRAP = 2 ** 31 - 1

    def __init__(
        self, client: AioInprocClient, handler: AioZmqMessageHandler, setup_control_notifying_error_handler: bool = True
    ):
        self._request_id = 1
        self._client = client
        self._reply_futures = dict()
        self._control_notifying_error_handler = (
            ControlNotifierHandler(self) if setup_control_notifying_error_handler else None
        )
        handler.add_handler(REPLY_MESSAGE, self._reply)

    def _reply(self, multipart_message: Sequence[bytes]):
        if len(multipart_message) > 2:
            future = self._reply_futures.pop(multipart_message[1], None)
            if future:
                try:
                    future.set_result(multipart_message[2:])
                except asyncio.InvalidStateError:
                    assert future.cancelled()  # abandoned

    async def request(self, multipart_message: Sequence[bytes], timeout=12.0):
        return await self._request([REQUEST_MESSAGE], multipart_message, timeout)

    def notify(self, multipart_message: Sequence[bytes]):
        """
        Send a request message without awaiting the reply.
        """
        self._client.send_multipart([REQUEST_MESSAGE, NOTIFY_ID, *multipart_message])

    async def _request(self, request: List[bytes], multipart_message: Sequence[bytes], timeout=12.0):
        request.append(r_id := uint2bytes(self._request_id))
        if self._request_id >= self._ID_WRAP:
            self._request_id = 1
        else:
            self._request_id += 1
        request.extend(multipart_message)
        self._reply_futures[r_id] = future = asyncio.get_event_loop().create_future()
        try:
            self._client.send_multipart(request)
            return await wait_for2.wait_for(future, timeout=timeout)
        finally:
            self._reply_futures.pop(r_id, None)

    def __enter__(self):
        if self._control_notifying_error_handler:
            logging.getLogger().addHandler(self._control_notifying_error_handler)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._control_notifying_error_handler:
            root_logger = logging.getLogger()
            if exc_val:
                root_logger.error("Exiting with exception", exc_info=exc_val)
            root_logger.removeHandler(self._control_notifying_error_handler)


class NoWorkerError(Exception):
    pass


class AioWorkerPoolRequester(AioClientRequester):
    """
    A simple and somewhat naive worker pool implementation governed by the worker's keepalive heartbeats.
    This can only be used with a client and handler of a ROUTER socket. The clients are required to send a group
    identifier with the greeting and each keepalive message. The requests can be made to any of the active groups.
    """

    def __init__(self, client: AioInprocClient, handler: AioZmqMessageHandler, keepalive_failure=KEEPALIVE_FAILURE):
        AioClientRequester.__init__(self, client, handler, setup_control_notifying_error_handler=False)
        self._keepalive_failure = keepalive_failure
        self._workers = dict()  # map of worker identity to their last heartbeat
        self._worker_iterators = dict()
        handler.add_handler(MSG_TYPE_GREETING, self._handle_greeting)
        handler.add_handler(MSG_TYPE_KEEPALIVE, self._handle_keepalive)
        handler.add_handler(MSG_TYPE_SHUTDOWN, self._handle_shutdown)

    def _reply(self, multipart_message: Sequence[bytes]):
        if len(multipart_message) > 3:
            future = self._reply_futures.pop(multipart_message[2], None)
            if future:
                try:
                    future.set_result(multipart_message[3:])
                except asyncio.InvalidStateError:
                    assert future.cancelled()  # abandoned

    def _handle_greeting(self, routed_message):
        self._client.send_multipart(routed_message)  # echo the greeting for ack

    def _handle_keepalive(self, routed_message):
        if len(routed_message) < 3:
            return
        server_prefix = routed_message[2].decode("utf-8")
        if (workers := self._workers.get(server_prefix)) is None:
            self._workers[server_prefix] = workers = dict()
        if routed_message[0] not in workers:
            workers[routed_message[0]] = time.perf_counter() + self._keepalive_failure
            self._worker_iterators[server_prefix] = iter(workers.items())
        else:
            workers[routed_message[0]] = time.perf_counter() + self._keepalive_failure
        self._client.send_multipart(routed_message)  # echo the keepalive for ack

    def _handle_shutdown(self, routed_message):
        for server_prefix, workers in self._workers.items():
            if routed_message[0] in workers:
                del workers[routed_message[0]]
                if not workers:
                    del self._worker_iterators[server_prefix]
                else:
                    self._worker_iterators[server_prefix] = iter(workers.items())

    def _get_worker_identity(self, server_prefix: str) -> bytes:
        if not (worker_iterator := self._worker_iterators.get(server_prefix)):
            raise NoWorkerError()
        t = time.perf_counter()
        while True:
            for worker_identity, last_keepalive in worker_iterator:
                if last_keepalive > t:
                    return worker_identity
                break
            else:
                self._worker_iterators[server_prefix] = worker_iterator = iter(self._workers[server_prefix].items())
                continue
            # remove expired worker
            workers = self._workers[server_prefix]
            del workers[worker_identity]
            if not workers:
                del self._worker_iterators[server_prefix]
                raise NoWorkerError()
            else:
                self._worker_iterators[server_prefix] = worker_iterator = iter(workers.items())

    async def request(self, multipart_message: Sequence[bytes], timeout=12.0, server_prefix: str = None):
        return await self._request(
            [self._get_worker_identity(server_prefix), REQUEST_MESSAGE], multipart_message, timeout
        )

    def notify(self, multipart_message: Sequence[bytes], server_prefix: str = None):
        self._client.send_multipart(
            [self._get_worker_identity(server_prefix), REQUEST_MESSAGE, NOTIFY_ID, *multipart_message]
        )
