from __future__ import annotations

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

import wait_for2
from clarabot_nano_common.utility.serialize import uint2bytes, msgpack_dumps

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 ZmqMessageHandler
from clarabot_nano.common.zmq.inproc_client import AioInprocClient


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: ZmqMessageHandler, 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, message_type: bytes = REQUEST_MESSAGE):
        return await self._request([message_type], multipart_message, timeout)

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

    async def _request(
        self,
        request: List[bytes],
        multipart_message: Sequence[bytes],
        timeout=12.0,
        timeout_future: Optional[asyncio.Future] = None,
    ):
        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)
            if timeout_future is None:
                return await wait_for2.wait_for(future, timeout=timeout)
            try:
                done, _ = await asyncio.wait(
                    (future, timeout_future), timeout=timeout, return_when=asyncio.FIRST_COMPLETED
                )
            except asyncio.CancelledError:
                future.cancel()
                raise
            else:
                if future in done:
                    return future.result()
                # actual timeout or timeout-future completed, cancel request-future
                future.cancel()
                raise asyncio.TimeoutError()
        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 WorkerConnection:
    __slots__ = ("identity", "alive_until", "closed")

    def __init__(self, identity: bytes, alive_until: float):
        self.identity = identity
        self.alive_until = alive_until
        self.closed = asyncio.get_running_loop().create_future()


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: ZmqMessageHandler, keepalive_failure=KEEPALIVE_FAILURE):
        AioClientRequester.__init__(self, client, handler, setup_control_notifying_error_handler=False)
        self._keepalive_failure = keepalive_failure
        self._workers: Dict[str, Dict[bytes, WorkerConnection]] = dict()  # worker tracking
        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]] = WorkerConnection(
                routed_message[0], time.perf_counter() + self._keepalive_failure
            )
            self._worker_iterators[server_prefix] = iter(workers.values())
        else:
            workers[routed_message[0]].alive_until = 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():
            worker = workers.pop(routed_message[0], None)
            if worker is not None:
                if not workers:
                    del self._worker_iterators[server_prefix]
                else:
                    self._worker_iterators[server_prefix] = iter(workers.values())
                if not worker.closed.done():
                    worker.closed.set_result(True)

    def _get_worker(self, server_prefix: str) -> WorkerConnection:
        if not (worker_iterator := self._worker_iterators.get(server_prefix)):
            raise NoWorkerError()
        t = time.perf_counter()
        while True:
            for worker in worker_iterator:
                if worker.alive_until > t:
                    return worker
                break
            else:
                self._worker_iterators[server_prefix] = worker_iterator = iter(self._workers[server_prefix].values())
                continue
            # remove expired worker
            workers = self._workers[server_prefix]
            del workers[worker.identity]
            if not worker.closed.done():
                worker.closed.set_result(True)
            if not workers:
                del self._worker_iterators[server_prefix]
                raise NoWorkerError()
            else:
                self._worker_iterators[server_prefix] = worker_iterator = iter(workers.values())

    async def request(
        self,
        multipart_message: Sequence[bytes],
        timeout=120.0,
        message_type: bytes = REQUEST_MESSAGE,
        server_prefix: str = None,
    ):
        worker = self._get_worker(server_prefix)
        return await self._request(
            [worker.identity, message_type], multipart_message, timeout, timeout_future=worker.closed
        )

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