#!/usr/bin/env python3
"""Run bounded serial or MPI Catch2 profiling experiments.

Every run receives its own combined stdout/stderr log.  A summary CSV is
updated after each run so useful measurements survive a later timeout or
interrupt.  Timed-out process groups are terminated and, after a grace
period, force-killed; this is important for MPI jobs whose launcher may
otherwise leave workers behind.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
from pathlib import Path
import re
import signal
import subprocess
import sys
import time
from typing import TextIO


def positive_integer(value: str) -> int:
    parsed = int(value)
    if parsed <= 0:
        raise argparse.ArgumentTypeError("value must be positive")
    return parsed


def positive_float(value: str) -> float:
    parsed = float(value)
    if not parsed > 0.0:
        raise argparse.ArgumentTypeError("value must be positive")
    return parsed


def rank_list(value: str) -> list[int]:
    ranks = [positive_integer(item.strip()) for item in value.split(",") if item.strip()]
    if not ranks:
        raise argparse.ArgumentTypeError("at least one rank count is required")
    if len(set(ranks)) != len(ranks):
        raise argparse.ArgumentTypeError("rank counts must be unique")
    return ranks


def safe_name(value: str) -> str:
    name = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_.")
    return name[:80] or "all_tests"


def terminate_process_group(process: subprocess.Popen[str], grace_seconds: float) -> None:
    if process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    try:
        process.wait(timeout=grace_seconds)
        return
    except subprocess.TimeoutExpired:
        pass
    try:
        os.killpg(process.pid, signal.SIGKILL)
    except ProcessLookupError:
        return
    process.wait()


def run_bounded(
    command: list[str],
    working_directory: Path,
    log: TextIO,
    timeout_seconds: float,
    grace_seconds: float,
) -> tuple[str, int, float]:
    start = time.perf_counter()
    process = subprocess.Popen(
        command,
        cwd=working_directory,
        stdout=log,
        stderr=subprocess.STDOUT,
        text=True,
        start_new_session=True,
    )
    try:
        return_code = process.wait(timeout=timeout_seconds)
        status = "passed" if return_code == 0 else "failed"
    except subprocess.TimeoutExpired:
        status = "timeout"
        terminate_process_group(process, grace_seconds)
        return_code = process.returncode if process.returncode is not None else -signal.SIGKILL
    except BaseException:
        terminate_process_group(process, grace_seconds)
        raise
    return status, return_code, time.perf_counter() - start


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("catch_filter", help="Catch2 test name or tag expression")
    parser.add_argument(
        "--executable",
        type=Path,
        default=Path("cmake-build-profile-homebrew-llvm/tests"),
        help="Catch2 executable relative to --working-directory",
    )
    parser.add_argument(
        "--mpi-executable",
        default="mpirun",
        help="MPI launcher used for rank counts greater than one",
    )
    parser.add_argument(
        "--ranks",
        type=rank_list,
        default=[1],
        help="comma-separated rank counts, for example 1,2,4",
    )
    parser.add_argument("--repeat", type=positive_integer, default=1)
    parser.add_argument("--timeout", type=positive_float, default=300.0, help="seconds per run")
    parser.add_argument("--kill-grace", type=positive_float, default=5.0, help="seconds before SIGKILL")
    parser.add_argument("--working-directory", type=Path, default=Path.cwd())
    parser.add_argument("--output-directory", type=Path, default=Path("profile-results"))
    command_line = sys.argv[1:]
    if "--" in command_line:
        separator = command_line.index("--")
        runner_arguments = command_line[:separator]
        catch_arguments = command_line[separator + 1 :]
    else:
        runner_arguments = command_line
        catch_arguments = []

    arguments = parser.parse_args(runner_arguments)
    arguments.catch_arguments = catch_arguments
    return arguments


def main() -> int:
    arguments = parse_arguments()
    working_directory = arguments.working_directory.resolve()
    executable = arguments.executable
    if not executable.is_absolute():
        executable = working_directory / executable
    executable = executable.resolve()

    if not executable.is_file():
        raise FileNotFoundError(f"test executable does not exist: {executable}")
    if not os.access(executable, os.X_OK):
        raise PermissionError(f"test executable is not executable: {executable}")

    output_directory = arguments.output_directory
    if not output_directory.is_absolute():
        output_directory = working_directory / output_directory
    output_directory.mkdir(parents=True, exist_ok=True)

    experiment = safe_name(arguments.catch_filter)
    summary_path = output_directory / f"{experiment}.csv"
    field_names = [
        "ranks",
        "repeat",
        "status",
        "exit_code",
        "wall_seconds",
        "timeout_seconds",
        "log",
        "command",
    ]

    failed = False
    with summary_path.open("w", newline="", encoding="utf-8") as summary_file:
        writer = csv.DictWriter(summary_file, fieldnames=field_names)
        writer.writeheader()
        summary_file.flush()

        for ranks in arguments.ranks:
            for repetition in range(1, arguments.repeat + 1):
                test_command = [str(executable), arguments.catch_filter, *arguments.catch_arguments]
                command = (
                    test_command
                    if ranks == 1
                    else [arguments.mpi_executable, "-n", str(ranks), *test_command]
                )
                log_path = output_directory / f"{experiment}.r{ranks}.run{repetition}.log"
                print(
                    f"[{ranks} rank{'s' if ranks != 1 else ''}, run {repetition}/{arguments.repeat}] "
                    f"timeout={arguments.timeout:.1f}s log={log_path}",
                    flush=True,
                )
                with log_path.open("w", encoding="utf-8") as log:
                    log.write(f"command: {json.dumps(command)}\n")
                    log.flush()
                    status, return_code, wall_seconds = run_bounded(
                        command,
                        working_directory,
                        log,
                        arguments.timeout,
                        arguments.kill_grace,
                    )

                writer.writerow(
                    {
                        "ranks": ranks,
                        "repeat": repetition,
                        "status": status,
                        "exit_code": return_code,
                        "wall_seconds": f"{wall_seconds:.9f}",
                        "timeout_seconds": f"{arguments.timeout:.3f}",
                        "log": str(log_path),
                        "command": json.dumps(command),
                    }
                )
                summary_file.flush()
                print(f"  {status}: {wall_seconds:.3f}s (exit {return_code})", flush=True)
                failed = failed or status != "passed"

    print(f"summary: {summary_path}")
    return 1 if failed else 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (FileNotFoundError, PermissionError, ValueError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(2) from error