#!/usr/bin/env python3
"""Run and score a frozen Frontier probe manifest, CPU only."""

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


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 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 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 execute(args: argparse.Namespace) -> dict[str, Any]:
    args.output = args.output.resolve()
    prepared = json.loads(args.prepared_manifest.read_text(encoding="utf-8"))
    if prepared["status"] != "PASS":
        raise RuntimeError("prepared simulator manifest did not pass")
    expected_runs = int(prepared.get("expected_runs", len(prepared["entries"])))
    if expected_runs != len(prepared["entries"]):
        raise RuntimeError(
            f"prepared manifest expected {expected_runs} runs but contains "
            f"{len(prepared['entries'])} entries"
        )
    driver = load_module(
        "simfid_execution_driver",
        args.replayserve_root
        / "runs/simfid_s2rb/results/execution_driver.py",
    )
    head = git_capture(args.frontier_root, "rev-parse", "HEAD").strip()
    status_short = git_capture(args.frontier_root, "status", "--short")
    aituner_root = Path(__file__).resolve().parents[2]
    aituner_head = git_capture(aituner_root, "rev-parse", "HEAD").strip()
    aituner_status_short = git_capture(aituner_root, "status", "--short")
    results = []
    failures = []
    gpu_visibility_disabled = True
    for sequence, entry in enumerate(prepared["entries"]):
        run_root = args.output / f"{sequence:03d}_{entry['fixture_id']}"
        scorer_path = run_root / "scorer_output.json"
        if scorer_path.is_file() and args.resume:
            scorer = json.loads(scorer_path.read_text(encoding="utf-8"))
            results.append({**entry, "sequence": sequence, "scorer": scorer, "resumed": True})
            continue
        run_root.mkdir(parents=True, exist_ok=True)
        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 = run_root / "frontier_metrics"
        run_id = f"fidelity_p1_frontier_{sequence:02d}_{entry['cell']}_{entry['role']}"
        knobs = config["frontier"]["knobs"]
        command = driver.build_command(
            trace_path=trace_path,
            metrics_root=metrics_root,
            run_id=run_id,
            knobs=knobs,
        )
        driver.audit_command(command, knobs)
        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)
        gpu_visibility_disabled = gpu_visibility_disabled and (
            environment.get("CUDA_VISIBLE_DEVICES") == ""
            and environment.get("NVIDIA_VISIBLE_DEVICES") == "void"
        )
        run_manifest = {
            "schema": "fidelity-p1-frontier-run-v1",
            "sequence": sequence,
            "cell": entry["cell"],
            "role": entry["role"],
            "anchor": entry["anchor"],
            "request_count": entry["selected_count"],
            "frontier": {
                "root": str(args.frontier_root.resolve()),
                "git_head": head,
                "git_status_short": status_short,
            },
            "runner": {
                "script": str(Path(__file__).resolve()),
                "script_sha256": sha256_file(Path(__file__).resolve()),
                "aituner_git_head": aituner_head,
                "aituner_git_status_short": aituner_status_short,
            },
            "inputs": {
                "config": str(config_path),
                "config_sha256": sha256_file(config_path),
                "fixture_manifest": str(fixture_manifest_path),
                "fixture_manifest_sha256": sha256_file(fixture_manifest_path),
                "frontier_csv": str(trace_path),
                "frontier_csv_sha256": sha256_file(trace_path),
                "sidecar": str(sidecar_path),
                "sidecar_sha256": sha256_file(sidecar_path),
            },
            "environment": {
                key: environment[key]
                for key in (
                    "PYTHONPATH",
                    "FRONTIER_EXECUTION_TIME_SCALE",
                    "CUDA_VISIBLE_DEVICES",
                    "NVIDIA_VISIBLE_DEVICES",
                    "FRONTIER_LOG_LEVEL",
                )
            },
            "command": command,
            "contains_prompt_text": False,
        }
        atomic_json(run_root / "run_manifest.json", run_manifest)
        start = time.time()
        with (run_root / "stdout.log").open("w", encoding="utf-8") as stdout, (
            run_root / "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 = time.time() - start
        if return_code != 0:
            failure = {
                "sequence": sequence,
                "cell": entry["cell"],
                "role": entry["role"],
                "return_code": return_code,
                "runtime_s": runtime,
            }
            failures.append(failure)
            atomic_json(run_root / "failure.json", failure)
            break
        system_path, request_path = driver.find_metrics(run_root)
        scorer = driver.score_trial(row, system_path, request_path)
        scorer["runtime_s"] = runtime
        atomic_json(scorer_path, scorer)
        results.append({**entry, "sequence": sequence, "scorer": scorer, "resumed": False})
        print(
            json.dumps(
                {
                    "sequence": sequence,
                    "cell": entry["cell"],
                    "role": entry["role"],
                    "runtime_s": runtime,
                    "sim_pass_rate": scorer["slo"]["pass_rate"],
                    "sim_feasible": scorer["slo"]["feasible"],
                },
                sort_keys=True,
            ),
            flush=True,
        )

    pass_rates = [float(item["scorer"]["slo"]["pass_rate"]) for item in results]
    throughputs = [
        float(item["scorer"]["throughput_requests_per_second_per_gpu"])
        for item in results
    ]
    runtimes = [float(item["scorer"]["runtime_s"]) for item in results]
    red_flags = []
    if failures:
        red_flags.append("frontier_run_failure")
    if len(results) != expected_runs:
        red_flags.append("runs_not_expected")
    if any(not 0.0 <= value <= 1.0 for value in pass_rates):
        red_flags.append("pass_rate_out_of_range")
    if any(value <= 0 for value in throughputs):
        red_flags.append("nonpositive_throughput")
    result = {
        "schema": "fidelity-p1-frontier-result-v1",
        "status": "PASS" if not red_flags else "STOP",
        "prepared_manifest": str(args.prepared_manifest.resolve()),
        "prepared_manifest_sha256": sha256_file(args.prepared_manifest),
        "frontier": {
            "root": str(args.frontier_root.resolve()),
            "git_head": head,
            "git_status_short": status_short,
        },
        "runner": {
            "script": str(Path(__file__).resolve()),
            "script_sha256": sha256_file(Path(__file__).resolve()),
            "aituner_git_head": aituner_head,
            "aituner_git_status_short": aituner_status_short,
        },
        "results": results,
        "failures": failures,
        "sanity": {
            "red_flags": red_flags,
            "n": len(results),
            "pass_rate": {
                "n": len(pass_rates),
                "min": min(pass_rates) if pass_rates else None,
                "max": max(pass_rates) if pass_rates else None,
                "distinct_n": len(set(pass_rates)),
            },
            "throughput_per_gpu": {
                "n": len(throughputs),
                "min": min(throughputs) if throughputs else None,
                "max": max(throughputs) if throughputs else None,
                "distinct_n": len(set(throughputs)),
            },
            "runtime_s": {
                "n": len(runtimes),
                "min": min(runtimes) if runtimes else None,
                "max": max(runtimes) if runtimes else None,
                "distinct_n": len(set(runtimes)),
            },
            "invariants": {
                "runs_expected": len(results) == expected_runs,
                "expected_runs": expected_runs,
                "zero_failures": not failures,
                "ratios_bounded": all(0.0 <= value <= 1.0 for value in pass_rates),
                "nonnegative_metrics": all(value > 0 for value in throughputs),
                "per_config_not_identical": len(set(pass_rates)) > 1,
                "gpu_visibility_disabled": gpu_visibility_disabled,
            },
        },
    }
    atomic_json(args.output / "metrics.json", result)
    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("--timeout-s", type=float, default=900.0)
    result.add_argument("--resume", action="store_true")
    return result


def main() -> None:
    result = execute(parser().parse_args())
    print(
        json.dumps(
            {
                "status": result["status"],
                "runs": len(result["results"]),
                "red_flags": result["sanity"]["red_flags"],
            },
            sort_keys=True,
        )
    )
    if result["status"] != "PASS":
        raise RuntimeError(result["sanity"]["red_flags"])


if __name__ == "__main__":
    main()
