#!/usr/bin/env python3
"""Run a build command and terminate its process group at a free-space floor."""

from __future__ import annotations

import argparse
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

GIB = 1024**3
STOPPED_FOR_DISK_EXIT = 75
HANDLED_SIGNALS = (signal.SIGINT, signal.SIGTERM, signal.SIGHUP)


class GuardSignal(Exception):
    def __init__(self, signum: int) -> None:
        super().__init__(signum)
        self.signum = signum


def free_bytes(path: Path) -> int:
    return shutil.disk_usage(path).free


def process_group_exists(pgid: int) -> bool:
    try:
        os.killpg(pgid, 0)
    except (ProcessLookupError, PermissionError):
        # Darwin can report EPERM for a process group whose final members are
        # exiting/zombies even though this process created the group.
        return False
    return True


def terminate_group(process: subprocess.Popen[bytes], timeout: float = 10.0) -> None:
    pgid = process.pid
    if process_group_exists(pgid):
        try:
            os.killpg(pgid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    deadline = time.monotonic() + timeout
    while process_group_exists(pgid) and time.monotonic() < deadline:
        process.poll()
        time.sleep(0.05)
    if process_group_exists(pgid):
        try:
            os.killpg(pgid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    process.wait()


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", required=True, type=Path)
    parser.add_argument("--minimum-free-gib", required=True, type=int)
    parser.add_argument("--poll-seconds", type=float, default=15.0)
    parser.add_argument("--cwd", type=Path)
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args(argv)
    if args.command[:1] == ["--"]:
        args.command = args.command[1:]
    if not args.command:
        parser.error("a command is required after --")
    if args.minimum_free_gib <= 0 or args.poll_seconds <= 0:
        parser.error("disk floor and poll interval must be positive")
    return args


def main(argv: list[str] | None = None) -> int:
    args = parse_args(sys.argv[1:] if argv is None else argv)
    args.path.mkdir(parents=True, exist_ok=True)
    floor = args.minimum_free_gib * GIB
    available = free_bytes(args.path)
    if available < floor:
        print(
            f"ERROR: only {available / GIB:.1f} GiB free at {args.path}; "
            f"the stop floor is {args.minimum_free_gib} GiB",
            file=sys.stderr,
        )
        return STOPPED_FOR_DISK_EXIT

    process: subprocess.Popen[bytes] | None = None
    previous_handlers: dict[int, signal.Handlers] = {}

    def handle_signal(signum: int, _frame: object) -> None:
        raise GuardSignal(signum)

    try:
        for signum in HANDLED_SIGNALS:
            previous_handlers[signum] = signal.getsignal(signum)
            signal.signal(signum, handle_signal)
        process = subprocess.Popen(args.command, cwd=args.cwd, start_new_session=True)
        while True:
            try:
                return process.wait(timeout=args.poll_seconds)
            except subprocess.TimeoutExpired:
                available = free_bytes(args.path)
                if available < floor:
                    print(
                        f"ERROR: free space fell to {available / GIB:.1f} GiB at "
                        f"{args.path}; stopping the build before the "
                        f"{args.minimum_free_gib} GiB floor is crossed",
                        file=sys.stderr,
                        flush=True,
                    )
                    terminate_group(process)
                    return STOPPED_FOR_DISK_EXIT
    except GuardSignal as caught:
        return 128 + caught.signum
    finally:
        for signum in previous_handlers:
            signal.signal(signum, signal.SIG_IGN)
        if process is not None and process.poll() is None:
            terminate_group(process)
        for signum, previous in previous_handlers.items():
            signal.signal(signum, previous)


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