#!/usr/bin/env python3
"""Detached, resumable solo-H20 controller for the oracle-gap frontier."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shlex
import shutil
import signal
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any

from analyze import CONFIGS, PHASES, atomic_json, score_trial, summarize_trials


SCHEMA = 1
AMENDMENT = "A-OG-3"
AMENDMENT_REASON = (
    "apply the frozen per-logical-trial quarantine rule to an isolated "
    "local transport-invalid attempt"
)
REMOTE_ROOT = Path("/home/admin/cpfs/wjh/oracle-gap-20260713")
RUN_ROOT = REMOTE_ROOT / "runs"
STATE = RUN_ROOT / "controller-state.json"
PRIVATE = Path("/home/admin/cpfs/wjh/opprof-phase3-private/manifests")
MODEL = Path("/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B")
SOURCE = Path("/home/admin/cpfs/wjh/opprof-phase2-dash0-20260711/vllm-v0.24.0")
VENV = Path("/tmp/wjh-opprof-phase2-dash0-20260711/.venv")
REPO = Path(os.environ.get("AITUNER_ORACLE_REPO", Path(__file__).resolve().parents[2]))
P5_CLIENT = REPO / "runs/opprof-phase5/opprof_phase5_client.py"
P3_CLIENT_DIR = REPO / "runs/opprof-phase3/provenance"
GPU = 0
CPU_MASK = "0-19"
PORT = 8820
GPU_HOUR_LIMIT = 6.0
PRIMARY_ORDER = ("C11", "C00", "C01", "C10")
CONFIRM_ORDER = tuple(reversed(PRIMARY_ORDER))
CONFIG_DETAILS = {
    "C00": {"mns": 1024, "mbt": 8192, "flags": []},
    "C10": {"mns": 64, "mbt": 8192, "flags": ["--max-num-seqs", "64"]},
    "C01": {
        "mns": 1024,
        "mbt": 2048,
        "flags": ["--max-num-batched-tokens", "2048"],
    },
    "C11": {
        "mns": 64,
        "mbt": 2048,
        "flags": [
            "--max-num-seqs", "64", "--max-num-batched-tokens", "2048"
        ],
    },
}
BASE_RATES = {
    "P01": (32.0, 26.0, 36.0, 28.0, 34.0, 30.0),
    "P06": (1.7, 1.4, 2.0, 1.5, 1.9, 1.6, 1.8),
}
UP_EXTENSIONS = {
    "P01": (38.0, 40.0, 42.0),
    "P06": (2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.8, 3.0),
}
DOWN_EXTENSIONS = {"P01": (24.0, 22.0, 20.0), "P06": (1.3, 1.2, 1.1)}
TIMELINE = {
    "P01": {"warmup": 60, "clean": 60, "drain": 120},
    "P06": {"warmup": 60, "clean": 120, "drain": 240},
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def run_text(command: list[str], *, check: bool = True) -> str:
    result = subprocess.run(
        command, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )
    if check and result.returncode:
        raise RuntimeError(
            f"command failed ({result.returncode}): {shlex.join(command)}\n{result.stdout}"
        )
    return result.stdout


def load_state(resume: bool) -> dict[str, Any]:
    if STATE.exists():
        if not resume:
            raise RuntimeError(f"state exists; use --resume: {STATE}")
        return json.loads(STATE.read_text())
    return {
        "schema": SCHEMA,
        "status": "created",
        "created_at": time.time(),
        "gpu_hours": 0.0,
        "completed_trials": [],
        "stages": {},
        "fingerprint": {},
        "owned_pgids": [],
    }


def save_state(state: dict[str, Any]) -> None:
    state["updated_at"] = time.time()
    state["controller_pid"] = os.getpid()
    atomic_json(STATE, state)


def compute_apps() -> list[dict[str, Any]]:
    output = run_text(
        [
            "nvidia-smi",
            "--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
            "--format=csv,noheader,nounits",
        ],
        check=False,
    )
    rows = []
    for line in output.splitlines():
        parts = [part.strip() for part in line.split(",", 3)]
        if len(parts) == 4 and parts[1].isdigit():
            rows.append(
                {
                    "gpu_uuid": parts[0],
                    "pid": int(parts[1]),
                    "process_name": parts[2],
                    "used_memory_mib": int(parts[3].split()[0]),
                }
            )
    return rows


def gpu_snapshot() -> dict[str, Any]:
    query = run_text(
        [
            "nvidia-smi",
            "--query-gpu=index,name,uuid,memory.used,utilization.gpu,clocks.sm,clocks.mem,power.draw",
            "--format=csv,noheader,nounits",
        ]
    )
    return {
        "time": time.time(),
        "gpus": query.splitlines(),
        "compute_apps": compute_apps(),
        "loadavg": list(os.getloadavg()),
    }


def assert_idle() -> None:
    deadline = time.monotonic() + 30
    while time.monotonic() < deadline:
        if not compute_apps():
            return
        time.sleep(1)
    raise RuntimeError(f"GPU host is not idle: {compute_apps()}")


def descendants(root_pid: int) -> set[int]:
    output = run_text(["ps", "-e", "-o", "pid=,ppid="], check=False)
    children: dict[int, list[int]] = {}
    for line in output.splitlines():
        parts = line.split()
        if len(parts) == 2:
            pid, parent = map(int, parts)
            children.setdefault(parent, []).append(pid)
    result = {root_pid}
    pending = [root_pid]
    while pending:
        for child in children.get(pending.pop(), []):
            if child not in result:
                result.add(child)
                pending.append(child)
    return result


def assert_only_server_apps(server_pid: int) -> None:
    allowed = descendants(server_pid)
    unexpected = [row for row in compute_apps() if int(row["pid"]) not in allowed]
    if unexpected:
        raise RuntimeError(f"unexpected GPU process during run: {unexpected}")


def wait_idle_after_stop() -> None:
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline:
        output = run_text(
            [
                "nvidia-smi", "--query-gpu=index,memory.used",
                "--format=csv,noheader,nounits",
            ]
        )
        memory = [int(line.split(",")[1].strip()) for line in output.splitlines()]
        if not compute_apps() and all(value == 0 for value in memory):
            return
        time.sleep(1)
    raise RuntimeError("GPU memory/processes did not return to zero")


def fingerprint() -> dict[str, Any]:
    manifests = {}
    for phase in PHASES:
        path = PRIVATE / f"{phase}.jsonl"
        summary = json.loads(path.with_suffix(path.suffix + ".summary.json").read_text())
        if int(summary["rows"]) != 32768 or summary["sha256"] != sha256_file(path):
            raise RuntimeError(f"manifest mismatch: {phase}")
        manifests[phase] = {"path": str(path), "sha256": summary["sha256"]}
    return {
        "controller_sha256": sha256_file(Path(__file__).resolve()),
        "analyzer_sha256": sha256_file(Path(__file__).with_name("analyze.py")),
        "protocol_sha256": sha256_file(
            REPO / "docs/opprof/oracle-gap-protocol.md"
        ),
        "p5_client_sha256": sha256_file(P5_CLIENT),
        "p3_client_sha256": sha256_file(P3_CLIENT_DIR / "opprof_phase3_client.py"),
        "repo_commit": run_text(["git", "-C", str(REPO), "rev-parse", "HEAD"]).strip(),
        "repo_tree": run_text(["git", "-C", str(REPO), "rev-parse", "HEAD^{tree}"]).strip(),
        "vllm_commit": run_text(["git", "-C", str(SOURCE), "rev-parse", "HEAD"]).strip(),
        "model": str(MODEL),
        "manifests": manifests,
        "runtime": run_text(
            [str(VENV / "bin/python"), "-c", "import torch,vllm; print(torch.__version__,torch.version.cuda,vllm.__version__)"]
        ).strip(),
        "driver": run_text(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"]).splitlines()[0],
        "config_details": CONFIG_DETAILS,
        "base_rates": {key: list(value) for key, value in BASE_RATES.items()},
        "up_extensions": {
            key: list(value) for key, value in UP_EXTENSIONS.items()
        },
    }


def resume_compatible(old: dict[str, Any], current: dict[str, Any]) -> bool:
    immutable = (
        "analyzer_sha256",
        "p5_client_sha256",
        "p3_client_sha256",
        "vllm_commit",
        "model",
        "manifests",
        "runtime",
        "driver",
        "config_details",
        "base_rates",
    )
    if not all(old.get(key) == current.get(key) for key in immutable):
        return False
    old_extensions = old.get("up_extensions")
    current_extensions = current.get("up_extensions")
    if old_extensions == current_extensions:
        return True
    # The only earlier compatible transition was A-OG-1, which appended the
    # five registered P06 anchors after the original 2.3-rps list.
    if not isinstance(old_extensions, dict) or not isinstance(current_extensions, dict):
        return False
    return (
        old_extensions.get("P01") == current_extensions.get("P01")
        and list(old_extensions.get("P06", [])) + [2.4, 2.5, 2.6, 2.8, 3.0]
        == current_extensions.get("P06")
    )


def ensure_provenance(current: dict[str, Any]) -> None:
    destination = RUN_ROOT / "provenance"
    destination.mkdir(parents=True, exist_ok=True)
    for source in (
        Path(__file__).resolve(),
        Path(__file__).with_name("analyze.py"),
        P5_CLIENT,
        P3_CLIENT_DIR / "opprof_phase3_client.py",
        REPO / "docs/opprof/oracle-gap-protocol.md",
    ):
        target = destination / source.name
        if target.exists() and sha256_file(target) != sha256_file(source):
            digest = sha256_file(source)
            target = destination / f"{source.stem}.{digest[:12]}{source.suffix}"
        if target.exists() and sha256_file(target) != sha256_file(source):
            raise RuntimeError(f"content-addressed provenance mismatch: {target}")
        if not target.exists():
            shutil.copy2(source, target)
    fingerprint_path = (
        destination / f"fingerprint.{current['repo_commit'][:8]}.json"
    )
    atomic_json(fingerprint_path, current)
    if not (destination / "fingerprint.json").exists():
        atomic_json(destination / "fingerprint.json", current)
    atomic_json(destination / "host-before.json", gpu_snapshot())
    (destination / "nvidia-smi-q.txt").write_text(run_text(["nvidia-smi", "-q"]))


def rate_label(rate: float) -> str:
    return f"{rate:.3f}".rstrip("0").rstrip(".")


def trial_key(phase: str, config: str, rate: float, repetition: int) -> str:
    return f"{phase}-{config}-r{rate_label(rate)}-rep{repetition}"


def derived_manifest(phase: str, rate: float, repetition: int) -> Path:
    label = f"{phase}-r{rate_label(rate)}-rep{repetition}"
    output = RUN_ROOT / "manifests" / f"{label}.jsonl"
    summary_path = output.with_suffix(output.suffix + ".summary.json")
    source = PRIVATE / f"{phase}.jsonl"
    domain = int.from_bytes(
        hashlib.sha256(f"oracle-gap:{label}".encode()).digest()[:4], "big"
    )
    if output.exists() and summary_path.exists():
        summary = json.loads(summary_path.read_text())
        if summary["sha256"] == sha256_file(output) and summary["source_sha256"] == sha256_file(source):
            return output
        raise RuntimeError(f"derived manifest mismatch: {output}")
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(f"{output.name}.tmp.{os.getpid()}")
    rows = 0
    input_sum = 0
    output_sum = 0
    with source.open() as src, temporary.open("w") as dst:
        for line in src:
            row = json.loads(line)
            row["token_seed"] = (int(row.get("token_seed", 0)) + domain * 1000003) & ((1 << 63) - 1)
            row["token_domain"] = domain
            dst.write(json.dumps(row, sort_keys=True, separators=(",", ":")) + "\n")
            rows += 1
            input_sum += int(row["input_tokens"])
            output_sum += int(row["output_tokens"])
        dst.flush()
        os.fsync(dst.fileno())
    os.replace(temporary, output)
    summary = {
        "schema": 1,
        "phase": phase,
        "rate": rate,
        "repetition": repetition,
        "rows": rows,
        "input_token_sum": input_sum,
        "output_token_sum": output_sum,
        "token_domain": domain,
        "source_sha256": sha256_file(source),
        "sha256": sha256_file(output),
        "invariants": {"rows_32768": rows == 32768, "positive_work": input_sum > 0 and output_sum > 0},
    }
    if not all(summary["invariants"].values()):
        raise RuntimeError(f"derived manifest invalid: {summary}")
    atomic_json(summary_path, summary)
    return output


def server_command(config: str) -> list[str]:
    return [
        "taskset", "-c", CPU_MASK, str(VENV / "bin/vllm"), "serve", str(MODEL),
        "--host", "127.0.0.1", "--port", str(PORT),
        "--tensor-parallel-size", "1", "--enable-chunked-prefill",
        "--enable-prefix-caching", "--shutdown-timeout", "120",
        *CONFIG_DETAILS[config]["flags"],
    ]


def wait_ready(process: subprocess.Popen[Any], timeout: float = 300) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if process.poll() is not None:
            raise RuntimeError(f"server exited before ready: {process.returncode}")
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=1) as response:
                if response.status == 200:
                    return
        except Exception:
            pass
        time.sleep(1)
    raise TimeoutError("server readiness timeout")


def validate_startup(log_path: Path, config: str) -> None:
    log = log_path.read_text(errors="replace")
    details = CONFIG_DETAILS[config]
    invariants = {
        "triton_moe": "Using TRITON Unquantized MoE backend" in log,
        "tp1": "tensor_parallel_size=1" in log,
        "mbt": (
            "Chunked prefill is enabled with max_num_batched_tokens=8192" in log
            if details["mbt"] == 8192
            else "'max_num_batched_tokens': 2048" in log
        ),
        "mns": details["mns"] == 1024 or "'max_num_seqs': 64" in log,
    }
    if not all(invariants.values()):
        raise RuntimeError(f"server startup invariants failed {config}: {invariants}")


def start_server(config: str, stage: str, state: dict[str, Any]) -> dict[str, Any]:
    assert_idle()
    directory = RUN_ROOT / "servers" / stage
    directory.mkdir(parents=True, exist_ok=True)
    command = server_command(config)
    echo = (
        f"RUN_ECHO stage={stage} host=dash0 gpu=0 cpus={CPU_MASK} config={config} "
        f"model={MODEL} source={SOURCE} manifests={PRIVATE}/P01,P06.jsonl "
        f"output={RUN_ROOT} expected_server_plus_trials=20-45min budget_cap={GPU_HOUR_LIMIT}H20h"
    )
    with (RUN_ROOT / "launch-echo.log").open("a") as handle:
        handle.write(echo + "\n")
    print(echo, flush=True)
    (directory / "command.txt").write_text(shlex.join(command) + "\n")
    atomic_json(directory / "gpu-before.json", gpu_snapshot())
    handle = (directory / "server.log").open("ab", buffering=0)
    environment = os.environ.copy()
    environment.update(
        {
            "CUDA_VISIBLE_DEVICES": str(GPU),
            "VLLM_OPPROF_DIR": str(directory / "opprof"),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
            "AITUNER_ORACLE_GAP_MARKER": stage,
        }
    )
    process = subprocess.Popen(
        command, cwd=SOURCE, env=environment, stdout=handle,
        stderr=subprocess.STDOUT, start_new_session=True,
    )
    state["owned_pgids"] = [process.pid]
    save_state(state)
    wait_ready(process)
    validate_startup(directory / "server.log", config)
    return {
        "config": config, "stage": stage, "dir": directory, "process": process,
        "handle": handle, "started_at": time.time(),
    }


def stop_server(entry: dict[str, Any], state: dict[str, Any]) -> None:
    process = entry["process"]
    if process.poll() is None:
        try:
            os.kill(process.pid, signal.SIGINT)
        except ProcessLookupError:
            pass
        try:
            process.wait(timeout=150)
        except subprocess.TimeoutExpired:
            for signum, timeout in ((signal.SIGTERM, 10), (signal.SIGKILL, 30)):
                if process.poll() is not None:
                    break
                try:
                    os.killpg(process.pid, signum)
                except ProcessLookupError:
                    pass
                try:
                    process.wait(timeout=timeout)
                except subprocess.TimeoutExpired:
                    continue
            if process.poll() is None:
                raise TimeoutError(f"server process group did not stop: {entry['stage']}")
    entry["handle"].close()
    elapsed = time.time() - float(entry["started_at"])
    state["gpu_hours"] = float(state["gpu_hours"]) + elapsed / 3600.0
    state["owned_pgids"] = []
    atomic_json(entry["dir"] / "gpu-after.json", gpu_snapshot())
    log = (entry["dir"] / "server.log").read_text(errors="replace")
    if "mode=drain timeout=120s" not in log:
        raise RuntimeError(f"server did not use drain shutdown: {entry['stage']}")
    wait_idle_after_stop()
    save_state(state)


def client_command(
    phase: str, rate: float, repetition: int, output: Path
) -> list[str]:
    timeline = TIMELINE[phase]
    return [
        "taskset", "-c", CPU_MASK, str(VENV / "bin/python"), str(P5_CLIENT), "run",
        "--manifest", str(derived_manifest(phase, rate, repetition)),
        "--base-url", f"http://127.0.0.1:{PORT}", "--model", str(MODEL),
        "--load-point", "moderate", "--fixed-request-rate", str(rate),
        "--max-concurrency", "256", "--ignore-eos", "--temperature", "0",
        "--warmup-seconds", str(timeline["warmup"]),
        "--clean-segment-seconds", str(timeline["clean"]),
        "--num-clean-segments", "1", "--post-clean-seconds", "0",
        "--drain-timeout-seconds", str(timeline["drain"]),
        "--workload-seed", "20260712", "--server-seed", "20260712",
        "--result-dir", str(output / "client"),
    ]


def run_trial(
    state: dict[str, Any], server: dict[str, Any], phase: str, rate: float,
    repetition: int, role: str,
) -> dict[str, Any]:
    config = str(server["config"])
    key = trial_key(phase, config, rate, repetition)
    output = RUN_ROOT / "trials" / f"{phase}-{config}" / f"rate-{rate_label(rate)}" / f"rep-{repetition}"
    score_path = output / "score.json"
    if key in state["completed_trials"]:
        if not score_path.exists():
            raise RuntimeError(f"completed state lacks score: {key}")
        return json.loads(score_path.read_text())
    output.mkdir(parents=True, exist_ok=True)
    command = client_command(phase, rate, repetition, output)
    (output / "command.txt").write_text(shlex.join(command) + "\n")
    if server["process"].poll() is not None:
        raise RuntimeError(f"server exited before trial: {key}")
    assert_only_server_apps(int(server["process"].pid))
    environment = os.environ.copy()
    environment["PYTHONPATH"] = str(P3_CLIENT_DIR) + os.pathsep + environment.get("PYTHONPATH", "")
    started = time.time()
    with (output / "client.log").open("ab", buffering=0) as handle:
        result = subprocess.run(
            command, cwd=REPO, env=environment, stdout=handle,
            stderr=subprocess.STDOUT, timeout=900,
        )
    assert_only_server_apps(int(server["process"].pid))
    if server["process"].poll() is not None:
        raise RuntimeError(f"server exited during trial: {key}")
    client_result = output / "client/result.json"
    client_sanity = output / "client/sanity.json"
    if not client_result.exists() or not client_sanity.exists():
        raise RuntimeError(f"client produced no result: {key} rc={result.returncode}")
    sanity = json.loads(client_sanity.read_text())["invariants"]
    failed = [name for name, passed in sanity.items() if not passed]
    allowed = {"moderate_offered_within_5pct"}
    if result.returncode != 0 and not failed:
        raise RuntimeError(f"client failed without sanity marker: {key}")
    if set(failed) - allowed:
        raise RuntimeError(f"client validity failure {key}: {failed}")
    score = score_trial(
        output / "client/requests.jsonl", client_result, phase=phase,
        config=config, target_rate=rate, repetition=repetition, role=role,
    )
    score["wall_seconds"] = time.time() - started
    score["client_returncode"] = result.returncode
    score["client_failed_invariants"] = failed
    score["manifest_sha256"] = sha256_file(derived_manifest(phase, rate, repetition))
    atomic_json(score_path, score)
    state["completed_trials"].append(key)
    state["last_trial"] = {
        "key": key, "feasible": score["feasible"], "pass_rate": score["pass_rate"],
        "goodput": score["slo_goodput_rps"], "completed_at": time.time(),
    }
    save_state(state)
    print(
        f"TRIAL {key} feasible={score['feasible']} pass={score['pass_rate']:.6f} "
        f"goodput={score['slo_goodput_rps']:.6f} lagmax={score['schedule_lag_ms']['max']:.3f}",
        flush=True,
    )
    return score


def primary_bracket(rows: list[dict[str, Any]]) -> tuple[float, float] | None:
    verdict = {float(row["target_rate_rps"]): bool(row["feasible"]) for row in rows}
    ordered = sorted(verdict)
    passes = [rate for rate in ordered if verdict[rate]]
    failures = [rate for rate in ordered if not verdict[rate]]
    if not passes or not failures:
        return None
    lower = max(passes)
    upper = min((rate for rate in failures if rate > lower), default=None)
    if upper is None or any(verdict[rate] for rate in ordered if rate > upper):
        raise RuntimeError(f"non-monotone primary frontier: {verdict}")
    return lower, upper


def run_primary_config(state: dict[str, Any], config: str) -> None:
    stage = f"primary-{config}"
    if state["stages"].get(stage, {}).get("status") == "complete":
        return
    state["stages"][stage] = {"status": "starting", "started_at": time.time()}
    save_state(state)
    server = start_server(config, stage, state)
    failure = None
    try:
        for phase in PHASES:
            phase_rows = [run_trial(state, server, phase, rate, 0, "primary") for rate in BASE_RATES[phase]]
            bracket = primary_bracket(phase_rows)
            if bracket is None:
                extension = UP_EXTENSIONS[phase] if all(row["feasible"] for row in phase_rows) else DOWN_EXTENSIONS[phase]
                for rate in extension:
                    phase_rows.append(run_trial(state, server, phase, rate, 0, "primary-extension"))
                    bracket = primary_bracket(phase_rows)
                    if bracket is not None:
                        break
            if bracket is None:
                raise RuntimeError(f"failed to bracket {phase}-{config}")
            state["stages"][stage].setdefault("brackets", {})[phase] = list(bracket)
            save_state(state)
    except Exception as error:
        failure = error
    finally:
        try:
            stop_server(server, state)
        except Exception as error:
            failure = failure or error
    if failure is not None:
        state["stages"][stage]["status"] = "failed"
        state["stages"][stage]["failure"] = repr(failure)
        state["status"] = "failed"
        save_state(state)
        raise failure
    state["stages"][stage]["status"] = "complete"
    state["stages"][stage]["completed_at"] = time.time()
    save_state(state)


def run_confirm_config(state: dict[str, Any], config: str) -> None:
    stage = f"confirm-{config}"
    if state["stages"].get(stage, {}).get("status") == "complete":
        return
    brackets = state["stages"][f"primary-{config}"]["brackets"]
    state["stages"][stage] = {"status": "starting", "started_at": time.time(), "brackets": brackets}
    save_state(state)
    server = start_server(config, stage, state)
    failure = None
    try:
        for phase in PHASES:
            lower, upper = (float(value) for value in brackets[phase])
            for rate in (upper, lower):
                for repetition in (1, 2):
                    run_trial(state, server, phase, rate, repetition, "boundary-confirmation")
    except Exception as error:
        failure = error
    finally:
        try:
            stop_server(server, state)
        except Exception as error:
            failure = failure or error
    if failure is not None:
        state["stages"][stage]["status"] = "failed"
        state["stages"][stage]["failure"] = repr(failure)
        state["status"] = "failed"
        save_state(state)
        raise failure
    state["stages"][stage]["status"] = "complete"
    state["stages"][stage]["completed_at"] = time.time()
    save_state(state)


def cleanup_recorded(state: dict[str, Any]) -> None:
    for pgid in state.get("owned_pgids", []):
        try:
            os.killpg(int(pgid), signal.SIGKILL)
        except ProcessLookupError:
            pass
    state["owned_pgids"] = []
    save_state(state)
    wait_idle_after_stop()


def execute(resume: bool) -> None:
    RUN_ROOT.mkdir(parents=True, exist_ok=True)
    state = load_state(resume)
    if resume and state.get("owned_pgids"):
        cleanup_recorded(state)
    assert_idle()
    current = fingerprint()
    if state["fingerprint"] and state["fingerprint"] != current:
        if not (
            resume
            and state.get("status") == "failed"
            and resume_compatible(state["fingerprint"], current)
        ):
            raise RuntimeError("resume fingerprint changed incompatibly")
        state.setdefault("amendments", []).append(
            {
                "id": AMENDMENT,
                "accepted_at": time.time(),
                "reason": AMENDMENT_REASON,
                "completed_trials_before": len(
                    state.get("completed_trials", [])
                ),
                "gpu_hours_before": state.get("gpu_hours"),
                "old_fingerprint": state["fingerprint"],
                "new_fingerprint": current,
            }
        )
    state["fingerprint"] = current
    state["status"] = "running"
    save_state(state)
    ensure_provenance(current)
    for config in PRIMARY_ORDER:
        if float(state["gpu_hours"]) >= GPU_HOUR_LIMIT:
            raise RuntimeError("GPU-hour hard cap reached before primary completion")
        run_primary_config(state, config)
    for config in CONFIRM_ORDER:
        if float(state["gpu_hours"]) >= GPU_HOUR_LIMIT:
            raise RuntimeError("GPU-hour hard cap reached before confirmations")
        run_confirm_config(state, config)
    score_paths = sorted((RUN_ROOT / "trials").glob("**/score.json"))
    summary = summarize_trials([json.loads(path.read_text()) for path in score_paths])
    summary["gpu_hours"] = state["gpu_hours"]
    summary["trial_files"] = len(score_paths)
    atomic_json(RUN_ROOT / "metrics.json", summary)
    state["status"] = "complete"
    state["completed_at"] = time.time()
    state["verdict"] = summary["verdict"]
    save_state(state)
    print(json.dumps({"status": "complete", "verdict": summary["verdict"], "gpu_hours": state["gpu_hours"]}, sort_keys=True))


def plan() -> dict[str, Any]:
    primary = sum(len(BASE_RATES[phase]) for phase in PHASES) * len(CONFIGS)
    confirmations = 2 * 2 * len(PHASES) * len(CONFIGS)
    return {
        "schema": 1,
        "placement": "serialized solo GPU0",
        "primary_order": list(PRIMARY_ORDER),
        "confirm_order": list(CONFIRM_ORDER),
        "primary_trials_without_extensions": primary,
        "confirmation_trials": confirmations,
        "expected_total_trials": primary + confirmations,
        "expected_h20_hours": "3.0-4.0",
        "hard_cap_h20_hours": GPU_HOUR_LIMIT,
        "rates": {key: list(value) for key, value in BASE_RATES.items()},
        "timelines": TIMELINE,
        "output": str(RUN_ROOT),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)
    run = sub.add_parser("run")
    run.add_argument("--resume", action="store_true")
    sub.add_parser("plan")
    sub.add_parser("status")
    args = parser.parse_args()
    if args.command == "run":
        execute(args.resume)
    elif args.command == "plan":
        print(json.dumps(plan(), indent=2, sort_keys=True))
    else:
        print(STATE.read_text() if STATE.exists() else json.dumps({"status": "not-started"}))


if __name__ == "__main__":
    main()