import os
import signal
import threading
from contextlib import nullcontext
from multiprocessing import Pipe
from multiprocessing.context import SpawnProcess
from typing import Any, Callable, ParamSpec

from .logger import logger

UNIX_SIGNALS = {
    getattr(signal, f"SIG{x}"): x
    for x in "HUP QUIT TTIN TTOU USR1 USR2 WINCH".split()
    if hasattr(signal, f"SIG{x}")
}

P = ParamSpec("P")


class ProcessParameters:
    def __init__(
        self,
        f: Callable[P, Any],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> None:
        self.f = f
        self.args = args
        self.kwargs = kwargs


class Process:
    def __init__(self, parameters: ProcessParameters) -> None:
        self.parameters = parameters
        self.parent_conn, self.child_conn = Pipe()
        self.process = SpawnProcess(target=self.target, daemon=True)

    def ping(self, timeout: float = 5) -> bool:
        try:
            self.parent_conn.send(b"ping")
            if self.parent_conn.poll(timeout):
                self.parent_conn.recv()
                return True
        except IOError:
            pass
        return False

    def pong(self) -> None:
        self.child_conn.recv()
        self.child_conn.send(b"pong")

    def always_pong(self) -> None:
        while True:
            self.pong()

    def target(self) -> Any:
        if os.name == "nt":
            # Windows doesn't support SIGTERM, so we use SIGBREAK instead.
            # And then we raise SIGTERM when SIGBREAK is received.
            # https://learn.microsoft.com/zh-cn/cpp/c-runtime-library/reference/signal?view=msvc-170
            signal.signal(
                signal.SIGBREAK, lambda sig, frame: signal.raise_signal(signal.SIGTERM)
            )

        threading.Thread(target=self.always_pong, daemon=True).start()
        return self.parameters.f(*self.parameters.args, **self.parameters.kwargs)

    def is_alive(self, timeout: float = 5) -> bool:
        if not self.process.is_alive():
            return False

        return self.ping(timeout)

    def start(self) -> None:
        self.process.start()
        logger.info("Started child process [{}]".format(self.process.pid))

    def terminate(self) -> None:
        if self.process.exitcode is not None:
            return
        assert self.process.pid is not None
        if os.name == "nt":
            # Windows doesn't support SIGTERM.
            # So send SIGBREAK, and then in process raise SIGTERM.
            os.kill(self.process.pid, signal.CTRL_BREAK_EVENT)
        else:
            os.kill(self.process.pid, signal.SIGTERM)
        logger.info("Terminated child process [{}]".format(self.process.pid))

        self.parent_conn.close()
        self.child_conn.close()

    def terminate_quickly(self) -> None:
        if self.process.exitcode is not None:
            return
        assert self.process.pid is not None
        if os.name == "nt":
            os.kill(self.process.pid, signal.CTRL_C_EVENT)
        else:
            os.kill(self.process.pid, signal.SIGINT)
        logger.info("Terminated quickly child process [{}]".format(self.process.pid))

        self.parent_conn.close()
        self.child_conn.close()

    def kill(self) -> None:
        # In Windows, the method will call `TerminateProcess` to kill the process.
        # In Unix, the method will send SIGKILL to the process.
        self.process.kill()

    def join(self) -> None:
        logger.info("Waiting for child process [{}]".format(self.process.pid))
        self.process.join()

    @property
    def pid(self) -> int | None:
        return self.process.pid


class MultiProcessManager:
    def __init__(self, processes_num: int, process_parameters: ProcessParameters):
        self.processes_num = processes_num
        self.process_parameters = process_parameters
        self.processes: list[Process] = []

        self.should_exit = threading.Event()

        self.signal_queue: list[int] = []
        for sig in UNIX_SIGNALS:
            signal.signal(sig, lambda sig, frame: self.signal_queue.append(sig))

        # Sent by Ctrl+C.
        signal.signal(signal.SIGINT, lambda sig, frame: self.handle_int())
        # Sent by `kill <pid>`. Not sent on Windows.
        signal.signal(signal.SIGTERM, lambda sig, frame: self.handle_term())
        if os.name == "nt":
            # Sent by `Ctrl+Break` on Windows.
            signal.signal(signal.SIGBREAK, lambda sig, frame: self.handle_break())

    def init_processes(self) -> None:
        for _ in range(self.processes_num):
            process = Process(self.process_parameters)
            process.start()
            self.processes.append(process)

    def terminate_all(self) -> None:
        for process in self.processes:
            process.terminate()

    def terminate_all_quickly(self) -> None:
        for process in self.processes:
            process.terminate_quickly()

    def join_all(self) -> None:
        for process in self.processes:
            process.join()

    def restart_all(self, use_kill: bool = True) -> None:
        for idx, process in enumerate(tuple(self.processes)):
            if use_kill:
                process.kill()
            else:
                process.terminate()
            process.join()
            del self.processes[idx]
            process = Process(self.process_parameters)
            process.start()
            self.processes.append(process)

    def mainloop(self) -> None:
        logger.info("Started parent process [{}]".format(os.getpid()))

        self.init_processes()

        while not self.should_exit.wait(0.5):
            self.handle_signals()
            self.keep_subprocess_alive()

        self.join_all()

        logger.info("Stopped parent process [{}]".format(os.getpid()))

    def keep_subprocess_alive(self) -> None:
        if self.should_exit.is_set():
            return  # parent process is exiting, no need to keep subprocess alive

        for idx, process in enumerate(tuple(self.processes)):
            if process.is_alive():
                continue

            process.kill()  # process is hung, kill it
            process.join()
            logger.info("Child process [{}] died".format(process.pid))
            del self.processes[idx]
            process = Process(self.process_parameters)
            process.start()
            self.processes.append(process)

    def handle_signals(self) -> None:
        for sig in tuple(self.signal_queue):
            self.signal_queue.remove(sig)
            sig_name = UNIX_SIGNALS[sig]
            sig_handler = getattr(self, f"handle_{sig_name.lower()}", None)
            if sig_handler is not None:
                sig_handler()
            else:
                logger.info(f"Received signal [{sig_name}], but nothing to do")

    def handle_int(self) -> None:
        logger.info("Received SIGINT, quickly exiting")
        self.should_exit.set()
        self.terminate_all_quickly()

    def handle_term(self) -> None:
        logger.info("Received SIGTERM, exiting")
        self.should_exit.set()
        self.terminate_all()

    def handle_break(self) -> None:
        logger.info("Received SIGBREAK, exiting")
        self.should_exit.set()
        self.terminate_all()

    def handle_hup(self) -> None:
        logger.info("Received SIGHUP, restarting processes")
        self.restart_all(use_kill=False)

    def handle_ttin(self) -> None:
        logger.info("Received SIGTTIN, increasing processes")
        self.processes_num += 1
        process = Process(self.process_parameters)
        process.start()
        self.processes.append(process)

    def handle_ttou(self) -> None:
        logger.info("Received SIGTTOU, decreasing processes")
        if self.processes_num <= 1:
            logger.info("Cannot decrease processes any more")
            return
        self.processes_num -= 1
        process = self.processes.pop()
        process.terminate()
        process.join()


def multiprocess(
    processes_num: int,
    process_parameters: ProcessParameters,
    watchfiles: str | None,
) -> None:
    processes_manager = MultiProcessManager(processes_num, process_parameters)

    if watchfiles is not None:
        from .reloader import listen_for_changes

        contextmanager = listen_for_changes(watchfiles, processes_manager.restart_all)
    else:
        contextmanager = nullcontext()

    with contextmanager:
        processes_manager.mainloop()
