#!/usr/bin/env python3
"""Prepare the 92 frozen S2-R-b probes with a replacement profile bundle."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any


PROFILE_KEYS = {
    "linear_op_input_file": "linear_op.csv",
    "atten_input_file": "attention.csv",
    "moe_input_file": "moe.csv",
}


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


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--resolved-plan", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--profile-root", type=Path, required=True)
    parser.add_argument("--cache-root", type=Path)
    parser.add_argument("--shards", type=int, default=1)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.shards < 1:
        raise SystemExit("--shards must be positive")
    source = json.loads(args.resolved_plan.read_text())
    source_rows = [row for row in source["runs"] if row["mode"] == "uncalibrated"]
    if len(source_rows) != 92:
        raise SystemExit(f"expected 92 uncalibrated probes, found {len(source_rows)}")

    output = args.output.resolve()
    profile_root = args.profile_root.resolve()
    cache_root = (args.cache_root or (output / "prediction-cache")).resolve()
    profile_hashes: dict[str, str] = {}
    for filename in PROFILE_KEYS.values():
        path = profile_root / filename
        if not path.is_file():
            raise SystemExit(f"missing frozen profile: {path}")
        profile_hashes[str(path)] = sha256(path)

    entries: list[dict[str, Any]] = []
    for row in source_rows:
        fixture_dir = Path(row["fixture_dir"]).resolve()
        fixture_manifest_path = fixture_dir / "fixture_manifest.json"
        fixture = json.loads(fixture_manifest_path.read_text())
        config = json.loads(Path(row["config_path"]).read_text())
        config["mode"] = "new-profile-only"
        config["config_id"] = f"{row['cell_id']}__new-profile-only"
        config["calibration"]["a_tp"] = 1.0
        knobs = config["frontier"]["knobs"]
        knobs["cache_dir"] = str(cache_root)
        knobs["no_cache"] = False
        for key, filename in PROFILE_KEYS.items():
            knobs[key] = str(profile_root / filename)

        target_config = output / "configs" / f"{row['fixture_id']}.json"
        write_json(target_config, config)
        entries.append(
            {
                "fixture_id": row["fixture_id"],
                "cell": row["cell_id"],
                "role": f"probe-{int(row['probe_index']):02d}",
                "anchor": float(row["sampling_u"]),
                "selected_count": int(row["request_count"]),
                "probe_index": int(row["probe_index"]),
                "sampling_u": float(row["sampling_u"]),
                "tensor_parallel_size": int(row["tensor_parallel_size"]),
                "config": str(target_config),
                "fixture_manifest": str(fixture_manifest_path),
                "frontier_csv": str(fixture_dir / "frontier.csv"),
                "sidecar": str(fixture_dir / "sidecar.jsonl"),
                "calibration_scale": 1.0,
                "request_count": int(fixture["request_count"]),
            }
        )

    base = {
        "schema": "frontier-qwen30-s2-profile-comparison-prepared.v1",
        "status": "PASS",
        "mode": "new-profile-only",
        "source": {
            "resolved_plan": str(args.resolved_plan.resolve()),
            "sha256": sha256(args.resolved_plan),
        },
        "profile_hashes": profile_hashes,
        "isolation": {
            "calibration_a_tp": 1.0,
            "prediction_cache": str(cache_root),
            "all_non_profile_knobs_inherited": True,
        },
        "suite_total_runs": len(entries),
    }
    write_json(
        output / "prepared-manifest.json",
        {**base, "expected_runs": len(entries), "entries": entries},
    )
    for shard in range(args.shards):
        shard_entries = entries[shard :: args.shards]
        write_json(
            output / f"prepared-manifest-shard-{shard:02d}-of-{args.shards:02d}.json",
            {
                **base,
                "shard": {"index": shard, "count": args.shards},
                "expected_runs": len(shard_entries),
                "entries": shard_entries,
            },
        )
    print(output / "prepared-manifest.json")


if __name__ == "__main__":
    main()