#!/usr/bin/env python3
"""Run the opt-in HTTP/2 TLS soak in a pinned Linux Go image."""

import argparse
import math
from pathlib import Path
import shutil
import signal
import subprocess
import uuid


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", default="golang:1.25")
    parser.add_argument("--scenario", choices=("single", "mixed"), default="single", help="single TLS connection or eight connections with paced upload/download clients")
    parser.add_argument("--goproxy", default="", help="optional Go module proxy passed to the container")
    parser.add_argument("--duration", help="soak duration (default: 30m, or 60s with --race)")
    parser.add_argument("--timeout", default="35m", help="go test timeout (default: 35m)")
    parser.add_argument("--gomaxprocs", type=int, default=2)
    parser.add_argument("--cpus", type=float, default=2, help="container CPU quota (default: 2)")
    parser.add_argument("--memory-mb", type=int, default=2048, help="container memory cap, with no swap (default: 2048 MiB)")
    parser.add_argument("--pids-limit", type=int, default=256, help="container process/thread cap (default: 256)")
    parser.add_argument("--wall-timeout", type=int, default=2400, help="total time limit including build/download, in seconds (default: 2400)")
    parser.add_argument(
        "--race",
        action="store_true",
        help="enable the race detector for short concurrency checks; its memory overhead is not a production baseline",
    )
    args = parser.parse_args()
    if args.duration is None:
        args.duration = "60s" if args.race else "30m"
    if not args.duration or not args.timeout or args.gomaxprocs < 1:
        parser.error("--duration and --timeout must be non-empty and --gomaxprocs must be positive")
    if not math.isfinite(args.cpus) or args.cpus <= 0 or args.memory_mb < 128 or args.pids_limit < 16 or args.wall_timeout < 1:
        parser.error("--cpus and --wall-timeout must be positive; --memory-mb >= 128 and --pids-limit >= 16 are required")

    docker = shutil.which("docker")
    if docker is None:
        raise SystemExit("docker is required for HTTP/2 Linux soak")
    root = Path(__file__).resolve().parent.parent
    name = "go-omni-http2-soak-" + uuid.uuid4().hex[:12]
    command = [
        docker, "run", "--rm", "--pull=missing", "--name", name,
        "--cpus", str(args.cpus), "--memory", f"{args.memory_mb}m",
        "--memory-swap", f"{args.memory_mb}m", "--pids-limit", str(args.pids_limit),
        "--log-opt", "max-size=10m", "--log-opt", "max-file=1",
        "-v", f"{root}:/workspace:ro", "-w", "/workspace",
        "-e", "GOWORK=off", "-e", f"GOMAXPROCS={args.gomaxprocs}",
        "-e", f"GO_OMNI_HTTP2_SOAK_DURATION={args.duration}",
    ]
    if args.goproxy:
        command += ["-e", f"GOPROXY={args.goproxy}"]
    command += [args.image, "go", "test"]
    if args.race:
        command.append("-race")
    command += [
        "-v", "-timeout", args.timeout, "./transport/http2", "-run",
        "^TestTLSMixedSoak$" if args.scenario == "mixed" else "^TestTLSSoak$", "-count=1",
    ]
    print(f"HTTP/2 Linux soak image: {args.image}", flush=True)
    print(f"HTTP/2 TLS soak duration: {args.duration}", flush=True)
    print(f"HTTP/2 TLS soak mode: {'race' if args.race else 'regular'}", flush=True)
    print(f"HTTP/2 TLS soak scenario: {args.scenario}", flush=True)
    print(f"HTTP/2 soak container: {name}; limits: {args.cpus:g} CPUs, {args.memory_mb} MiB, {args.pids_limit} tasks", flush=True)

    def interrupt(signum, frame):
        raise KeyboardInterrupt

    previous_handlers = {
        sig: signal.signal(sig, interrupt) for sig in (signal.SIGTERM, signal.SIGHUP)
    }
    try:
        result = subprocess.run(command, cwd=root, timeout=args.wall_timeout)
    except subprocess.TimeoutExpired:
        raise SystemExit("HTTP/2 Linux soak exceeded --wall-timeout; stopping its container")
    except KeyboardInterrupt:
        raise SystemExit("HTTP/2 Linux soak interrupted; stopping its container")
    finally:
        for sig, handler in previous_handlers.items():
            signal.signal(sig, handler)
        # Only this run's unique container is targeted, including on timeout or
        # a disconnected SSH session. Do not prune other jobs.
        try:
            cleanup = subprocess.run(
                [docker, "rm", "--force", name], capture_output=True, text=True, timeout=15,
            )
            if cleanup.returncode != 0 and "No such container" not in cleanup.stderr:
                raise SystemExit(f"Could not clean up {name}: {cleanup.stderr.strip()}")
        except subprocess.TimeoutExpired:
            raise SystemExit(f"Cleanup timed out; inspect container {name} before another run")
    if result.returncode != 0:
        raise SystemExit(f"HTTP/2 Linux soak failed with exit code {result.returncode}")
    print("PASS HTTP/2 TLS soak", flush=True)


if __name__ == "__main__":
    main()