#!/usr/bin/env python3

import logging
import os
import signal
import sys
from argparse import Namespace
from ctypes import CDLL, byref, c_int, create_string_buffer
from ctypes.util import find_library
from logging import getLogger
from pathlib import Path
from typing import Any, Generator, List

# Constant defined in prctl.h
# See prctl(2) for more details
PR_SET_NAME = 15
PR_SET_CHILD_SUBREAPER = 36


def get_libc() -> str:
    """Find libc.so from the user's system."""
    return find_library("c") or ""


def _get_pids() -> Generator[int, Any, None]:
    yield from (int(pid.name) for pid in Path("/proc").glob("*") if pid.name.isdigit())


def get_pstree_from_pid(root_pid: int) -> set[int]:
    """Get descendent PIDs of a PID."""
    descendants: set[int] = set()
    pid_to_ppid: dict[int, int] = {}

    for pid in _get_pids():
        try:
            path = Path(f"/proc/{pid}/status")
            with path.open(mode="r", encoding="utf-8") as file:
                st_ppid = next(line for line in file if line.startswith("PPid:"))
                st_ppid = st_ppid.removeprefix("PPid:").strip()
                pid_to_ppid[pid] = int(st_ppid)
        except (FileNotFoundError, ProcessLookupError, ValueError):
            continue

    current_pid: list[int] = [root_pid]
    while current_pid:
        current = current_pid.pop()
        # Ignore. mypy flags [arg-type] due to the reuse of pid variable
        for pid, ppid in pid_to_ppid.items():  # type: ignore
            if ppid == current and pid not in descendants:
                descendants.add(pid)  # type: ignore
                current_pid.append(pid)  # type: ignore

    return descendants


def subreaper(args: Namespace, other: List[str]) -> int:
    logger = getLogger("subreaper")
    logging.basicConfig(
        format="[%(name)s] %(levelname)s: %(message)s",
        level=logging.DEBUG if args.debug else logging.INFO,
        stream=sys.stderr,
    )

    logger.debug("command: %s", args)
    logger.debug("arguments: %s", other)

    def signal_handler(sig, frame):
        logger.info("Caught '%s' signal.", signal.strsignal(sig))
        pstree = get_pstree_from_pid(os.getpid())
        for p in pstree:
            os.kill(p, sig)

    command: List[str] = [args.command, *other]
    workdir: str = args.workdir
    child_status: int = 0

    libc: str = get_libc()
    prctl = CDLL(libc).prctl
    prctl.restype = c_int
    prctl.argtypes = [
        c_int,
        # c_ulong,
        # c_ulong,
        # c_ulong,
        # c_ulong,
    ]

    proc_name = b"reaper"
    buff = create_string_buffer(len(proc_name) + 1)
    buff.value = proc_name
    prctl_ret = prctl(PR_SET_NAME, byref(buff), 0, 0, 0)
    logger.debug("prctl PR_SET_NAME exited with status: %s", prctl_ret)

    prctl_ret = prctl(PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0, 0)
    logger.debug("prctl PR_SET_CHILD_SUBREAPER exited with status: %s", prctl_ret)

    pid = os.fork()  # pylint: disable=E1101
    if pid == -1:
        logger.error("Fork failed")

    if pid == 0:
        sys.stdout.flush()
        sys.stderr.flush()
        os.chdir(workdir)
        os.execvp(command[0], command)  # noqa: S606
    else:
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)

    while True:
        try:
            child_pid, child_status = os.wait()  # pylint: disable=E1101
            logger.info("Child %s exited with wait status: %s", child_pid, child_status)
        except ChildProcessError as e:
            logger.info(e)
            break

    return child_status


if __name__ == "__main__":
    sep = sys.argv.index("--")
    argv = sys.argv[sep + 1 :]
    args = Namespace(command=argv.pop(0), workdir=os.getcwd(), debug=True)
    subreaper(args, argv)


__all__ = ["subreaper"]
