"""填圆任务的诊断运行 harness（宿主的固定工具，候选不得修改）。

一次诊断运行是一个容器、一次正式运行的名额，里面按顺序跑若干个诊断实验（docs/方法.md 4.3"诊断"一段）。
每个实验各开一个子进程：实验崩了、超时了，只影响它自己，后面的实验照常跑。
**它不判定解是否合法**：每个实验交出的答案，事后由独立校验器 ``validate.py`` 逐个重算。
读候选程序、写原始解、整理 metrics.json 都用 ``harness.py`` 里现成的函数，两份 harness 在这些事上不分叉。

命令行（由宿主的运行定义给出，实验员不用自己跑）：
  python -B harness-diagnosis.py --program <程序> --baseline <目录> --answers <目录>
         --seconds-each <秒> --report <文件> --experiments <实验编号> [<实验编号> ...]

每个实验：
  - 调用候选程序里的 ``run_experiment(experiment_id, baseline_dir)``，返回 (centers, radii, reported_sum)，
    与 ``run_packing()`` 的返回值同形。``experiment_id`` 是实验编号原文，``baseline_dir`` 是只读的原始产物目录
    （相对工作目录），里面是产出当前基线的那次正式运行的 ``solution.json``。
  - 分到的时限是 ``--seconds-each``（总时限 ÷ 实验数），到时被直接杀掉。
  - 随机种子同正式运行：环境变量 INFORSI_SEED，每个实验的子进程各自从头设一次。
  - 答案写进 ``<answers>/<实验编号>/solution.json``。指标写法同 TASK.md"指标怎么交给宿主"：
    在 ``run_experiment`` 返回之前往工作目录里的 ``metrics.json`` 写，写完后 harness 把它整理好、
    搬到 ``<answers>/<实验编号>/metrics.json``，下一个实验开始前工作目录里不留它。

每个实验结束后，harness 重写一次 ``--report``：每个实验一项
``{"id", "status", "wall_seconds"}``，status 是 finished（子进程退出码 0）、crashed（非零退出）、
timeout（超过分到的时限，被杀）。跑到哪写到哪：整个容器被总时限杀掉时，已经跑完的实验还留着记录，
没跑到的实验没有记录。
"""

import argparse
import hashlib
import importlib
import json
import os
import random
import signal
import subprocess
import sys
import time
from typing import Any, cast

import harness

SOLUTION_FILE = "solution.json"


def run_experiment_child(args) -> int:
    """子进程：设种子、载入候选程序、跑一个实验、落盘答案和指标。"""
    np = importlib.import_module("numpy")

    seed = int.from_bytes(
        hashlib.sha256(os.environ["INFORSI_SEED"].encode()).digest()[:4], "big"
    )
    random.seed(seed)
    np.random.seed(seed)

    module = harness.load_program(args.program)
    run_experiment = getattr(module, "run_experiment", None)
    if not callable(run_experiment):
        raise RuntimeError(f"{args.program} 没有定义可调用的 run_experiment(experiment_id, baseline_dir)")
    centers, radii, reported_sum = cast(Any, run_experiment(args.id, args.baseline))
    solution = {
        "centers": [[float(x), float(y)] for x, y in centers],
        "radii": [float(r) for r in radii],
        "reported_sum": float(reported_sum),
    }
    harness.write_json_atomically(args.output, solution)
    harness.finalize_metrics()
    return 0


def run_experiment_with_limit(args, experiment_id: str) -> tuple[str, float]:
    """在子进程里跑一个实验，交回结局和用时。"""
    answer_dir = os.path.join(args.answers, experiment_id)
    os.makedirs(answer_dir)
    # 上一个实验或候选目录里遗留的指标文件，不能被当成这个实验的。
    if os.path.exists(harness.METRICS_FILE):
        os.remove(harness.METRICS_FILE)
    command = [
        sys.executable, "-B", os.path.abspath(__file__), "--child",
        "--program", args.program, "--baseline", args.baseline,
        "--id", experiment_id, "--output", os.path.join(answer_dir, SOLUTION_FILE),
    ]
    start = time.monotonic()
    # 新开一个进程组：超时时连实验自己开的子进程一起杀掉。
    child = subprocess.Popen(command, start_new_session=True)
    try:
        status = "finished" if child.wait(timeout=args.seconds_each) == 0 else "crashed"
    except subprocess.TimeoutExpired:
        os.killpg(child.pid, signal.SIGKILL)
        child.wait()
        status = "timeout"
    wall = time.monotonic() - start
    if os.path.exists(harness.METRICS_FILE):
        if status == "finished":
            os.replace(harness.METRICS_FILE, os.path.join(answer_dir, harness.METRICS_FILE))
        else:
            os.remove(harness.METRICS_FILE)
    return status, wall


def main() -> int:
    parser = argparse.ArgumentParser(description="按顺序运行若干个诊断实验，各自限时、互不影响。")
    parser.add_argument("--program", required=True)
    parser.add_argument("--baseline", required=True)
    parser.add_argument("--child", action="store_true")
    parser.add_argument("--id")
    parser.add_argument("--output")
    parser.add_argument("--answers")
    parser.add_argument("--seconds-each", type=float)
    parser.add_argument("--report")
    parser.add_argument("--experiments", nargs="+")
    args = parser.parse_args()
    if args.child:
        return run_experiment_child(args)

    report: list[dict] = []
    for experiment_id in args.experiments:
        status, wall = run_experiment_with_limit(args, experiment_id)
        print(f"[诊断实验 {experiment_id}] {status}，用时 {wall:.2f} 秒", file=sys.stderr, flush=True)
        report.append({"id": experiment_id, "status": status, "wall_seconds": wall})
        harness.write_json_atomically(args.report, report)
    return 0


if __name__ == "__main__":
    sys.exit(main())