import os
import subprocess
import sys
from pathlib import Path

import psutil

from clarabot_nano.common.spawn.server import write_tight_python_arguments
from clarabot_nano.core.spawn.base import AbstractSpawnWithThread
from clarabot_nano.core.spawn.track import SubprocessTracker


class TightlyControlledSpawn(AbstractSpawnWithThread):
    """
    Spawn a subprocess that is tightly integrated with the Core. It will greet, receive config and periodically send
    keepalive messages to indicate it is healthy. The socket communication is implemented in the process, this is only
    the data structure.
    """

    def __init__(
        self,
        index: int,
        pid_dir: Path,
        pid_prefix: str,
        core_port: int,
        core_auth_data: bytes,
        popen_args,
        gui_type=False,
    ):
        AbstractSpawnWithThread.__init__(self)
        self._spawned_pid = None
        self._index = index
        self._pid_dir = pid_dir
        self._pid_prefix = pid_prefix
        self._core_port = core_port
        self._core_auth_data = core_auth_data
        self._popen_args = popen_args
        self._gui_type = gui_type
        self.router_peer_identity = None

    def greet(self, pid: int, identity: bytes):
        if self.router_peer_identity is None and self._spawned_pid and not self.is_disposed():
            # Some platforms cannot spawn the children directly and there are intermediate processes.
            try:
                if pid == self._spawned_pid or any(
                    cp.pid == pid for cp in psutil.Process(self._spawned_pid).children()
                ):
                    self.set_healthy()
                    self.router_peer_identity = identity
                    return {"pid": os.getpid()}
            except psutil.NoSuchProcess:
                pass  # ignore child errors here, they should be temporary anyways

    def dispose(self):
        AbstractSpawnWithThread.dispose(self)
        self.router_peer_identity = None

    def kill(self):
        """
        The "stop" command already includes a kill after a timeout so nothing to be done here.
        """
        pass

    def get_pid_prefix(self) -> str:
        return f"{self._pid_prefix}-{self._index}"

    def _stop_spawn(self, *args):
        for _ in SubprocessTracker.cleanup(
            self._pid_dir,
            kill_timeout=10.0,
            pid_file_prefix=self.get_pid_prefix(),
            gui_type=True if sys.platform.startswith("win") else self._gui_type,
        ):
            pass  # we design our processes to shutdown gracefully after they receive a SIGTERM, so just return here

    def _start_spawn(self, *args, metadata=None) -> psutil.Popen:
        p = self._start_reap_popen(self._popen_args, self._gui_type, stdin=subprocess.PIPE, bufsize=0)
        self._spawned_pid = p.pid
        try:
            write_tight_python_arguments(p.stdin, self._core_port, self._core_auth_data)
        except OSError:  # pragma: no cover
            pass  # normal cleanup will work later
        SubprocessTracker.record(self._pid_dir, p, metadata=metadata, pid_file_prefix=self.get_pid_prefix())
        return p
