#!/usr/bin/env python3
"""Prepare old/new profile-only Frontier manifests from the frozen P1 probes."""

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("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--mode", choices=("old-profile-only", "new-profile-only"), required=True)
    parser.add_argument("--profile-root", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    source = json.loads(args.source.read_text())
    if source.get("status") != "PASS" or len(source.get("entries", [])) != 12:
        raise SystemExit("source manifest must contain 12 passing P1 probes")
    if args.mode == "new-profile-only" and args.profile_root is None:
        raise SystemExit("--profile-root is required for new-profile-only")

    output = args.output.resolve()
    config_root = output / "configs"
    cache_root = output / "prediction-cache"
    entries: list[dict[str, Any]] = []
    profile_hashes: dict[str, str] = {}
    if args.profile_root is not None:
        profile_root = args.profile_root.resolve()
        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)

    for entry in source["entries"]:
        config_path = Path(entry["config"])
        config = json.loads(config_path.read_text())
        config["mode"] = args.mode
        config["config_id"] = f"{config['cell_id']}__{args.mode}"
        config["calibration"]["a_tp"] = 1.0
        knobs = config["frontier"]["knobs"]
        knobs["cache_dir"] = str(cache_root)
        knobs["no_cache"] = False
        if args.mode == "new-profile-only":
            for key, filename in PROFILE_KEYS.items():
                knobs[key] = str((args.profile_root.resolve() / filename))

        target_config = config_root / f"{entry['fixture_id']}.json"
        write_json(target_config, config)
        updated_entry = dict(entry)
        updated_entry["config"] = str(target_config)
        updated_entry["calibration_scale"] = 1.0
        entries.append(updated_entry)

    prepared = {
        "schema": "frontier-qwen30-profile-comparison-prepared.v1",
        "status": "PASS",
        "mode": args.mode,
        "source": {
            "manifest": str(args.source.resolve()),
            "sha256": sha256(args.source),
        },
        "profile_hashes": profile_hashes,
        "isolation": {
            "calibration_a_tp": 1.0,
            "prediction_cache": str(cache_root),
            "all_non_profile_knobs_inherited": True,
        },
        "entries": entries,
    }
    write_json(output / "prepared-manifest.json", prepared)
    print(output / "prepared-manifest.json")


if __name__ == "__main__":
    main()