# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

from __future__ import annotations

import collections
import concurrent.futures
import contextlib
import logging
import multiprocessing
import os
import signal
import socket
import sys
import threading
import time
import typing
import uuid

from cotyledon import _service_worker
from cotyledon import _utils
from cotyledon import types as t


if typing.TYPE_CHECKING:
    import types

    from cotyledon import _service

LOG = logging.getLogger(__name__)


class WorkerInfo(typing.NamedTuple):
    worker_id: t.WorkerId
    started_event: multiprocessing.synchronize.Event


ProcessInfo = typing.NewType("ProcessInfo", dict[multiprocessing.Process, WorkerInfo])
RunningServices = typing.NewType("RunningServices", dict[t.ServiceId, ProcessInfo])

OnTerminateHook: typing.TypeAlias = typing.Callable[[], None]
OnReloadHook: typing.TypeAlias = typing.Callable[[], None]
OnDeadWorkerHook: typing.TypeAlias = typing.Callable[
    [t.ServiceId, t.WorkerId, int],
    None,
]


class Hooks(typing.TypedDict):
    terminate: list[OnTerminateHook]
    reload: list[OnReloadHook]
    new_worker: list[_service_worker.OnNewWorkerHook]
    dead_worker: list[OnDeadWorkerHook]


class ServiceManager(_utils.SignalManager):
    """Manage lifetimes of services

    :py:class:`ServiceManager` acts as a master process that controls the
    lifetime of children processes and restart them if they die unexpectedly.
    It also propagate some signals (SIGTERM, SIGALRM, SIGINT and SIGHUP) to
    them.

    Each child process (:py:class:`ServiceWorker`) runs an instance of
    a :py:class:`Service`.

    An application must create only one :py:class:`ServiceManager` class and
    use :py:meth:`ServiceManager.run()` as main loop of the application.



    Usage::

        class MyService(Service):
            def __init__(self, worker_id, myconf):
                super(MyService, self).__init__(worker_id)
                preparing_my_job(myconf)
                self.running = True

            def run(self):
                while self.running:
                    do_my_job()

            def terminate(self):
                self.running = False
                gracefully_stop_my_jobs()

            def reload(self):
                restart_my_job()


        class MyManager(ServiceManager):
            def __init__(self):
                super(MyManager, self).__init__()
                self.register_hooks(on_reload=self.reload)

                conf = {'foobar': 2}
                self.service_id = self.add(MyService, 5, conf)

            def reload(self):
                self.reconfigure(self.service_id, 10)

        MyManager().run()

    This will create 5 children processes running the service MyService.

    """

    _process_runner_already_created = False

    def __init__(
        self,
        wait_interval: float = 0.01,
        graceful_shutdown_timeout: int = 60,
        mp_context: multiprocessing.context.BaseContext | None = None,
    ) -> None:
        """Creates the ServiceManager object

        :param wait_interval: time between each new process spawn
        :type wait_interval: float
        :param graceful_shutdown_timeout: timeout for graceful shutdown
        :type graceful_shutdown_timeout: int
        :param mp_context: multiprocessing context to use (defaults to 'fork')
        :type mp_context: multiprocessing.context.BaseContext | None

        By default, workers are spawned using the 'fork' start method.
        To use a different start method (e.g., 'spawn' for macOS compatibility),
        pass the context via the mp_context parameter.

        Example::

            import multiprocessing
            # Use 'spawn' start method (useful on macOS)
            manager = ServiceManager(mp_context=multiprocessing.get_context('spawn'))
            manager.add(MyService, workers=5)
            manager.run()

        """

        if self._process_runner_already_created:
            msg = "Only one instance of ServiceManager per application is allowed"
            raise RuntimeError(msg)
        ServiceManager._process_runner_already_created = True
        super().__init__()

        # We use OrderedDict to start services in adding order
        self._services: dict[
            t.ServiceId,
            _service_worker.ServiceConfig[typing.Any, typing.Any],
        ] = collections.OrderedDict()
        self._running_services: RunningServices = collections.defaultdict(dict)  # type: ignore[assignment]
        self._forktimes: list[float] = []
        self._graceful_shutdown_timeout: int = graceful_shutdown_timeout
        self._wait_interval: float = wait_interval

        self._dead = threading.Event()
        # NOTE(sileht): Set it on startup, so first iteration
        # will spawn initial workers
        self._got_sig_chld = threading.Event()
        self._got_sig_chld.set()

        self._child_supervisor: threading.Thread | None = None

        self._hooks: Hooks = {
            "terminate": [],
            "reload": [],
            "new_worker": [],
            "dead_worker": [],
        }

        _utils.set_process_title(
            "{}: master process [{}]".format(
                _utils.get_process_name(),
                " ".join(sys.argv),
            ),
        )

        # Try to create a session id if possible
        with contextlib.suppress(OSError, AttributeError):
            os.setsid()

        # Default to multiprocessing (fork-compatible) if not provided
        self.mp_context: multiprocessing.context.BaseContext = (
            mp_context or multiprocessing.get_context()
        )

        self._death_detection_pipe = self.mp_context.Pipe(duplex=False)

        if os.name == "posix":
            signal.signal(signal.SIGCHLD, self._signal_catcher)

    def register_hooks(
        self,
        on_terminate: OnTerminateHook | None = None,
        on_reload: OnReloadHook | None = None,
        on_new_worker: _service_worker.OnNewWorkerHook | None = None,
        on_dead_worker: OnDeadWorkerHook | None = None,
    ) -> None:
        """Register hook methods

        This can be callable multiple times to add more hooks, hooks are
        executed in added order. If a hook raised an exception, next hooks
        will be not executed.

        :param on_terminate: method called on SIGTERM
        :type on_terminate: callable()
        :param on_reload: method called on SIGHUP
        :type on_reload: callable()
        :param on_new_worker: method called in the child process when this one
                              is ready
        :type on_new_worker: callable(service_id, worker_id, service_obj)
        :param on_new_worker: method called when a child died
        :type on_new_worker: callable(service_id, worker_id, exit_code)

        If window support is planned, hooks callable must support
        to be pickle.pickle(). See CPython multiprocessing module documentation
        for more detail.
        """

        if on_terminate is not None:
            _utils.check_callable(on_terminate, "on_terminate")
            self._hooks["terminate"].append(on_terminate)
        if on_reload is not None:
            _utils.check_callable(on_reload, "on_reload")
            self._hooks["reload"].append(on_reload)
        if on_new_worker is not None:
            _utils.check_callable(on_new_worker, "on_new_worker")
            self._hooks["new_worker"].append(on_new_worker)
        if on_dead_worker is not None:
            _utils.check_callable(on_dead_worker, "on_dead_worker")
            self._hooks["dead_worker"].append(on_dead_worker)

    @typing.overload
    def _run_hooks(
        self,
        name: typing.Literal["dead_worker"],
        service_id: t.ServiceId,
        worker_id: t.WorkerId,
        exit_code: int,
    ) -> None: ...

    @typing.overload
    def _run_hooks(
        self,
        name: typing.Literal["new_worker"],
        service_id: t.ServiceId,
        worker_id: t.WorkerId,
        service: _service.Service,
    ) -> None: ...

    @typing.overload
    def _run_hooks(
        self,
        name: typing.Literal["reload"],
    ) -> None: ...

    @typing.overload
    def _run_hooks(
        self,
        name: typing.Literal["terminate"],
    ) -> None: ...

    def _run_hooks(
        self,
        name: typing.Literal["terminate", "reload", "new_worker", "dead_worker"],
        *args: typing.Any,
        **kwargs: typing.Any,
    ) -> None:
        _utils.run_hooks(name, self._hooks[name], *args, **kwargs)

    def add(
        self,
        service: type[_service.Service],
        workers: int = 1,
        args: typing.Any = None,  # noqa: ANN401
        kwargs: typing.Any = None,  # noqa: ANN401
    ) -> t.ServiceId:
        """Add a new service to the ServiceManager

        :param service: callable that return an instance of :py:class:`Service`
        :type service: callable
        :param workers: number of processes/workers for this service
        :type workers: int
        :param args: additional positional arguments for this service
        :type args: tuple
        :param kwargs: additional keywoard arguments for this service
        :type kwargs: dict

        :return: a service id
        :rtype: uuid.uuid4
        """
        _utils.check_callable(service, "service")
        _utils.check_workers(workers, 1)
        service_id = t.ServiceId(uuid.uuid4())
        self._services[service_id] = _service_worker.ServiceConfig(
            service_id,
            service,
            workers,
            args,
            kwargs,
        )
        return service_id

    def reconfigure(self, service_id: t.ServiceId, workers: int) -> None:
        """Reconfigure a service registered in ServiceManager

        :param service_id: the service id
        :type service_id: uuid.uuid4
        :param workers: number of processes/workers for this service
        :type workers: int
        :raises: ValueError
        """
        try:
            sc = self._services[service_id]
        except KeyError:
            msg = f"{service_id} service id doesn't exists"
            raise ValueError(msg) from None
        else:
            _utils.check_workers(workers, minimum=(1 - sc.workers))
            sc.workers = workers
            # Reset forktimes to respawn services quickly
            self._forktimes = []

    def run(self) -> None:
        """Start and supervise services workers

        This method will start and supervise all children processes
        until the master process asked to shutdown by a SIGTERM.

        All spawned processes are part of the same unix process group.
        """

        self._systemd_notify_once()
        self._child_supervisor = _utils.spawn(
            self._child_supervisor_thread,
        )
        self._wait_forever()

    def _child_supervisor_thread(self) -> None:
        while not self._dead.is_set():
            self._got_sig_chld.wait()
            self._got_sig_chld.clear()

            info = self._get_last_worker_died()
            while info is not None:
                if self._dead.is_set():
                    return
                service_id, worker_id = info
                self._start_worker(service_id, worker_id)
                info = self._get_last_worker_died()

            self._adjust_workers()

    def _on_signal_received(self, sig: int) -> None:
        if sig == _utils.SIGALRM:
            self._alarm()
        elif sig == signal.SIGINT:
            self._fast_exit()
        elif sig == signal.SIGTERM:
            self._shutdown()
        elif sig == _utils.SIGHUP:
            self._reload()
        elif sig == _utils.SIGCHLD:
            self._got_sig_chld.set()
        else:
            LOG.debug("unhandled signal %s", sig)

    def _alarm(self) -> None:
        self._fast_exit(
            reason="Graceful shutdown timeout exceeded, "
            "instantaneous exiting of master process",
        )

    def _reload(self) -> None:
        """reload all children

        posix only
        """
        self._run_hooks("reload")

        # Reset forktimes to respawn services quickly
        self._forktimes = []
        signal.signal(signal.SIGHUP, signal.SIG_IGN)
        os.killpg(0, signal.SIGHUP)
        signal.signal(signal.SIGHUP, self._signal_catcher)

    def shutdown(self) -> None:
        LOG.info("Manager shutdown requested")
        os.kill(os.getpid(), signal.SIGTERM)
        self._dead.wait()

    def _shutdown(self) -> None:
        LOG.info("Caught SIGTERM signal, graceful exiting of master process")
        signal.signal(signal.SIGTERM, signal.SIG_IGN)

        if self._graceful_shutdown_timeout > 0:
            if os.name == "posix":
                signal.alarm(self._graceful_shutdown_timeout)
            else:
                threading.Timer(self._graceful_shutdown_timeout, self._alarm).start()

        # NOTE(sileht): Stop the child supervisor
        self._dead.set()
        self._got_sig_chld.set()

        if self._child_supervisor is not None:
            self._child_supervisor.join()

        # NOTE(sileht): During startup if we receive SIGTERM, python
        # multiprocess may fork the process after we send the killpg(0)
        # To workaround the issue we sleep a bit, so multiprocess can finish
        # its work.
        with concurrent.futures.ThreadPoolExecutor() as pool:
            futures = [
                pool.submit(worker_info.started_event.wait, timeout=1)
                for processes in self._running_services.values()
                for worker_info in processes.values()
            ]
            concurrent.futures.wait(futures)

        self._run_hooks("terminate")

        LOG.debug("Killing services with signal SIGTERM")
        # NOTE(sileht): we don't have killpg so we
        # kill all known processes instead
        # NOTE(sileht): We should use CTRL_BREAK_EVENT on windows
        # when CREATE_NEW_PROCESS_GROUP will be set on child
        # process
        for process in self._child_processes:
            process.terminate()

        LOG.debug("Waiting services to terminate")
        for process in self._child_processes:
            process.join()
            process.close()

        LOG.debug("Shutdown finish")
        sys.exit(0)

    @property
    def _child_processes(self) -> list[multiprocessing.Process]:
        return [
            process
            for processes in self._running_services.values()
            for process in processes
        ]

    def _adjust_workers(self) -> None:
        for service_id, conf in self._services.items():
            running_workers = len(self._running_services[service_id])
            if running_workers < conf.workers:
                for worker_id in range(running_workers, conf.workers):
                    self._start_worker(service_id, t.WorkerId(worker_id))
            elif running_workers > conf.workers:
                for worker_id in range(conf.workers, running_workers):
                    self._stop_worker(service_id, t.WorkerId(worker_id))

    def _get_last_worker_died(self) -> tuple[t.ServiceId, t.WorkerId] | None:
        """Return the last died worker information or None"""
        for service_id in list(self._running_services.keys()):
            # We copy the list to clean the orignal one
            processes = list(self._running_services[service_id].items())
            for process, worker_info in processes:
                if not process.is_alive():
                    if process.exitcode is None:
                        raise RuntimeError("Dead process without exitcode")  # noqa: TRY003, EM101
                    self._run_hooks(
                        "dead_worker",
                        service_id,
                        worker_info.worker_id,
                        process.exitcode,
                    )
                    if process.exitcode < 0:
                        sig = _utils.signal_to_name(process.exitcode)
                        LOG.info(
                            "Child %(pid)d killed by signal %(sig)s",
                            {"pid": process.pid, "sig": sig},
                        )
                    else:
                        LOG.info(
                            "Child %(pid)d exited with status %(code)d",
                            {"pid": process.pid, "code": process.exitcode},
                        )
                    del self._running_services[service_id][process]
                    return service_id, worker_info.worker_id
        return None

    @staticmethod
    def _fast_exit(
        signo: int | None = None,
        frame: types.FrameType | None = None,
        reason: str = "Caught SIGINT signal, instantaneous exiting",
    ) -> None:
        if os.name == "posix":
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            signal.signal(signal.SIGALRM, signal.SIG_IGN)
            LOG.info(reason)
            os.killpg(0, signal.SIGINT)
            # Wait a bit as os._exit(1) will break the child pipe and the child
            # will end with SIGTERM instead of SIGINT
            time.sleep(0.1)
        else:
            # NOTE(sileht): On windows killing the master process
            # with SIGINT kill automatically children
            LOG.info(reason)
        os._exit(1)

    def _slowdown_respawn_if_needed(self) -> None:
        # Limit ourselves to one process a second (over the period of
        # number of workers * 1 second). This will allow workers to
        # start up quickly but ensure we don't fork off children that
        # die instantly too quickly.
        expected_children = sum(s.workers for s in self._services.values())
        if len(self._forktimes) > expected_children:
            if time.time() - self._forktimes[0] < expected_children:
                LOG.info("Forking too fast, sleeping")
                time.sleep(5)
            self._forktimes.pop(0)
        else:
            time.sleep(self._wait_interval)
        self._forktimes.append(time.time())

    def _start_worker(self, service_id: t.ServiceId, worker_id: t.WorkerId) -> None:
        self._slowdown_respawn_if_needed()

        started_event = self.mp_context.Event()

        # Create and run a new service
        p = _utils.spawn_process(
            _service_worker.ServiceWorker.create_and_wait,
            self.mp_context,
            started_event,
            self._services[service_id],
            service_id,
            worker_id,
            self._death_detection_pipe,
            self._hooks["new_worker"],
            self._graceful_shutdown_timeout,
        )

        self._running_services[service_id][p] = WorkerInfo(worker_id, started_event)

    def _stop_worker(self, service_id: t.ServiceId, worker_id: t.WorkerId) -> None:
        for process, worker_info in self._running_services[service_id].items():
            if worker_info.worker_id == worker_id:
                # NOTE(sileht): We should use CTRL_BREAK_EVENT on windows
                # when CREATE_NEW_PROCESS_GROUP will be set on child process
                process.terminate()

    @staticmethod
    def _systemd_notify_once() -> None:
        """Send notification once to Systemd that service is ready.

        Systemd sets NOTIFY_SOCKET environment variable with the name of the
        socket listening for notifications from services.
        This method removes the NOTIFY_SOCKET environment variable to ensure
        notification is sent only once.
        """

        notify_socket = os.getenv("NOTIFY_SOCKET")
        if notify_socket:
            if notify_socket.startswith("@"):
                # abstract namespace socket
                notify_socket = f"\0{notify_socket[1:]}"
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
            with contextlib.closing(sock):
                try:
                    sock.connect(notify_socket)
                    sock.sendall(b"READY=1")
                    del os.environ["NOTIFY_SOCKET"]
                except OSError:
                    LOG.debug("Systemd notification failed", exc_info=True)