import copy
import importlib.util
import itertools
import logging
import os
import sys
import threading
import time
import traceback
from typing import Iterable, Dict, Optional, Sequence, List, Type, Callable

from msgpack import UnpackException

from clarabot_nano import __version__
from clarabot_nano.common.core.message import (
    MSG_TYPE_GREETING,
    MSG_TYPE_KEEPALIVE,
    MSG_TYPE_REQUEST_CONFIG_UPDATE,
    MSG_TYPE_UPDATED_CONFIG,
    MSG_TYPE_SHUTDOWN,
    MSG_TYPE_REQUEST_STATE_UPDATE,
    MSG_TYPE_GATEWAY_LOGIN,
    MSG_TYPE_GATEWAY_LOGOUT,
    MSG_TYPE_INDEXER_SYNC,
    MSG_TYPE_REQUEST_CONFIG_RESET,
    MSG_TYPE_GATEWAY_CHANGE_PASSWORD,
    MSG_TYPE_CONTROL_REQ_STATS,
    MSG_TYPE_CONTROL_INDEXER_WORK,
    MSG_TYPE_CONTROL_SERVICE_ERROR,
)
from clarabot_nano.common.log import LoggingBootstrap, CustomCallbackHandler
from clarabot_nano.common.native_process import NativeProcessAdapter
from clarabot_nano.common.zmq.auth import setup_server_auth
from clarabot_nano.common.zmq.constants import REQUEST_MESSAGE, REPLY_MESSAGE, NOTIFY_ID
from clarabot_nano.common.zmq.router_server import BiPairZmqRouterServerThread
from clarabot_nano.common.zmq_utils import zmq_authenticator
from clarabot_nano.core.config.handler import DROP_STATE_KEY, Config, ConfigLockedByRunningInstance
from clarabot_nano.core.config.schema import DEFAULT_CONFIG
from clarabot_nano.core.config.update.base import AbstractRemoteConfigUpdateHandler, RemoteConfigUpdateError
from clarabot_nano.core.config.update_handler import CONFIG_UPDATE_HANDLER_REGISTRY
from clarabot_nano.core.spawn.base import AbstractSpawn
from clarabot_nano.core.spawn.picoture import PicotureSpawn
from clarabot_nano.core.spawn.solr import SolrSpawn
from clarabot_nano.core.spawn.tight import TightlyControlledSpawn
from clarabot_nano.core.spawn.tika import TikaSpawn
from clarabot_nano.core.spawner.base import AbstractSpawner
from clarabot_nano.core.spawner.picoture import PicotureSpawner
from clarabot_nano.core.spawner.solr import SolrSpawner
from clarabot_nano.core.spawner.tight import TightPythonSpawner
from clarabot_nano.core.spawner.tika import TikaSpawner
from clarabot_nano.core.state.update.base import AbstractStateUpdateHandler
from clarabot_nano.core.state.update_handler import STATE_UPDATE_HANDLER_REGISTRY
from clarabot_nano.utility.serialize import msgpack_loads, msgpack_dumps

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


