#!/usr/bin/env python3
"""Replay frozen Frontier fixtures with state-detail outputs enabled.

This runner does not modify the frozen fixture/config inputs.  It reuses the
audited SimFid command and only turns on existing Frontier batch/ledger output
flags in a separate result root.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


HERE = Path(__file__).resolve().parent
AITUNER_ROOT = HERE.parents[1]
sys.path.insert(0, str(HERE))

from common_state import numeric, summarize_frontier  # noqa: E402, F401


def load_module(name: str, path: Path):
    module_root = str(path.parent.resolve())
    if module_root not in sys.path:
        sys.path.insert(0, module_root)
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise ImportError(path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


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 git_capture(root: Path, *arguments: str) -> str:
    return subprocess.run(
        ["git", "-C", str(root), *arguments],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    ).stdout


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")
    os.replace(temporary, path)


def enable_state_outputs(command: list[str]) -> list[str]:
    result = list(command)
    disabled_ledger = "--no-metrics_config_store_frontier_stage_batch_ledger"
    enabled_ledger = "--metrics_config_store_frontier_stage_batch_ledger"
    enabled_batches = "--metrics_config_keep_individual_batch_metrics"
    disabled_batches = "--no-metrics_config_keep_individual_batch_metrics"
    if result.count(disabled_ledger) != 1 or enabled_ledger in result:
        raise ValueError("base command must explicitly disable one full Frontier ledger")
    result[result.index(disabled_ledger)] = enabled_ledger
    if disabled_batches in result:
        if result.count(disabled_batches) != 1:
            raise ValueError("duplicate disabled individual-batch flag")
        result[result.index(disabled_batches)] = enabled_batches
    elif enabled_batches not in result:
        result.append(enabled_batches)
    if result.count(enabled_ledger) != 1 or result.count(enabled_batches) != 1:
        raise ValueError("state-detail flags were not enabled exactly once")
    return result


def find_state_metrics(run_root: Path) -> dict[str, Path]:
    patterns = {
        "system": "frontier_metrics/**/system_metrics.json",
        "requests": "frontier_metrics/**/request_metrics.csv",
        "batches": "frontier_metrics/**/monolithic_batch_metrics.csv",
        "ledger": "frontier_metrics/**/frontier_stage_batch_ledger.jsonl",
    }
    result = {}
    for name, pattern in patterns.items():
        matches = sorted(run_root.glob(pattern))
        if len(matches) != 1:
            raise ValueError(
                f"{run_root}: expected exactly one {name} artifact, found {len(matches)}"
            )
        result[name] = matches[0]
    if len({path.parent for path in result.values()}) != 1:
        raise ValueError("Frontier state artifacts are not colocated")
    return result


def execute(args: argparse.Namespace) -> dict[str, Any]:
    prepared = json.loads(args.prepared_manifest.read_text(encoding="utf-8"))
    if prepared.get("status") != "PASS":
        raise RuntimeError("prepared simulator manifest did not pass")
    matches = [
        entry
        for entry in prepared["entries"]
        if entry["cell"] == args.cell and entry["role"] == args.role
    ]
    if len(matches) != 1:
        raise ValueError(
            f"expected one prepared entry for {args.cell}/{args.role}, found {len(matches)}"
        )
    entry = matches[0]
    if args.output.exists() and any(args.output.iterdir()):
        raise FileExistsError(f"non-empty output already exists: {args.output}")
    args.output.mkdir(parents=True, exist_ok=True)

    driver = load_module(
        "telemetry_residual_execution_driver",
        args.replayserve_root / "runs/simfid_s2rb/results/execution_driver.py",
    )
    config_path = Path(entry["config"])
    config = json.loads(config_path.read_text(encoding="utf-8"))
    fixture_manifest_path = Path(entry["fixture_manifest"])
    fixture = json.loads(fixture_manifest_path.read_text(encoding="utf-8"))
    trace_path = Path(entry["frontier_csv"])
    sidecar_path = Path(entry["sidecar"])
    metrics_root = args.output / "frontier_metrics"
    run_id = f"telemetry_residual_r0_{args.cell}_{args.role}"
    knobs = config["frontier"]["knobs"]
    base_command = driver.build_command(
        trace_path=trace_path,
        metrics_root=metrics_root,
        run_id=run_id,
        knobs=knobs,
    )
    driver.audit_command(base_command, knobs)
    command = enable_state_outputs(base_command)
    row = {
        "hook_path": config["calibration"]["hook_path"],
        "applied_a_tp": config["calibration"]["a_tp"],
        "sidecar_path": str(sidecar_path),
        "request_count": int(fixture["request_count"]),
        "tensor_parallel_size": int(fixture["tensor_parallel_size"]),
    }
    environment = driver.environment_for(row)
    if environment.get("CUDA_VISIBLE_DEVICES") != "":
        raise ValueError("Frontier state replay must hide CUDA devices")
    manifest = {
        "schema": "telemetry-residual-frontier-state-run-v1",
        "entry": {
            "cell": args.cell,
            "role": args.role,
            "anchor": entry["anchor"],
            "request_count": entry["selected_count"],
        },
        "inputs": {
            "prepared_manifest": str(args.prepared_manifest.resolve()),
            "prepared_manifest_sha256": sha256_file(args.prepared_manifest),
            "config": str(config_path.resolve()),
            "config_sha256": sha256_file(config_path),
            "fixture_manifest": str(fixture_manifest_path.resolve()),
            "fixture_manifest_sha256": sha256_file(fixture_manifest_path),
            "frontier_csv": str(trace_path.resolve()),
            "frontier_csv_sha256": sha256_file(trace_path),
            "sidecar": str(sidecar_path.resolve()),
            "sidecar_sha256": sha256_file(sidecar_path),
        },
        "frontier": {
            "root": str(args.frontier_root.resolve()),
            "git_head": git_capture(args.frontier_root, "rev-parse", "HEAD").strip(),
            "git_status_short": git_capture(args.frontier_root, "status", "--short"),
        },
        "runner": {
            "script": str(Path(__file__).resolve()),
            "script_sha256": sha256_file(Path(__file__).resolve()),
            "aituner_git_head": git_capture(AITUNER_ROOT, "rev-parse", "HEAD").strip(),
            "aituner_git_status_short": git_capture(AITUNER_ROOT, "status", "--short"),
        },
        "environment": {
            key: environment[key]
            for key in (
                "PYTHONPATH",
                "FRONTIER_EXECUTION_TIME_SCALE",
                "CUDA_VISIBLE_DEVICES",
                "NVIDIA_VISIBLE_DEVICES",
                "FRONTIER_LOG_LEVEL",
            )
        },
        "command": command,
        "state_outputs": {
            "individual_batch_metrics": True,
            "full_stage_batch_ledger": True,
        },
        "contains_prompt_text": False,
    }
    atomic_json(args.output / "run_manifest.json", manifest)
    start = time.monotonic()
    with (args.output / "stdout.log").open("w", encoding="utf-8") as stdout, (
        args.output / "stderr.log"
    ).open("w", encoding="utf-8") as stderr:
        try:
            process = subprocess.run(
                command,
                cwd=args.frontier_root,
                env=environment,
                stdout=stdout,
                stderr=stderr,
                timeout=args.timeout_s,
            )
            return_code = int(process.returncode)
        except subprocess.TimeoutExpired:
            return_code = 124
    runtime_s = time.monotonic() - start
    if return_code != 0:
        failure = {
            "status": "STOP",
            "return_code": return_code,
            "runtime_s": runtime_s,
        }
        atomic_json(args.output / "failure.json", failure)
        raise RuntimeError(f"Frontier state replay failed: {failure}")

    paths = find_state_metrics(args.output)
    summary = summarize_frontier(
        system_metrics_path=paths["system"],
        request_metrics_path=paths["requests"],
        batch_metrics_path=paths["batches"],
        ledger_path=paths["ledger"],
    )
    atomic_json(args.output / "common-state.json", summary)
    scorer = driver.score_trial(row, paths["system"], paths["requests"])
    atomic_json(args.output / "scorer_output.json", scorer)
    sizes = {name: path.stat().st_size for name, path in paths.items()}
    red_flags = []
    if summary["interval"]["request_count"] != int(entry["selected_count"]):
        red_flags.append("request_count_mismatch")
    if summary["sanity"]["batch_rows"] <= 0:
        red_flags.append("no_batch_rows")
    if summary["sanity"]["ledger_rows"] <= 0:
        red_flags.append("no_ledger_rows")
    result = {
        "schema": "telemetry-residual-frontier-state-result-v1",
        "status": "PASS" if not red_flags else "STOP",
        "runtime_s": runtime_s,
        "return_code": return_code,
        "paths": {name: str(path.resolve()) for name, path in paths.items()},
        "bytes": sizes,
        "common_state": str((args.output / "common-state.json").resolve()),
        "red_flags": red_flags,
        "sanity": {
            "request_rows": summary["sanity"]["request_rows"],
            "batch_rows": summary["sanity"]["batch_rows"],
            "ledger_rows": summary["sanity"]["ledger_rows"],
            "artifact_bytes": {
                "n": len(sizes),
                "min": min(sizes.values()),
                "max": max(sizes.values()),
                "distinct_n": len(set(sizes.values())),
            },
            "invariants": {
                "zero_failures": return_code == 0,
                "gpu_visibility_disabled": True,
                "request_count_match": not red_flags,
                "nonnegative_file_sizes": all(size >= 0 for size in sizes.values()),
                "state_values_not_all_identical": summary["sanity"]["invariants"][
                    "batch_values_not_all_identical"
                ],
            },
        },
    }
    atomic_json(args.output / "result.json", result)
    if result["status"] != "PASS":
        raise RuntimeError(red_flags)
    return result


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--prepared-manifest", type=Path, required=True)
    result.add_argument("--output", type=Path, required=True)
    result.add_argument("--replayserve-root", type=Path, required=True)
    result.add_argument("--frontier-root", type=Path, required=True)
    result.add_argument("--cell", required=True)
    result.add_argument("--role", choices=("low1", "high1"), required=True)
    result.add_argument("--timeout-s", type=float, default=300.0)
    return result


def main() -> None:
    result = execute(parser().parse_args())
    print(
        json.dumps(
            {
                "status": result["status"],
                "runtime_s": result["runtime_s"],
                "request_rows": result["sanity"]["request_rows"],
                "batch_rows": result["sanity"]["batch_rows"],
                "ledger_rows": result["sanity"]["ledger_rows"],
                "red_flags": result["red_flags"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
