#!/usr/bin/env python3
"""Freeze the crossed-constraint action-aware development pilot."""

from __future__ import annotations

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


SCHEMA_V0 = "action-aware-constraint-pilot-manifest-v0"
SCHEMA_V1 = "action-aware-constraint-pilot-manifest-v1"


def configs(token_source_mbbt: int) -> tuple[dict[str, Any], ...]:
    return (
        {
            "id": "b_base",
            "mns": 64,
            "mbbt": token_source_mbbt,
            "repetition_order": [1, 2, 3],
        },
        {"id": "a_base", "mns": 16, "mbbt": 8192, "repetition_order": [2, 3, 1]},
        {"id": "shared", "mns": 64, "mbbt": 8192, "repetition_order": [3, 1, 2]},
        {
            "id": "b_mns",
            "mns": 128,
            "mbbt": token_source_mbbt,
            "repetition_order": [1, 3, 2],
        },
        {"id": "a_mbbt", "mns": 16, "mbbt": 16384, "repetition_order": [2, 1, 3]},
    )


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


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    os.replace(temporary, path)


def build(
    base_path: Path,
    *,
    token_source_mbbt: int = 256,
    prior_attempt_h20_hours: float = 0.0,
    prior_attempt_artifact: str | None = None,
) -> dict[str, Any]:
    if token_source_mbbt <= 0:
        raise ValueError("token source MBBT must be positive")
    if prior_attempt_h20_hours < 0.0 or prior_attempt_h20_hours >= 8.0:
        raise ValueError("prior attempt cost must be in [0, 8)")
    base = json.loads(base_path.read_text(encoding="utf-8"))
    if base.get("schema") != "intervention-response-phase-aware-pilot-manifest-v3":
        raise ValueError("unexpected base manifest schema")
    if base.get("status") != "PASS":
        raise ValueError("base manifest did not pass its preflight")
    if sorted(int(key) for key in base["repetitions"]) != [1, 2, 3]:
        raise ValueError("base manifest must contain exactly three repetitions")

    repetitions = {}
    selection_hashes = []
    for repetition in (1, 2, 3):
        source = base["repetitions"][str(repetition)]
        selection = dict(source["selections"]["mid"])
        selection_hashes.append(selection["request_id_order_sha256"])
        repetitions[str(repetition)] = {
            "study": source["study"],
            "study_sha256": source["study_sha256"],
            "selection": selection,
            "merged_trace": source["merged_trace"],
        }

    frozen_configs = configs(token_source_mbbt)
    config_ids = [str(config["id"]) for config in frozen_configs]
    schema = (
        SCHEMA_V0
        if token_source_mbbt == 256 and prior_attempt_h20_hours == 0.0
        else SCHEMA_V1
    )
    payload = {
        "schema": schema,
        "status": "PASS",
        "source": {
            "base_manifest": str(base_path.resolve()),
            "base_manifest_sha256": sha256_file(base_path),
            "window_id": base["source"]["window_id"],
            "source_trace": base["source"]["source_trace"],
            "source_trace_sha256": base["source"]["source_trace_sha256"],
        },
        "engine": {
            "tp": 4,
            "duration_s": 300.0,
            "disable_slo_early_stop": True,
            "client_timeout_s": 450.0,
            "burnin_max_elapsed_s": 90.0,
        },
        "burnin": base["burnin"],
        "repetitions": repetitions,
        "configs": [dict(config) for config in frozen_configs],
        "regimes": {
            "A": {
                "source": "a_base",
                "actions": {"mns": "shared", "mbbt": "a_mbbt"},
            },
            "B": {
                "source": "b_base",
                "actions": {"mns": "b_mns", "mbbt": "shared"},
            },
        },
        "budget": {
            "global_hard_cap_h20_hours": 8.0,
            "hard_cap_h20_hours": 8.0 - prior_attempt_h20_hours,
            "prior_attempt_h20_hours": prior_attempt_h20_hours,
            "prior_attempt_artifact": prior_attempt_artifact,
            "session_estimate_h20_hours": 1.35,
            "safety_h20_hours": 0.25,
            "expected_h20_hours": [6.0, 7.2],
            "expected_wall_minutes": [90, 110],
        },
        "gates": {
            "minimum_relative_winner_margin": 0.10,
            "minimum_exclusive_fraction": 0.10,
            "minimum_exclusive_ratio": 5.0,
            "phase_fractions": [0.25, 0.50, 0.75, 1.0],
            "material_kv_usage": 0.90,
        },
        "sanity": {
            "invariants": {
                "five_unique_configs": len(config_ids) == len(set(config_ids)) == 5,
                "three_disjoint_repetitions": len(set(selection_hashes)) == 3,
                "same_load_all_repetitions": len(
                    {
                        float(item["selection"]["offered_req_s_per_gpu"])
                        for item in repetitions.values()
                    }
                )
                == 1,
                "all_repetition_orders_are_permutations": all(
                    sorted(config["repetition_order"]) == [1, 2, 3]
                    for config in frozen_configs
                ),
            }
        },
    }
    payload["sanity"]["invariants"]["shared_endpoint_reused_by_both_regimes"] = (
        payload["regimes"]["A"]["actions"]["mns"]
        == payload["regimes"]["B"]["actions"]["mbbt"]
        == "shared"
    )
    payload["sanity"]["red_flags"] = [
        name
        for name, passed in payload["sanity"]["invariants"].items()
        if not passed
    ]
    if payload["sanity"]["red_flags"]:
        payload["status"] = "FAIL"
    return payload


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--token-source-mbbt", type=int, default=256)
    parser.add_argument("--prior-attempt-h20-hours", type=float, default=0.0)
    parser.add_argument("--prior-attempt-artifact")
    args = parser.parse_args()
    payload = build(
        args.base_manifest,
        token_source_mbbt=args.token_source_mbbt,
        prior_attempt_h20_hours=args.prior_attempt_h20_hours,
        prior_attempt_artifact=args.prior_attempt_artifact,
    )
    atomic_json(args.output, payload)
    print(json.dumps(payload["sanity"], sort_keys=True))
    if payload["status"] != "PASS":
        raise SystemExit("manifest preflight failed")


if __name__ == "__main__":
    main()