from __future__ import annotations

import asyncio
import logging
import os
from contextlib import asynccontextmanager
from functools import partial
from typing import Sequence, Awaitable, Type, Optional

from clarabot_nano.common.config.base import AbstractCoreZmqClientMixin
from clarabot_nano.common.config.extract.base import AbstractExtractedConfig
from clarabot_nano.common.core.message import MSG_TYPE_UPDATED_CONFIG, MSG_TYPE_SHUTDOWN
from clarabot_nano.common.log import LoggingBootstrap
from clarabot_nano.common.native_process import NativeProcessAdapter
from clarabot_nano.common.zmq.dealer_client import AioZmqDealerClientThread
from clarabot_nano.common.zmq.handler import QueuedAioZmqMessageHandler
from clarabot_nano.common.zmq.inproc_client import AioInprocClient
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.utility.serialize import msgpack_loads

NANO_INFO_LOGGER = logging.getLogger("nano/info")


class AioCoreZmqClient(AbstractCoreZmqClientMixin, AioZmqDealerClientThread):
    """
    Asyncio-integrated core zmq client.
    """

    def __init__(self, core_port, core_auth_data, inproc_bind_name, loop):
        AbstractCoreZmqClientMixin.__init__(self)
        AioZmqDealerClientThread.__init__(self, loop, core_port, core_auth_data, inproc_bind_name)


class AbstractAioConfiguredProcess(object):
    """
    The config should be acquired by .get_config() for any message handler that will need it. The returned
    """

    def __init__(self, process_name: str, log_with_pid: bool = True):
        self._loop = asyncio.get_event_loop()
        self.__config: dict = None
        self.__config_revision = 1
        self._config_change_event = asyncio.Event()
        self._core_exit_event = asyncio.Event()
        self._process_name = process_name
        # initialise core connection
        self.core, initial_config = AioCoreZmqClient.create(
            self._core_exit_callback, self.__class__.__name__, self._loop
        )
        # update signal handler with asyncio-event
        LoggingBootstrap.setup_custom_sigint_sigterm_handler(self._signal_handler)
        self.__raw_set_config(initial_config)
        self._core_client: AioInprocClient = None
        self._core_handler = QueuedAioZmqMessageHandler(
            {
                MSG_TYPE_SHUTDOWN: self._handle_core_shutdown,
                MSG_TYPE_UPDATED_CONFIG: self._handle_core_updated_config,
            }
        )
        # setup logging
        LoggingBootstrap.setup_logging(self.__config["local"]["log_config"], process_name, with_pid=log_with_pid)
        NANO_INFO_LOGGER.info("Joined Core Instance, starting... (parent: %s)", os.getppid())

    def __raw_set_config(self, config):
        """
        Called from a different thread than the asyncio event loop when initialising, no buffering or event handling.
        """
        self.__config_revision += 1
        config["_local_revision"] = self.__config_revision
        self.__config = config

    def __set_config(self, config):
        """
        Called from the asyncio event loop thread.
        """
        self.__raw_set_config(config)
        self._config_change_event.set()
        self._config_change_event.clear()

    def get_config(self):
        """
        The config shall be acquired only by this method. The returned object is not going to change.
        Any thread, process or message handler that will need the config must get it at the beginning of their work
        and use that config throughout. In case the work would take a long time and configuration update tracking
        is needed for consistency/security, they must track it themselves. They can observe the config change event or
        compare the "_local_revision" value of the old config with a newly acquired config.
        """
        return self.__config

    def _signal_handler(self, *args):
        """
        IMPORTANT NOTE: Read the notes about signal handling callbacks at CoreProcess._server_exit_callback !
        Stopping the core thread will eventually force this process to be terminated, but for better behaviour this
        can be overridden to implement additional exit handling.
        """
        self.core.stop(0)
        self._core_exit_callback()  # faster response to signal by setting the core-exit event in advance

    def _core_exit_callback(self, *args):
        """
        This is invoked from the core-connection thread or a native process adapter.
        """
        if not self._loop.is_closed():
            self._loop.call_soon_threadsafe(partial(self._core_exit_event.set))

    async def run(self):
        with NativeProcessAdapter.start(self._process_name, self._core_exit_callback):
            self._core_client = self.core.create_inproc_client()
            try:
                await self._main_wrapper()
            finally:
                self.core.stop()
                self._core_client.close()

    async def _main_wrapper(self):
        NANO_INFO_LOGGER.info("Entering main task processing...")
        main_tasks = [
            self._log_running(),
            self._process_core_messages(),
            self._core_exit_event.wait(),
            self._main(),
        ]
        async with self._create_main_tasks() as custom_main_tasks:
            main_tasks.extend(custom_main_tasks)
            main_tasks = [asyncio.ensure_future(coro_or_future) for coro_or_future in main_tasks]
            early_stopping_tasks = {main_tasks[0]}  # the log_running task is expected to exit soon
            main_tasks = set(main_tasks)
            # wait until any main task completes
            while True:
                done, pending = await asyncio.wait(main_tasks, return_when=asyncio.FIRST_COMPLETED)
                if early_stopping_tasks:
                    all_done_are_early_tasks = True
                    for t in done:
                        if t not in early_stopping_tasks or t.exception() is not None:
                            all_done_are_early_tasks = False
                            break
                    if all_done_are_early_tasks:
                        early_stopping_tasks.difference_update(done)
                        main_tasks.difference_update(done)
                        continue
                break
            self._signal_handler()
            for t in pending:
                t.cancel()
            await asyncio.gather(*main_tasks, return_exceptions=True)
            for t in done:
                e = t.exception()
                if e:
                    raise Exception("A main task failed") from e
            NANO_INFO_LOGGER.info("Shutting down, bye!")

    async def _log_running(self):
        await asyncio.sleep(1.5 * self.core.keepalive_failure_seconds)
        NANO_INFO_LOGGER.info("All systems seem to work, going into the background.")

    @asynccontextmanager
    async def _create_main_tasks(self) -> Sequence[Awaitable]:
        """
        Additional main tasks may be defined for the process here.
        """
        yield []

    async def _process_core_messages(self):
        async with self._core_handler:
            while True:
                messages = await self._core_client.recv_multipart_messages()
                for message in messages:
                    await self._core_handler.handle(message)

    def _handle_core_shutdown(self, message):
        self._signal_handler()

    async def _handle_core_updated_config(self, message):
        if len(message) != 2:
            raise Exception("invalid config message")
        await self._apply_new_config(msgpack_loads(message[1]))

    async def _apply_new_config(self, config):
        """
        There are two opposing goals that we need to reconcile here.
            - A config change can be critical to security, so applying it ASAP may be important.
            - Most changes are okay to be applied after the current running operations complete.
        It is also a question if the updated config is relevant at all for the current process.

        The default implementation just replaces the config on the process. This should be used by subclasses to
        change the config if is desired.
        """
        self.__set_config(config)

    def create_core_requester(self) -> AioClientRequester:
        """
        Having a requester is optional, not all subprocesses will need to have requests for the core.
        """
        return AioClientRequester(self._core_client, self._core_handler)

    async def _main(self):
        raise NotImplementedError()