class CoreProcess(object):
    """
    Subprocess management and server thread.
    The goal was to be as simple as possible to reduce possible issues. This is not designed to handle a large
    number of subprocesses optimally.
    NOTE: There is nothing in the core or service processes to prevent the computer from going to idle-sleep. This
    must be done by the deployment configuration to specific systems, like SystemD service configs on Linux, etc.
    If Nano is started with a GUI it does try to inhibit idle-sleep.
    """

    @classmethod
    def startup(
        cls,
        local_config: Optional[str],
        headless: bool,
        bootstrap_config_dir: Optional[str],
        bootstrap_data_dir: Optional[str],
        *args,
        **kwargs,
    ):
        use_gui = importlib.util.find_spec("wx") is not None and not headless
        try:
            with Config.lock_and_load(local_config, bootstrap_config_dir, bootstrap_data_dir) as config:
                LoggingBootstrap.setup_logging(config.log_config, "core")
                NANO_INFO_LOGGER.info(f"Starting Clarabot Nano {__version__} - {config.local['source']}")
                proc = cls(config, use_gui, *args, **kwargs)
                proc.run()
                NANO_INFO_LOGGER.info("Goodbye!")  # seeing this in the log indicates clean termination
        except ConfigLockedByRunningInstance:
            sys.exit(3)
        except Exception as e:
            if use_gui:
                # When the application should have a GUI, errors should be shown using the available GUI as well.
                # This is especially true for initial startup errors, otherwise there would be no apparent feedback
                # to the user that there is some issue.
                try:
                    from clarabot_nano.core.error import ErrorApp

                    ErrorApp.display("".join(traceback.format_exception(type(e), e, e.__traceback__)))
                except Exception as wxe:
                    raise wxe from e
            raise e

    def __init__(self, config: Config, use_gui: bool):
        config.check_dependencies()  # raise if environment is not set-up or broken
        self._exit = False
        self._exit_by_explicit_user_request = False  # NOTE: Not always correct on darwin.
        self._system_session_ending = False  # NOTE: Relevant on Windows.
        LoggingBootstrap.setup_custom_sigint_sigterm_handler(self._server_exit_callback)
        self._config = config
        self._inproc_client = None
        self._introduced: Dict[bytes, TightlyControlledSpawn] = dict()
        # spawners are not extensible by config for now
        self._server_spawner = self._create_server_spawner()  # server is special for being able to set the spawn num.
        self._gateway_spawner = self._create_gateway_spawner()  # gateway is special for command requests
        self._indexer_spawner = self._create_indexer_spawner()  # indexer is special for command requests
        self._control_spawner = None
        self._spawners = [self._server_spawner, self._gateway_spawner, self._indexer_spawner]
        self._spawners.extend(self._create_spawners())
        if use_gui:
            self._control_spawner = self._create_control_spawner()
            self._spawners.append(self._control_spawner)
        self._tight_spawners = [s for s in self._spawners if isinstance(s, TightPythonSpawner)]
        self._apply_and_broadcast_config()  # apply only (no subprocess yet)
        self._request_handler_map: Dict[bytes, Callable[[TightlyControlledSpawn, Sequence[bytes]], None]] = {
            MSG_TYPE_REQUEST_CONFIG_UPDATE: self._handle_remote_config_update,
            MSG_TYPE_REQUEST_STATE_UPDATE: self._handle_state_update,
            MSG_TYPE_CONTROL_REQ_STATS: self._route_message_to_control,
            MSG_TYPE_CONTROL_INDEXER_WORK: self._route_message_to_control,
            MSG_TYPE_CONTROL_SERVICE_ERROR: self._route_message_to_control,
            MSG_TYPE_INDEXER_SYNC: self._route_message_to_indexer,
            MSG_TYPE_GATEWAY_LOGIN: self._route_message_to_gateway,
            MSG_TYPE_GATEWAY_LOGOUT: self._route_message_to_gateway,
            MSG_TYPE_GATEWAY_CHANGE_PASSWORD: self._route_message_to_gateway,
            MSG_TYPE_SHUTDOWN: self._handle_instance_shutdown_request,
            MSG_TYPE_REQUEST_CONFIG_RESET: self._handle_remote_config_reset,
        }

    def _create_server_spawner(self) -> AbstractSpawner:
        return TightPythonSpawner(
            "nano-gene-server", self._spawn_general_server_args(), dispose_callback=self._tight_dispose_callback
        )

    def _create_gateway_spawner(self) -> AbstractSpawner:
        return TightPythonSpawner(
            "nano-gateway",
            self._spawn_gateway_args(),
            dispose_callback=self._gateway_dispose_callback,
            only_spawn_when_clean=True,
        )

    def _create_indexer_spawner(self) -> AbstractSpawner:
        return TightPythonSpawner(
            "nano-indexer",
            self._spawn_indexer_args(),
            dispose_callback=self._tight_dispose_callback,
            only_spawn_when_clean=True,
        )

    def _create_control_spawner(self) -> AbstractSpawner:
        return TightPythonSpawner(
            "nano-control",
            self._spawn_control_args(),
            dispose_callback=self._tight_dispose_callback,
            only_spawn_when_clean=True,
            gui_type=True,
        )

    def _create_spawners(self) -> List[AbstractSpawner]:
        return [
            SolrSpawner(self._solr_spawn_callback, self._solr_dispose_callback, spawn_cls=self._spawn_solr_cls()),
            TikaSpawner(self._tika_spawn_callback, self._tika_dispose_callback, spawn_cls=self._spawn_tika_cls()),
            PicotureSpawner(
                self._picoture_spawn_callback, self._picoture_dispose_callback, spawn_cls=self._spawn_picoture_cls()
            ),
            TightPythonSpawner(
                "nano-fstx-server",
                self._spawn_fstx_server_args(),
                dispose_callback=self._tight_dispose_callback,
                only_spawn_when_clean=True,
            ),
        ]

    def _spawn_general_server_args(self) -> List[str]:
        return ["clarabot_nano.server.general"]

    def _spawn_gateway_args(self) -> List[str]:
        return ["clarabot_nano.gateway"]

    def _spawn_indexer_args(self) -> List[str]:
        return ["clarabot_nano.indexer"]

    def _spawn_control_args(self) -> List[str]:
        return ["clarabot_nano.control"]

    def _spawn_fstx_server_args(self) -> List[str]:
        return ["clarabot_nano.server.fstx"]

    def _spawn_solr_cls(self) -> Type[SolrSpawn]:
        return SolrSpawn

    def _spawn_tika_cls(self) -> Type[TikaSpawn]:
        return TikaSpawn

    def _spawn_picoture_cls(self) -> Type[PicotureSpawn]:
        return PicotureSpawn

    def _solr_spawn_callback(self, instance: SolrSpawn):
        self._set_state_values(solr_port=instance.port, solr_key=instance.key)

    def _solr_dispose_callback(self, instance: SolrSpawn):
        self._set_state_values(solr_port=DROP_STATE_KEY, solr_key=DROP_STATE_KEY)

    def _tika_spawn_callback(self, instance: TikaSpawn):
        self._set_state_values(tika_port=instance.port)

    def _tika_dispose_callback(self, instance: TikaSpawn):
        self._set_state_values(tika_port=DROP_STATE_KEY)

    def _picoture_spawn_callback(self, instance: PicotureSpawn):
        self._set_state_values(picoture_port=instance.port)

    def _picoture_dispose_callback(self, instance: PicotureSpawn):
        self._set_state_values(picoture_port=DROP_STATE_KEY)

    def _gateway_dispose_callback(self, instance: TightlyControlledSpawn):
        self._set_state_values(gateway_port=DROP_STATE_KEY, gateway_key=DROP_STATE_KEY)
        self._tight_dispose_callback(instance)

    def _tight_dispose_callback(self, instance: TightlyControlledSpawn):
        self._introduced.pop(instance.router_peer_identity, None)

    def _tight_instances(self) -> Iterable[TightlyControlledSpawn]:
        return itertools.chain(*[spawner.instances() for spawner in self._tight_spawners])

    def _instances(self) -> Iterable[AbstractSpawn]:
        return itertools.chain(spawner.instances() for spawner in self._spawners)

    def run(self):
        with NativeProcessAdapter.start(
            "core", self._system_exit_callback, restart_argv=self._create_restart_argv()
        ) as work_end:
            self._run(work_end)

    def _create_restart_argv(self):
        restart_argv = list(sys.argv)
        if restart_argv:
            if sys.platform.startswith("win"):
                exec_name = os.path.basename(restart_argv[0].lower())
                if exec_name not in ("clarabot-nano.exe", "clarabot-nanow.exe"):
                    # Restarting is only supported for the bundled application. Restarting a source-deployment using
                    # python may not work reliably because of the lack of environment variables and cwd support here.
                    # The frozen executables can setup their own environment without other configs being required.
                    return
            return restart_argv

    def _run(self, work_end: threading.Event):
        with zmq_authenticator() as auth:
            server_thread = BiPairZmqRouterServerThread("server", self._config.localhost_network_security)
            setup_server_auth(auth, server_thread)
            try:
                # preload
                server_port = server_thread.start()  # callback=self._server_exit_callback)
                for spawner in self._spawners:
                    if isinstance(spawner, TightPythonSpawner):
                        spawner.prepare(server_port, server_thread.get_client_auth_info())
                # enter blocking configuration/migration state if needed
                # TODO: some changes?
                # enter main
                self._work(server_thread)
            finally:
                # If we are terminating aggressively (e.g. Windows session ending) then we can acknowledge the exit
                # here. All child processes have been cleaned up as well as they could.
                work_end.set()
                server_thread.stop(20)

    def _system_exit_callback(self, session_ending: bool, *args) -> bool:
        """
        Returns a boolean indicating if shutdown is by user request. This case shall unregister restart hooks if any
        and such is applicable, but be aware that on darwin, this boolean is not always correct.
        NOTE: see notes of _server_exit_callback
        """
        if session_ending:
            self._system_session_ending = True
        return self._server_exit_callback()

    def _server_exit_callback(self, *args) -> bool:
        """
        This is invoked by different mechanisms and possible threading environments on different platforms in response
        to signals on POSIX and console events on Windows.
        IMPORTANT NOTE: (primarily on Windows) This has some undesired behaviour quirks. For example: emitting a log
        record from here with the ConcurrentLogHandler deadlocks the file-lock in the logging facility until some kind
        of timeout occurs. Until the timeout (90 seconds) all processes are get stuck on the first log call they
        encounter. After they continue all heartbeat systems have already failed, so the whole instance will crumble
        like dust. As a precaution keep the code here absolutely minimal and simple.
        """
        self._exit = True
        return self._exit_by_explicit_user_request

    def _work(self, server_thread):
        self._inproc_client = server_thread.create_inproc_client()
        logging.getLogger().addHandler(control_notify_handler := CustomCallbackHandler(self._notify_control_error))
        try:
            spawn_interval = AbstractSpawn.DEFAULT_HEALTH_TIMEOUT / 2
            next_check = time.perf_counter()
            while server_thread.is_alive():
                sockets = dict(self._inproc_client.pollin.poll(max(100, 1000 * int(next_check - time.perf_counter()))))
                if self._exit:
                    break
                if self._inproc_client.socket in sockets:
                    self._handle_unauthenticated_message(self._inproc_client.socket.recv_multipart())
                if time.perf_counter() > next_check:
                    next_check = time.perf_counter() + spawn_interval
                    for spawner in self._spawners:
                        spawner.manage(self._config)
        except Exception:
            logger.exception("Failure in primary thread.")
            raise
        finally:
            NANO_INFO_LOGGER.info("Shutting down!")
            try:
                if self._system_session_ending and sys.platform.startswith("win"):
                    # On Windows we shall exit without waiting for the child processes when the session is ending.
                    # The OS will terminate all child processes regardless of what we do.
                    pass
                else:
                    # cleanup what we can
                    cleanup_successful = False
                    try:
                        cleanup_successful = self._work_graceful_cleanup(server_thread)
                    except Exception:
                        logger.exception("Failure during final cleanup.")
                        raise
                    finally:
                        if not cleanup_successful:
                            # subprocess cleanup has not completed, use force
                            for s in self._spawners:
                                s.kill()
            finally:
                logging.getLogger().removeHandler(control_notify_handler)
                self._inproc_client.close()

    def _work_graceful_cleanup(self, server_thread) -> bool:
        if server_thread.is_alive():
            # do graceful shutdown if server is still running
            self._broadcast_message((MSG_TYPE_SHUTDOWN,))
            for _ in range(15):  # 15 seconds wait at most
                if all(s.cleanup(0.0) for s in self._spawners):
                    return True
                time.sleep(1)
        return False

    def _handle_unauthenticated_message(self, routed_message: Sequence[bytes]):
        if len(routed_message) < 2:
            return
        if routed_message[1] == MSG_TYPE_GREETING:
            self._handle_greeting(routed_message)
        else:
            instance = self._get_introduced(routed_message[0])
            if instance:
                self._handle_authenticated_message(instance, routed_message)

    def _handle_greeting(self, routed_message: Sequence[bytes]):
        """
        Greeting by subprocess is critically important as this is the two-way authentication between the processes
        using the secret keys that was passed on the stdin stream.
        """
        if len(routed_message) != 3 or routed_message[0] in self._introduced:
            return
        try:
            greet = msgpack_loads(routed_message[2])
        except (UnpackException, ValueError):
            return
        if not isinstance(greet, dict) or not isinstance(greet.get("pid"), int):
            return
        for instance in self._tight_instances():
            response = instance.greet(greet["pid"], routed_message[0])
            if response:
                self._introduced[routed_message[0]] = instance
                self._inproc_client.send_multipart(
                    (
                        routed_message[0],
                        MSG_TYPE_GREETING,
                        msgpack_dumps(response),
                        self._config.serialize(),
                    )
                )
                break

    def _get_introduced(self, identity) -> Optional[TightlyControlledSpawn]:
        """
        Get the authenticated (introduced) tight instance if it is still alive.
        """
        instance = self._introduced.get(identity)
        if instance and not instance.is_disposed():
            return instance

    def _is_introduced(self, identity) -> bool:
        """
        Determine if the client identity has been introduced properly by the greet protocol and the subprocess is alive.
        """
        return bool(self._get_introduced(identity))

    def _handle_authenticated_message(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        if routed_message[1] == MSG_TYPE_KEEPALIVE:
            self._handle_keepalive(instance, routed_message)
        elif routed_message[1] == REQUEST_MESSAGE:
            self._handle_request_message(instance, routed_message)

    def _handle_keepalive(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        if len(routed_message) != 2:
            return
        instance.last_healthy = time.perf_counter()
        self._inproc_client.send_multipart(
            (
                routed_message[0],
                MSG_TYPE_KEEPALIVE,
            )
        )

    def _handle_request_message(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        if len(routed_message) != 5:
            return
        handler = self._request_handler_map.get(t := routed_message[3])
        if handler:
            handler(instance, routed_message)
        else:
            logger.warning("Unknown core request type: %s", t)

    def _handle_remote_config_update(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        response = {}
        try:
            request = msgpack_loads(routed_message[4])
        except (UnpackException, ValueError):
            logger.exception("Failed to unpack config update request")
        else:
            if isinstance(request, dict):
                try:
                    with self._config.update_remote_config() as new_remote:
                        handler: Optional[
                            AbstractRemoteConfigUpdateHandler
                        ] = CONFIG_UPDATE_HANDLER_REGISTRY.get_handler(
                            request.get("_handler_key"),
                            DEFAULT_CONFIG["remote"],
                            new_remote,
                            copy.deepcopy(self._config.local),  # keep it safe, only pass a copy
                        )
                        handler.execute(request, instance.get_pid_prefix())  # will raise if handler is None
                except RemoteConfigUpdateError as e:
                    logger.exception("Failed to update config: %s %s", request.get("_handler_key"), e.info)
                    response["info"] = e.info
                except Exception:
                    logger.exception("Failed to update config: %s", request.get("_handler_key"))
                else:
                    success = True
        if success:
            self._config.reload_remote_config()
            self._apply_and_broadcast_config()
        response["ok"] = success
        self._reply_to_request(routed_message, msgpack_dumps(response))

    def _handle_remote_config_reset(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        try:
            self._config.reset_remote_config()
        except Exception:
            logger.exception("Failed to reset config")
        else:
            success = True
        if success:
            self._apply_and_broadcast_config()
        self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _handle_state_update(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        state_kvs = None
        try:
            request = msgpack_loads(routed_message[4])
        except (UnpackException, ValueError):
            logger.exception("Failed to unpack state update request")
        else:
            if isinstance(request, dict):
                try:
                    handler: Optional[AbstractStateUpdateHandler] = STATE_UPDATE_HANDLER_REGISTRY.get_handler(
                        request.get("_handler_key"), self._config.state
                    )
                    state_kvs = handler.get_new_state_kvs(request)  # will raise if handler is None
                    self._config.check_state_update_valid(state_kvs)
                except Exception:
                    logger.exception("Failed to update state")
                else:
                    success = True
        if success and state_kvs:
            self._set_state_values(**state_kvs)
        self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _route_message_to_indexer(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        indexer_inst = self._indexer_spawner.instances()
        if len(indexer_inst) == 1 and not indexer_inst[0].is_disposed() and indexer_inst[0].router_peer_identity:
            self._inproc_client.send_multipart(
                (indexer_inst[0].router_peer_identity, routed_message[3], routed_message[4])
            )
            success = True
        if routed_message[2] != NOTIFY_ID:
            self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _route_message_to_control(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        if self._control_spawner is None:
            success = True  # TODO: sender do not care about headless mode for now, but they might later...
        else:
            control_inst = self._control_spawner.instances()
            if len(control_inst) == 1 and not control_inst[0].is_disposed() and control_inst[0].router_peer_identity:
                self._inproc_client.send_multipart(
                    (control_inst[0].router_peer_identity, routed_message[3], routed_message[4])
                )
                success = True
        if routed_message[2] != NOTIFY_ID:
            self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _route_message_to_gateway(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        gateway_inst = self._gateway_spawner.instances()
        if len(gateway_inst) == 1 and not gateway_inst[0].is_disposed() and gateway_inst[0].router_peer_identity:
            self._inproc_client.send_multipart(
                (gateway_inst[0].router_peer_identity, routed_message[3], routed_message[4])
            )
            success = True
        self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _handle_instance_shutdown_request(self, instance: TightlyControlledSpawn, routed_message: Sequence[bytes]):
        success = False
        # TODO: add test case for denying this request
        if self._control_spawner and instance in self._control_spawner.instances():
            # only the control process may perform an instance-level shutdown
            if routed_message[-1] == b"explicit-user-request":
                # The explicitness of the shutdown is relevant on Windows, to unregister the Restart Manager hook.
                # NOTE: It can't even be determined if the request is explicit or not on darwin with current wxWidgets
                # hook setup, so this should only be utilized on Windows.
                self._exit_by_explicit_user_request = True
            self._server_exit_callback()
            success = True
        self._reply_to_request(routed_message, msgpack_dumps({"ok": success}))

    def _reply_to_request(self, routed_message: Sequence[bytes], reply: bytes):
        self._inproc_client.send_multipart(
            (
                routed_message[0],
                REPLY_MESSAGE,
                routed_message[2],
                reply,
            )
        )

    def _notify_control_error(self, msg):
        self._route_message_to_control(
            None, [b"", b"", NOTIFY_ID, MSG_TYPE_CONTROL_SERVICE_ERROR, msgpack_dumps({"msg": msg})]
        )

    def _set_state_values(self, **kv):
        self._config.update_state(kv)
        self._apply_and_broadcast_config()

    def _apply_and_broadcast_config(self):
        # apply configs that control something in the core
        self._server_spawner.set_instance_number(self._config.remote["server_process_number"])
        # send config to subprocesses
        self._broadcast_message((MSG_TYPE_UPDATED_CONFIG, self._config.serialize()))

    def _broadcast_message(self, message: tuple):
        # push message to all subprocesses (broadcast)
        for instance in self._tight_instances():
            identity = instance.router_peer_identity
            if identity and self._is_introduced(identity):
                self._inproc_client.send_multipart((identity,) + message)
