#!/usr/bin/env python3
"""Bound one process tree with Linux PID-namespace/subreaper semantics."""

from __future__ import annotations

import ctypes
import errno
import os
import pathlib
import shutil
import signal
import subprocess
import sys
import time


SIGTERM = signal.SIGTERM
SIGKILL = signal.SIGKILL
PR_SET_CHILD_SUBREAPER = 36
RECEIPT_ENV = "BEAGLE_BOUNDED_COMPLETION_RECEIPT"
NAMESPACE_ENV = "BEAGLE_PLATFORM_SUPERVISOR_NAMESPACE"
UNSHARE_ENV = "BEAGLE_PLATFORM_UNSHARE_USABLE"
interrupted_signal: int | None = None


def fail(detail: str) -> "NoReturn":
    print(f"beagle platform supervisor: {detail}", file=sys.stderr, flush=True)
    raise SystemExit(2)


def positive_seconds(text: str, name: str) -> int:
    try:
        value = int(text)
    except ValueError:
        fail(f"{name} must be a positive integer")
    if value <= 0:
        fail(f"{name} must be a positive integer")
    return value


def parse_arguments() -> tuple[int, int, list[str]]:
    arguments = sys.argv[1:]
    if len(arguments) < 4 or arguments[2] != "--":
        fail("expected SECONDS KILL-GRACE -- COMMAND [ARG ...]")
    return (
        positive_seconds(arguments[0], "deadline"),
        positive_seconds(arguments[1], "kill grace"),
        arguments[3:],
    )


def probe_unshare(unshare: str) -> bool:
    completed = subprocess.run(
        [
            unshare,
            "--user",
            "--map-current-user",
            "--pid",
            "--fork",
            "--kill-child",
            "true",
        ],
        stdin=subprocess.DEVNULL,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    return completed.returncode == 0


def enter_pid_namespace() -> None:
    if os.environ.get(NAMESPACE_ENV) == "1" or os.getpid() == 1:
        return
    unshare = shutil.which("unshare")
    memo = os.environ.get(UNSHARE_ENV)
    usable = bool(unshare) and (memo == "1" or (memo is None and probe_unshare(unshare)))
    os.environ[UNSHARE_ENV] = "1" if usable else "0"
    if not usable:
        return
    environment = os.environ.copy()
    environment[NAMESPACE_ENV] = "1"
    os.execve(
        unshare,
        [
            unshare,
            "--user",
            "--map-current-user",
            "--pid",
            "--fork",
            "--kill-child",
            sys.executable,
            str(pathlib.Path(__file__).resolve()),
            *sys.argv[1:],
        ],
        environment,
    )


def become_subreaper() -> None:
    libc = ctypes.CDLL(None, use_errno=True)
    if libc.prctl(PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0:
        failure = ctypes.get_errno()
        fail(f"could not become a child subreaper: {os.strerror(failure)}")


def parent_of(pid: int) -> int | None:
    try:
        for line in pathlib.Path(f"/proc/{pid}/status").read_text().splitlines():
            if line.startswith("PPid:"):
                return int(line.split()[1])
    except (FileNotFoundError, PermissionError, ProcessLookupError, ValueError):
        return None
    return None


def adopted_children() -> list[int]:
    own_pid = os.getpid()
    children: list[int] = []
    try:
        entries = os.listdir("/proc")
    except OSError:
        return children
    for entry in entries:
        if entry.isdigit():
            pid = int(entry)
            if pid != own_pid and parent_of(pid) == own_pid:
                children.append(pid)
    return children


def signal_owned(process_group: int | None, sig: signal.Signals) -> None:
    if os.getpid() == 1:
        try:
            os.kill(-1, sig)
        except ProcessLookupError:
            pass
        return
    if process_group is not None:
        try:
            os.killpg(process_group, sig)
        except ProcessLookupError:
            pass
    for pid in adopted_children():
        try:
            os.kill(pid, sig)
        except ProcessLookupError:
            pass


def reap_nonblocking() -> tuple[bool, bool]:
    reaped = False
    while True:
        try:
            pid, _ = os.waitpid(-1, os.WNOHANG)
        except ChildProcessError:
            return reaped, False
        except InterruptedError:
            continue
        if pid == 0:
            return reaped, True
        reaped = True


def shutdown_descendants(process_group: int | None, kill_grace: int) -> None:
    signal_owned(process_group, SIGTERM)
    grace_deadline = time.monotonic() + kill_grace
    while time.monotonic() < grace_deadline:
        _, children_remain = reap_nonblocking()
        if not children_remain:
            return
        time.sleep(0.01)
    signal_owned(process_group, SIGKILL)
    quiet_rounds = 0
    while quiet_rounds < 100:
        reaped, children_remain = reap_nonblocking()
        if not children_remain:
            return
        signal_owned(process_group, SIGKILL)
        quiet_rounds = 0 if reaped else quiet_rounds + 1
        time.sleep(0.01)
    fail("descendants remained after SIGKILL")


def write_receipt(receipt: str | None, timed_out: bool, status: int) -> None:
    if receipt is None:
        return
    pathlib.Path(receipt).write_text(
        f"subtree-reaped-v0 {'timeout' if timed_out else 'exit'} status={status}\n"
    )


def record_interruption(signum: int, _call_state: object) -> None:
    global interrupted_signal
    interrupted_signal = signum


def main() -> int:
    seconds, kill_grace, command = parse_arguments()
    receipt = os.environ.get(RECEIPT_ENV)
    if receipt:
        try:
            pathlib.Path(receipt).unlink()
        except FileNotFoundError:
            pass
    enter_pid_namespace()
    become_subreaper()
    executable = shutil.which(command[0])
    if executable is None and pathlib.Path(command[0]).is_file():
        executable = command[0]
    if executable is None:
        fail(f"command is unavailable: {command[0]}")
    label = pathlib.Path(executable).name
    print(
        f"beagle platform supervisor: {label} START deadline={seconds}s "
        f"kill-grace={kill_grace}s",
        file=sys.stderr,
        flush=True,
    )
    try:
        child = subprocess.Popen([executable, *command[1:]], start_new_session=True)
    except OSError as error:
        fail(f"could not start {command[0]}: {error}")
    for caught_signal in (signal.SIGHUP, signal.SIGINT, signal.SIGTERM):
        signal.signal(caught_signal, record_interruption)
    process_group = child.pid
    timed_out = False
    deadline = time.monotonic() + seconds
    while True:
        status = child.poll()
        if status is not None:
            break
        if interrupted_signal is not None:
            signal_owned(process_group, SIGTERM)
            try:
                child.wait(timeout=kill_grace)
            except subprocess.TimeoutExpired:
                signal_owned(process_group, SIGKILL)
                child.wait()
            status = 128 + interrupted_signal
            break
        if time.monotonic() >= deadline:
            timed_out = True
            signal_owned(process_group, SIGTERM)
            try:
                child.wait(timeout=kill_grace)
            except subprocess.TimeoutExpired:
                signal_owned(process_group, SIGKILL)
                child.wait()
            status = child.returncode
            break
        time.sleep(0.01)
    shutdown_descendants(process_group, kill_grace)
    outcome = 124 if timed_out else status
    write_receipt(receipt, timed_out, outcome)
    print(
        f"beagle platform supervisor: {label} "
        f"{'TIMEOUT' if timed_out else 'END'} status={outcome}",
        file=sys.stderr,
        flush=True,
    )
    return outcome


if __name__ == "__main__":
    raise SystemExit(main())
