#!/usr/bin/env python3
"""Replay the two Qwen235 Fixed-PD A1 cells with Frontier state outputs.

The frozen command is the experimental input.  This runner changes only the
metrics output/run id and the two existing state-observation flags, then checks
that request-level scorer inputs are byte-identical to the frozen A1 run.
"""

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
REPO_ROOT = HERE.parents[1]
TELEMETRY_ROOT = REPO_ROOT / "runs/telemetry-residual"
sys.path.insert(0, str(TELEMETRY_ROOT))

from common_state import summarize_frontier  # noqa: E402
from run_frontier_state import enable_state_outputs, find_state_metrics  # noqa: E402


CONFIGS = ("tp4_ep1_mns64", "tp8_ep8_mns64")
ALLOWED_FLAG_CHANGES = {
    "--metrics_config_output_dir",
    "--metrics_config_run_id",
    "--no-metrics_config_store_frontier_stage_batch_ledger",
    "--metrics_config_store_frontier_stage_batch_ledger",
    "--no-metrics_config_keep_individual_batch_metrics",
    "--metrics_config_keep_individual_batch_metrics",
    "--metrics_config_enable_op_level_tracing",
}


def load_module(name: str, path: Path):
    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 replace_option(command: list[str], option: str, value: str) -> list[str]:
    result = list(command)
    if result.count(option) != 1:
        raise ValueError(f"expected exactly one {option}, found {result.count(option)}")
    index = result.index(option)
    if index + 1 >= len(result) or result[index + 1].startswith("--"):
        raise ValueError(f"{option} has no value")
    result[index + 1] = value
    return result


def option_value(command: list[str], option: str) -> str:
    if command.count(option) != 1:
        raise ValueError(f"expected exactly one {option}, found {command.count(option)}")
    return command[command.index(option) + 1]


def _semantic_options(command: list[str]) -> dict[str, tuple[str, ...]]:
    """Parse CLI tokens sufficiently to audit this controlled command edit."""

    result: dict[str, tuple[str, ...]] = {}
    index = 0
    while index < len(command):
        token = command[index]
        if not token.startswith("--"):
            index += 1
            continue
        values: list[str] = []
        index += 1
        while index < len(command) and not command[index].startswith("--"):
            values.append(command[index])
            index += 1
        if token in result:
            raise ValueError(f"duplicate option in frozen command: {token}")
        result[token] = tuple(values)
    return result


def transform_command(
    command: list[str], *, metrics_root: Path, run_id: str, op_trace: bool = False
) -> list[str]:
    result = replace_option(
        command, "--metrics_config_output_dir", str(metrics_root.resolve())
    )
    result = replace_option(result, "--metrics_config_run_id", run_id)
    result = enable_state_outputs(result)
    op_trace_flag = "--metrics_config_enable_op_level_tracing"
    if op_trace:
        if op_trace_flag in result:
            raise ValueError(f"frozen command already enables {op_trace_flag}")
        result.append(op_trace_flag)

    before = _semantic_options(command)
    after = _semantic_options(result)
    changed = {
        option
        for option in set(before) | set(after)
        if before.get(option) != after.get(option)
    }
    if not changed <= ALLOWED_FLAG_CHANGES:
        raise ValueError(f"unapproved command changes: {sorted(changed - ALLOWED_FLAG_CHANGES)}")
    required = {
        "--metrics_config_output_dir",
        "--metrics_config_run_id",
        "--no-metrics_config_store_frontier_stage_batch_ledger",
        "--metrics_config_store_frontier_stage_batch_ledger",
        "--metrics_config_keep_individual_batch_metrics",
    }
    if not required <= changed:
        raise ValueError(f"required controlled changes missing: {sorted(required - changed)}")
    if op_trace and op_trace_flag not in changed:
        raise ValueError("op-trace observation flag was not audited as a command change")
    return result


