#!/usr/bin/env python3
"""Run a command in its own process group and clean it up on timeout/cancel."""

from __future__ import annotations

import os
import signal
import subprocess
import sys
import time


GRACE_SECONDS = 10


def main() -> int:
    if len(sys.argv) < 3:
        print("usage: run-with-timeout.py <seconds> <command> [args...]", file=sys.stderr)
        return 2

    try:
        timeout_seconds = float(sys.argv[1])
    except ValueError:
        print("timeout must be a number of seconds", file=sys.stderr)
        return 2

    process = subprocess.Popen(sys.argv[2:], start_new_session=True)
    stop_signal: int | None = None
    stop_started: float | None = None
    timed_out = False

    def request_stop(signum: int, _frame: object) -> None:
        nonlocal stop_signal, stop_started
        if stop_signal is None:
            stop_signal = signum
            stop_started = time.monotonic()
            try:
                os.killpg(process.pid, signal.SIGTERM)
            except ProcessLookupError:
                pass

    for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal.signal(signum, request_stop)

    deadline = time.monotonic() + timeout_seconds
    while True:
        now = time.monotonic()
        if stop_signal is None and not timed_out and now >= deadline:
            timed_out = True
            stop_started = now
            try:
                os.killpg(process.pid, signal.SIGTERM)
            except ProcessLookupError:
                pass
        elif stop_started is not None and now - stop_started >= GRACE_SECONDS:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            stop_started = None

        try:
            result = process.wait(timeout=0.2)
            if stop_signal is not None:
                try:
                    os.killpg(process.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                return 128 + stop_signal
            if timed_out:
                try:
                    os.killpg(process.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                return 124
            return result
        except subprocess.TimeoutExpired:
            continue


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