#!/usr/bin/env python3
"""Extract simulator-visible graph/KV metadata from Qwen235 server starts."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path


CONFIGS = ("tp4_ep1_mns64", "tp4_ep1_mns128", "tp8_ep8_mns64", "tp8_ep8_mns128")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--case-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    configs = {}
    for config in CONFIGS:
        log = args.case_root / "real" / config / "trial1" / "logs" / "server.log"
        text = log.read_text(errors="replace")
        token_matches = re.findall(r"GPU KV cache size:\s*([0-9,]+) tokens", text)
        capture_matches = re.findall(r"cudagraph_capture_sizes': \[([^]]+)\]", text)
        if not token_matches or not capture_matches:
            raise ValueError(f"missing runtime metadata in {log}")
        tokens = int(token_matches[-1].replace(",", ""))
        if tokens % 16:
            raise ValueError(f"KV token count is not block aligned: {tokens}")
        capture = [int(value.strip()) for value in capture_matches[-1].split(",")]
        configs[config] = {
            "num_gpu_blocks": tokens // 16,
            "capture_sizes": capture,
            "source_log": str(log.resolve()),
        }
    payload = {"schema": "qwen235-v020-runtime-contract-v1", "configs": configs}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload, sort_keys=True))


if __name__ == "__main__":
    main()