import os
import time

import psutil
import zmq
from clarabot_nano_common.common.log import LoggingBootstrap
from clarabot_nano_common.utility.serialize import msgpack_dumps, msgpack_loads

from clarabot_nano.common.core.message import MSG_TYPE_GREETING
from clarabot_nano.common.spawn.client import read_tight_python_arguments
from clarabot_nano.common.zmq_utils import polling_send_multipart


class AbstractCoreZmqClientMixin(object):
    """
    Contains the client-implementation free initialization of the core socket connection.

    Every major process in a Nano instance has to acquire and synchronize the active configuration. The current
    configuration is stored and determined by the Core. This connection is also used by the Core to track
    all subprocesses and manage them based on keepalive messages.

    Some subprocesses may even ask the Core requests if it allows it. An instance of this class will be by design
    a single-instance entity within a subprocess and shall live as long as the process runs.

    Inherit on top of ZmqDealerClientThreadMixin and AbstractZmqSocketThread or their subclasses.
    """

    def __init__(self):
        self.__initial_config = None

    def _thread_init_client_socket(self, client_socket):
        polling_send_multipart(client_socket, [MSG_TYPE_GREETING, msgpack_dumps({"pid": (pid := os.getpid())})])
        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 Core")
        message = client_socket.recv_multipart()
        if len(message) != 3 or message[0] != MSG_TYPE_GREETING:
            raise Exception("Invalid connect response from Core")
        response = msgpack_loads(message[1])
        ppid = response.get("pid")
        # Some platforms cannot spawn the children directly and there are intermediate processes.
        if not isinstance(ppid, int) or (
            ppid != os.getppid() and all(pp.pid != ppid for pp in psutil.Process(pid).parents())
        ):
            raise Exception("Could not finish handshake with Core")
        initial_config = msgpack_loads(message[2])
        if not isinstance(initial_config, dict):  # could make more strict
            raise Exception("Initial config is invalid")
        self.__initial_config = initial_config

    @classmethod
    def create(cls, callback, name_prefix, *args):
        """
        Starting the core zmq client involves:
            - acquiring the secret arguments from the parent by reading the stdin
            - launching the socket-thread (it will try to connect and hopefully succeeds)
            - acquiring the initial config from the socket thread
        """
        core_port, core_auth_data = read_tight_python_arguments()
        self = cls(core_port, core_auth_data, "core", *args)
        return self, self._init_wait_config(callback, name_prefix)

    def _init_wait_config(self, callback, name_prefix):
        LoggingBootstrap.setup_custom_sigint_sigterm_handler(self._init_signal_handler)
        self.start(callback=callback, name_prefix=name_prefix)
        while True:  # wait until initial config is acquired (and core connection is established)
            if not self.is_alive():
                raise Exception("Core connection setup failed")
            elif self.__initial_config is not None:
                break
            time.sleep(0.25)
        initial_config = self.__initial_config
        del self.__initial_config
        return initial_config

    def _init_signal_handler(self, *args):
        """
        IMPORTANT NOTE: Read the notes about signal handling callbacks at CoreProcess._server_exit_callback !
        """
        self.stop(0)