def environment(frontier_source: Path, python_deps: Path) -> dict[str, str]:
    result = os.environ.copy()
    result.update(
        {
            "PYTHONPATH": ":".join([str(python_deps), str(frontier_source)]),
            "CUDA_VISIBLE_DEVICES": "",
            "NVIDIA_VISIBLE_DEVICES": "void",
            "WANDB_DISABLED": "true",
            "VIDUR_DISABLE_WANDB": "1",
            "FRONTIER_LOG_LEVEL": "WARNING",
            "PYTHONDONTWRITEBYTECODE": "1",
        }
    )
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-sim-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--frontier-source", type=Path, required=True)
    parser.add_argument("--python-deps", type=Path, required=True)
    parser.add_argument("--config", action="append", choices=CONFIGS)
    parser.add_argument("--timeout-seconds", type=float, default=1200)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--op-trace", action="store_true")
    return parser.parse_args()


def run_cell(
    *,
    config: str,
    base_sim_root: Path,
    output_root: Path,
    frontier_source: Path,
    python_deps: Path,
    timeout_seconds: float,
    resume: bool,
    op_trace: bool,
    q30: Any,
) -> dict[str, Any]:
    base_run = base_sim_root / "runs" / config / "eval"
    base_command_path = base_run / "command.json"
    base_result_path = base_run / "result.json"
    if not base_command_path.is_file() or not base_result_path.is_file():
        raise FileNotFoundError(f"frozen A1 inputs missing under {base_run}")
    base_command = json.loads(base_command_path.read_text())
    base_result = json.loads(base_result_path.read_text())
    if base_result.get("status") != "completed":
        raise ValueError(f"frozen A1 result is not completed: {base_result_path}")

    run_root = output_root / config
    result_path = run_root / "result.json"
    if resume and result_path.is_file():
        previous = json.loads(result_path.read_text())
        if (
            previous.get("status") == "PASS"
            and previous.get("inputs", {}).get("base_command_sha256")
            == sha256_file(base_command_path)
            and previous.get("inputs", {}).get("base_result_sha256")
            == sha256_file(base_result_path)
            and bool(previous.get("op_trace_enabled")) == op_trace
        ):
            return previous
    if run_root.exists() and any(run_root.iterdir()):
        raise FileExistsError(f"refusing non-empty output: {run_root}")
    run_root.mkdir(parents=True, exist_ok=True)

    command = transform_command(
        base_command,
        metrics_root=run_root / "frontier_metrics",
        run_id=f"qwen235_fixed_pd_state_{config}",
        op_trace=op_trace,
    )
    atomic_json(run_root / "command.json", command)
    manifest = {
        "schema": "qwen235-fixed-pd-state-replay-v1",
        "config": config,
        "inputs": {
            "base_run": str(base_run.resolve()),
            "base_command_sha256": sha256_file(base_command_path),
            "base_result_sha256": sha256_file(base_result_path),
            "base_request_metrics_sha256": base_result["request_metrics_sha256"],
        },
        "controlled_changes": [
            "metrics output directory",
            "metrics run id",
            "full Frontier stage/batch ledger enabled",
            "individual batch metrics enabled",
        ] + (["op-level tracing enabled"] if op_trace else []),
        "command_sha256": sha256_file(run_root / "command.json"),
        "frontier": {
            "source": str(frontier_source),
            "git_head": subprocess.check_output(
                ["git", "-C", str(frontier_source), "rev-parse", "HEAD"], text=True
            ).strip(),
        },
        "environment": {
            "PYTHONPATH": ":".join([str(python_deps), str(frontier_source)]),
            "CUDA_VISIBLE_DEVICES": "",
            "NVIDIA_VISIBLE_DEVICES": "void",
            "FRONTIER_LOG_LEVEL": "WARNING",
        },
        "op_trace_enabled": op_trace,
    }
    atomic_json(run_root / "run_manifest.json", manifest)

    started = time.monotonic()
    with (run_root / "stdout.log").open("w") as stdout, (
        run_root / "stderr.log"
    ).open("w") as stderr:
        try:
            completed = subprocess.run(
                command,
                cwd=frontier_source,
                env=environment(frontier_source, python_deps),
                stdout=stdout,
                stderr=stderr,
                timeout=timeout_seconds,
                check=False,
            )
            returncode = int(completed.returncode)
        except subprocess.TimeoutExpired:
            returncode = 124
    elapsed_seconds = time.monotonic() - started
    if returncode != 0:
        failure = {
            "status": "STOP",
            "config": config,
            "returncode": returncode,
            "elapsed_seconds": elapsed_seconds,
        }
        atomic_json(run_root / "failure.json", failure)
        raise RuntimeError(f"state replay failed: {failure}")

    fallback_evidence = q30.collective_fallback_evidence(run_root)
    if fallback_evidence:
        raise RuntimeError(f"collective-profile fallback detected: {fallback_evidence}")
    paths = find_state_metrics(run_root)
    request_metrics_sha256 = sha256_file(paths["requests"])
    scorer_equivalent = request_metrics_sha256 == base_result["request_metrics_sha256"]
    if not scorer_equivalent:
        raise RuntimeError(
            f"observation changed scorer input for {config}: "
            f"{request_metrics_sha256} != {base_result['request_metrics_sha256']}"
        )
    state = summarize_frontier(
        system_metrics_path=paths["system"],
        request_metrics_path=paths["requests"],
        batch_metrics_path=paths["batches"],
        ledger_path=paths["ledger"],
    )
    atomic_json(run_root / "common-state.json", state)
    if op_trace:
        op_trace_matches = sorted(run_root.glob("frontier_metrics/**/op_traces.jsonl"))
        if len(op_trace_matches) != 1:
            raise ValueError(
                f"expected one op trace for {config}, found {len(op_trace_matches)}"
            )
        paths["op_trace"] = op_trace_matches[0]
    result = {
        "schema": "qwen235-fixed-pd-state-replay-result-v1",
        "status": "PASS",
        "config": config,
        "elapsed_seconds": elapsed_seconds,
        "returncode": returncode,
        "inputs": manifest["inputs"],
        "scorer_equivalence": {
            "request_metrics_byte_identical": scorer_equivalent,
            "request_metrics_sha256": request_metrics_sha256,
            "base_metrics": base_result["metrics"],
        },
        "state_artifacts": {
            name: {"path": str(path.resolve()), "sha256": sha256_file(path)}
            for name, path in paths.items()
        },
        "common_state": state,
        "collective_fallback_evidence": fallback_evidence,
        "op_trace_enabled": op_trace,
    }
    atomic_json(result_path, result)
    return result


def main() -> None:
    args = parse_args()
    for name in ("base_sim_root", "output_root", "frontier_source", "python_deps"):
        setattr(args, name, getattr(args, name).resolve())
    selected = tuple(args.config or CONFIGS)
    q30 = load_module("q235_state_q30", HERE / "run_frontier_qwen30_exact_trace_surface.py")
    results = []
    for config in selected:
        result = run_cell(
            config=config,
            base_sim_root=args.base_sim_root,
            output_root=args.output_root,
            frontier_source=args.frontier_source,
            python_deps=args.python_deps,
            timeout_seconds=args.timeout_seconds,
            resume=args.resume,
            op_trace=args.op_trace,
            q30=q30,
        )
        results.append(result)
        print(json.dumps({"config": config, "status": result["status"]}), flush=True)
    aggregate = {
        "schema": "qwen235-fixed-pd-state-replay-aggregate-v1",
        "status": "PASS" if all(result["status"] == "PASS" for result in results) else "STOP",
        "results": results,
    }
    atomic_json(args.output_root / "state_replay.json", aggregate)


if __name__ == "__main__":
    main()