#!/usr/bin/env python3
"""Replay one TP2 trace with structured attention and measured prefill MoE."""

from __future__ import annotations

import argparse
import importlib.util
import json
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
S3_REAL = REPO / "runs/frontier-s3-real-v0"
BASE_REFERENCE = (
    REPO
    / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json"
)
BASE_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2"
EXPERIMENT_COMMIT = "1f8900a4ac64e45754b03d0aa7c1dddab65785cf"
ATTENTION_PATCH = (
    REPO
    / "runs/frontier-attn-structured-v0/"
    "0001-Experiment-with-structured-attention-prefill-predict.patch"
)
VERDICT = ROOT / "results/serving-smoke-verdict.json"


def load_s3_module():
    spec = importlib.util.spec_from_file_location(
        "s3_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
    )
    module = importlib.util.module_from_spec(spec)
    sys.path.insert(0, str(S3_REAL))
    spec.loader.exec_module(module)
    return module


def git(checkout: Path, *args: str) -> str:
    return subprocess.check_output(
        ["git", "-C", str(checkout), *args], text=True
    ).strip()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--trace", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--max-tokens", type=int, required=True)
    parser.add_argument("--duration-s", type=float)
    parser.add_argument("--cache-root", type=Path, required=True)
    parser.add_argument(
        "--frontier-checkout",
        type=Path,
        default=Path("/tmp/frontier-attn-structured-v0"),
    )
    parser.add_argument(
        "--attention-profile",
        type=Path,
        default=REPO
        / "runs/frontier-prefill-kvgrowth-fix-v0/profiles/"
        "profile-v5-kvgrowth/attention.csv",
    )
    args = parser.parse_args()
    args.config = "tp2_mns16"

    frontier = args.frontier_checkout.resolve()
    profile = args.attention_profile.resolve()
    if git(frontier, "rev-parse", "HEAD") != EXPERIMENT_COMMIT:
        raise SystemExit(f"unexpected experiment checkout HEAD: {frontier}")
    if git(frontier, "rev-parse", "HEAD^") != BASE_COMMIT:
        raise SystemExit("experiment commit is not directly based on frozen Frontier")
    if git(frontier, "status", "--porcelain"):
        raise SystemExit("experiment Frontier checkout must be clean")
    if not profile.is_file():
        raise SystemExit(f"attention profile missing: {profile}")

    smoke = json.loads(VERDICT.read_text())
    moe_scale = float(smoke["counterfactual"]["moe_scale"])
    reference = json.loads(BASE_REFERENCE.read_text())
    reference["frontier_checkout"] = str(frontier)
    reference["frontier_commit"] = EXPERIMENT_COMMIT
    generated_reference = ROOT / "frontier-reference.json"
    generated_reference.write_text(json.dumps(reference, indent=2))

    module = load_s3_module()
    module.REFERENCE = generated_reference
    module.EXPECTED_FRONTIER_COMMIT = EXPERIMENT_COMMIT
    original_replace = module.replace_flag
    injection_added = False

    def replace_and_override(argv: list[str], flag: str, value: str) -> None:
        nonlocal injection_added
        original_replace(argv, flag, value)
        if flag.endswith("trace_file"):
            attention_flag = (
                "--random_forrest_execution_time_predictor_config_atten_input_file"
            )
            original_replace(argv, attention_flag, str(profile))
            no_cache = (
                "--random_forrest_execution_time_predictor_config_no_cache"
            )
            if no_cache in argv:
                argv.remove(no_cache)
            if not injection_added:
                argv.extend(
                    (
                        "--random_forrest_execution_time_predictor_config_"
                        "moe_grouped_gemm_calibration_scale",
                        str(moe_scale),
                        "--random_forrest_execution_time_predictor_config_"
                        "decode_phase_moe_grouped_gemm_calibration_scale",
                        "1.0",
                    )
                )
                injection_added = True

    module.replace_flag = replace_and_override
    module.parse_args = lambda: args
    module.main()

    manifest_path = args.output_root / "manifest.json"
    manifest = json.loads(manifest_path.read_text())
    manifest.update(
        {
            "schema": "frontier-tp2-prefill-serving-replay-v1",
            "frontier_base_commit": BASE_COMMIT,
            "frontier_experiment_commit": EXPERIMENT_COMMIT,
            "structured_attention_patch": str(ATTENTION_PATCH.resolve()),
            "structured_attention_patch_sha256": module.sha256(ATTENTION_PATCH),
            "attention_profile_override": str(profile),
            "attention_profile_sha256": module.sha256(profile),
            "prefill_moe_grouped_gemm_scale": moe_scale,
            "decode_phase_moe_grouped_gemm_scale": 1.0,
            "moe_scale_source": str(VERDICT.resolve()),
            "moe_scale_source_sha256": module.sha256(VERDICT),
            "model_cache_enabled": True,
        }
    )
    manifest_path.write_text(json.dumps(manifest, indent=2))
    print(f"TP2 prefill-MoE replay done: {args.output_root}")


if __name__ == "__main__":
    main()