from multiprocessing import get_context
from multiprocessing.synchronize import Event as EventT
from typing import Callable, Optional
from typing_extensions import ParamSpec, Concatenate

import forcen_public_utils.checked as ch
from forcen_public_utils.loggers.console_logger import ConsoleLogger
from forcen_public_utils.rate import Rate
from forcen_public_utils.try_for import try_for

P = ParamSpec("P")


class ManagedProcess:
    def __init__(
        self,
        target: Callable[Concatenate[EventT, P], None],
        *args: P.args,
        **kwargs: P.kwargs,
    ):  
        self._ctx = get_context('spawn')
        self._stop_event = self._ctx.Event()
        self._handle = self._ctx.Process(target=target, args=(self._stop_event,) + args, kwargs=kwargs)
        self._handle.start()

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, exception_traceback):
        self.stop(10.0)

    def get_pid(self) -> Optional[int]:
        return self._handle.pid

    def stop(self, join_timeout: float) -> ch.MaybeVal[Optional[int]]:
        self._stop_event.set()
        self._handle.join(join_timeout)

        if not self._handle.is_alive():
            return ch.ok(self._handle.exitcode, message="exited cleanly")

        self._handle.terminate()
        if try_for(
            self._handle.is_alive,
            join_timeout,
            Rate(10),
        ):
            ConsoleLogger.get().info(f"Server terminate result: {self._handle.exitcode}\n\n")
            return ch.ok(self._handle.exitcode, message=f"Server process terminated")

        self._handle.kill()
        if try_for(
            self._handle.is_alive,
            join_timeout,
            Rate(10),
        ):
            ConsoleLogger.get().info(f"Server kill result: {self._handle.exitcode}\n\n")
            return ch.ok(self._handle.exitcode, message=f"Server process killed")

        return ch.bad(message="Failed to destroy process even with SIGTERM and SIGKILL. Giving up")