class AbstractAioConfiguredWorkerProcess(AbstractAioConfiguredProcess):
    """
    Add a reusable worker task management by event processing. This can be used to soft-restart in response to
    some command or config change.
    """

    def __init__(
        self, process_name, log_with_pid=True, extracted_config_cls: Optional[Type[AbstractExtractedConfig]] = None
    ):
        # main event queue to manage things that must reset the work task
        self.__main_event_queue = asyncio.Queue()
        self.__extracted_config_cls = extracted_config_cls
        self._extracted_config = None

        AbstractAioConfiguredProcess.__init__(self, process_name, log_with_pid=log_with_pid)

    async def _apply_new_config(self, config):
        await AbstractAioConfiguredProcess._apply_new_config(self, config)
        if self.__extracted_config_cls is not None:
            ec = self.__extracted_config_cls(config)
            if not ec.matches(self._extracted_config):
                old_extracted_config = self._extracted_config
                self._extracted_config = ec
                await self._on_new_extracted_config(old_extracted_config)

    async def _on_new_extracted_config(self, old_extracted_config):
        """
        A new config was received and extracted. It is available as self._extracted_config along with the old one.
        """

    def queue_main_event(self, event):
        self.__main_event_queue.put_nowait(event)

    async def _main_wrapper(self):
        """
        Initialise the extracted config from the base config before starting main tasks.
        """
        await self._apply_new_config(self.get_config())
        await AbstractAioConfiguredProcess._main_wrapper(self)

    async def _main(self, *args):
        """
        Work task management and main event processing.
        """
        await self._process_main_events()  # there could already be some events made from initial config
        event_future = asyncio.ensure_future(self.__main_event_queue.get())
        work_future = asyncio.ensure_future(self._work(*args))
        try:
            while True:
                done, pending = await asyncio.wait([event_future, work_future], return_when=asyncio.FIRST_COMPLETED)
                if work_future not in done:
                    # We can never reuse a work coroutine. Either it failed by some error or config may change
                    # by the event queue.
                    work_future.cancel()
                try:
                    await work_future
                except asyncio.CancelledError:
                    pass
                if event_future in done:
                    await self._process_main_events(await event_future)
                    event_future = asyncio.ensure_future(self.__main_event_queue.get())
                # after event processing restart the worker
                work_future = asyncio.ensure_future(self._work(*args))
        finally:
            event_future.cancel()
            work_future.cancel()
            await asyncio.gather(event_future, work_future, return_exceptions=True)

    async def _process_main_events(self, triggering_event=None):
        """
        Called from the main loop when the work task is not running.
        """
        if triggering_event:
            await self._process_main_event(triggering_event)
        try:
            while True:
                await self._process_main_event(self.__main_event_queue.get_nowait())
        except asyncio.QueueEmpty:
            pass

    async def _process_main_event(self, event):
        """
        Work task is not running, execute command requests or events made from config changes.
        """
        raise NotImplementedError()

    async def _work(self, *args):
        """
        Work based on extracted config, the async task of this will be terminated between config changes.
        NOTE: The self._extracted_config is volatile, so all access shall be done through a local reference to it.
        """
        raise NotImplementedError()
