#!/usr/bin/env python3
"""Render the fleet queue for a frozen T0 real-surface plan."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def quoted(value: object) -> str:
    return json.dumps(str(value))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--plan", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--config", action="append", dest="configs")
    parser.add_argument("--rates", nargs="+", type=float)
    parser.add_argument("--artifact-root", default="artifacts/t0-real-surface-v1")
    parser.add_argument("--name-suffix", default="")
    args = parser.parse_args()
    plan = json.loads(args.plan.read_text())
    if plan.get("schema") != "qwen235b-t0-real-plan-v1" or len(plan.get("cells") or []) != 8:
        raise ValueError("invalid or incomplete T0 real plan")
    lattice = {0.10, 0.20, 0.40, 0.80, 1.20, 1.60, 2.40, 3.20}
    if args.rates and any(rate not in lattice for rate in args.rates):
        raise ValueError("rate override must stay on the frozen T0 lattice")

    indexed_cells = list(enumerate(plan["cells"]))
    if args.configs:
        requested = set(args.configs)
        known = {cell["config"]["name"] for _, cell in indexed_cells}
        if not requested <= known:
            raise ValueError(f"unknown configs: {sorted(requested - known)}")
        indexed_cells = [
            (index, cell)
            for index, cell in indexed_cells
            if cell["config"]["name"] in requested
        ]

    lines = [
        "# Generated from the frozen T0 real plan; do not edit rates in place.",
        "version = 1",
        "",
    ]
    for index, cell in indexed_cells:
        config = cell["config"]
        suffix = f"-{args.name_suffix}" if args.name_suffix else ""
        name = f"qwen235b-t0-real-{config['name']}{suffix}-20260716-v1"
        artifact = f"{args.artifact_root.rstrip('/')}/{config['name']}"
        rates = args.rates or cell["rates"]
        lines.extend(
            [
                "[[jobs]]",
                f"name = {quoted(name)}",
                f"gpus = {int(config['tp'])}",
                'gpu_model = "H20"',
                'hosts = ["dash0"]',
                'command = "timeout --signal=TERM --kill-after=60s 7200 bash run_t0_real_config.sh"',
                f"artifacts = [{quoted(artifact)}]",
                "",
                "[jobs.env]",
                f"OUTPUT_ROOT = {quoted('/home/admin/cpfs/wjh/aituner/aituner-fidelity-v1/' + artifact)}",
                f"TP = {quoted(config['tp'])}",
                f"MNS = {quoted(config['mns'])}",
                f"MBT = {quoted(config['mbt'])}",
                f"RATES = {quoted(' '.join(f'{rate:.2f}' for rate in rates))}",
                f"SERVER_PORT = {quoted(18920 + index)}",
                'VENV_ROOT = "/tmp/wjh-frontier-vllm0102-smoke/.venv"',
                'MODEL_ROOT = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8"',
                "",
            ]
        )
    args.output.write_text("\n".join(lines))
    print(args.output)


if __name__ == "__main__":
    main()