#!/usr/bin/env python3
"""Positive preflight and completion checks for the bounded single-GPU verl GRPO run."""

from __future__ import annotations

import argparse
import hashlib
import importlib
import importlib.metadata
import json
import math
import os
import re
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path


EXPECTED_COMMIT = "ddd86f527a4af75095e4677b02b5aa272913a088"
EXPECTED_DATASET_REVISION = "740312add88f781978c0658806c59bc2815b9866"
EXPECTED_MODEL_MANIFEST = "0ea1d330342b4b9efbd1c3648360fbc4c2e3b1d5abd6120e49ef837649001f7b"
REQUIRED_SCALARS = (
    "training/global_step",
    "actor/pg_loss",
    "actor/grad_norm",
    "actor/lr",
    "critic/score/mean",
)
FATAL_PATTERNS = (
    r"Traceback \(most recent call last\)",
    r"Error executing job with overrides",
    r"CUDA out of memory",
    r"RayTaskError",
    r"NCCL[^\n]*(?:error|unhandled|failed)",
    r"Segmentation fault",
    r"Killed(?:\s|$)",
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def run(*command: str) -> str:
    return subprocess.run(command, check=True, text=True, capture_output=True).stdout.strip()


def package_record(distribution: str, module_name: str | None = None) -> dict[str, object]:
    module = importlib.import_module(module_name or distribution.replace("-", "_"))
    try:
        distribution_version = importlib.metadata.version(distribution)
    except importlib.metadata.PackageNotFoundError:
        distribution_version = None
    return {
        "distribution_version": distribution_version,
        "module_version": getattr(module, "__version__", None),
        "module_file": getattr(module, "__file__", None),
    }


def atomic_json(path: Path, data: dict[str, object]) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    os.replace(temporary, path)


def preflight(args: argparse.Namespace) -> int:
    import pyarrow.parquet as pq
    import torch

    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    source_commit = run("git", "-C", str(args.verl_root), "rev-parse", "HEAD")
    if source_commit != EXPECTED_COMMIT:
        raise RuntimeError(f"Unexpected verl commit: {source_commit}")
    if run("git", "-C", str(args.verl_root), "status", "--porcelain"):
        raise RuntimeError("The pinned verl checkout is not clean")

    model_manifest_hash = sha256(args.model_root / "SHA256SUMS")
    if model_manifest_hash != EXPECTED_MODEL_MANIFEST:
        raise RuntimeError(f"Unexpected model manifest: {model_manifest_hash}")

    dataset_manifest = json.loads((args.data_root / "dataset-manifest.json").read_text(encoding="utf-8"))
    if dataset_manifest.get("revision") != EXPECTED_DATASET_REVISION:
        raise RuntimeError("Unexpected processed GSM8K revision")
    datasets = {}
    for split, expected_rows in (("train", 7473), ("test", 1319)):
        path = args.data_root / f"{split}.parquet"
        table = pq.read_table(path)
        if table.num_rows != expected_rows:
            raise RuntimeError(f"Unexpected {split} rows: {table.num_rows}")
        sample = table.slice(0, 1).to_pylist()[0]
        if sample["data_source"] != "openai/gsm8k" or sample["reward_model"]["style"] != "rule":
            raise RuntimeError(f"Invalid {split} semantic fields")
        datasets[split] = {"path": str(path), "rows": table.num_rows, "bytes": path.stat().st_size}

    if torch.cuda.device_count() != 1:
        raise RuntimeError(f"Expected exactly one visible GPU, got {torch.cuda.device_count()}")
    properties = torch.cuda.get_device_properties(0)
    gpu = {
        "name": properties.name,
        "capability": list(torch.cuda.get_device_capability(0)),
        "total_memory_bytes": properties.total_memory,
        "torch_cuda_version": torch.version.cuda,
    }
    nvidia_smi = run(
        "nvidia-smi",
        "--query-gpu=driver_version,name,memory.total",
        "--format=csv,noheader,nounits",
    ).splitlines()
    mount = json.loads(
        run("findmnt", "-J", "-T", str(args.storage_root), "-o", "TARGET,SOURCE,FSTYPE,OPTIONS")
    )["filesystems"][0]
    if mount["fstype"] == "overlay" or "rw" not in mount["options"].split(","):
        raise RuntimeError(f"Storage root is not a writable shared mount: {mount}")

    packages = {
        "torch": package_record("torch"),
        "ray": package_record("ray"),
        "transformers": package_record("transformers"),
        "sglang": package_record("sglang"),
        "pyarrow": package_record("pyarrow"),
        "tensorboard": package_record("tensorboard"),
        "verl": package_record("verl"),
    }
    record = {
        "status": "passed",
        "phase": "preflight",
        "checked_at": datetime.now(timezone.utc).isoformat(),
        "run_id": args.run_id,
        "target_steps": args.target_steps,
        "train_batch_size": args.train_batch_size,
        "rollout_n": args.rollout_n,
        "python": {"version": sys.version.split()[0], "executable": sys.executable},
        "nvidia_smi": nvidia_smi,
        "gpu": gpu,
        "mount": mount,
        "packages": packages,
        "verl_source": {"path": str(args.verl_root), "commit": source_commit, "clean": True},
        "model": {"path": str(args.model_root), "manifest_sha256": model_manifest_hash},
        "dataset": {"path": str(args.data_root), "manifest": dataset_manifest, "splits": datasets},
    }
    atomic_json(output_dir / "preflight.json", record)
    print(json.dumps(record, indent=2, ensure_ascii=False))
    return 0


def load_scalars(tensorboard_root: Path) -> tuple[list[Path], dict[str, list[dict[str, float | int]]]]:
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

    event_files = sorted(tensorboard_root.rglob("events.out.tfevents.*"))
    if not event_files or any(path.stat().st_size == 0 for path in event_files):
        raise RuntimeError("No non-empty TensorBoard event file was produced")
    scalars: dict[str, list[dict[str, float | int]]] = {}
    for event_file in event_files:
        accumulator = EventAccumulator(str(event_file), size_guidance={"scalars": 0})
        accumulator.Reload()
        for tag in accumulator.Tags().get("scalars", []):
            scalars.setdefault(tag, []).extend(
                {"step": event.step, "value": event.value, "wall_time": event.wall_time}
                for event in accumulator.Scalars(tag)
            )
    return event_files, scalars


def postflight(args: argparse.Namespace) -> int:
    output_dir = args.output_dir.resolve()
    log_text = args.training_log.read_text(encoding="utf-8", errors="replace")
    for pattern in FATAL_PATTERNS:
        if re.search(pattern, log_text, flags=re.IGNORECASE):
            raise RuntimeError(f"Fatal training marker found: {pattern}")
    required_log_markers = (
        "[validate_config] All configuration checks passed successfully!",
        f"Total training steps: {args.target_steps}",
        f"global_step_{args.target_steps}",
    )
    missing_markers = [marker for marker in required_log_markers if marker not in log_text]
    if missing_markers:
        raise RuntimeError(f"Training log is missing completion markers: {missing_markers}")

    event_files, scalars = load_scalars(output_dir / "tensorboard")
    metric_summary = {}
    for tag in REQUIRED_SCALARS:
        events = scalars.get(tag, [])
        if not events:
            raise RuntimeError(f"TensorBoard scalar is missing: {tag}")
        latest = max(events, key=lambda event: int(event["step"]))
        if int(latest["step"]) < args.target_steps or not math.isfinite(float(latest["value"])):
            raise RuntimeError(f"TensorBoard scalar did not reach a finite target step: {tag}={latest}")
        metric_summary[tag] = latest

    expected_rollouts = args.train_batch_size * args.rollout_n
    rollout_summary = {}
    for step in range(1, args.target_steps + 1):
        path = output_dir / "rollouts" / f"{step}.jsonl"
        rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
        if len(rows) != expected_rollouts:
            raise RuntimeError(f"Unexpected rollout count at step {step}: {len(rows)}")
        for row in rows:
            if row.get("step") != step or not row.get("input") or not isinstance(row.get("output"), str):
                raise RuntimeError(f"Invalid rollout row at step {step}")
            if not math.isfinite(float(row["score"])):
                raise RuntimeError(f"Non-finite rollout score at step {step}")
        rollout_summary[str(step)] = {
            "path": str(path),
            "rows": len(rows),
            "bytes": path.stat().st_size,
            "score_min": min(float(row["score"]) for row in rows),
            "score_max": max(float(row["score"]) for row in rows),
        }

    checkpoint_root = output_dir / "checkpoints"
    checkpoint = checkpoint_root / f"global_step_{args.target_steps}"
    actor = checkpoint / "actor"
    required_checkpoint_files = (
        actor / "model_world_size_1_rank_0.pt",
        actor / "optim_world_size_1_rank_0.pt",
        actor / "extra_state_world_size_1_rank_0.pt",
        checkpoint / "data.pt",
        checkpoint_root / "latest_checkpointed_iteration.txt",
    )
    checkpoint_files = {}
    for path in required_checkpoint_files:
        if not path.is_file() or path.stat().st_size == 0:
            raise RuntimeError(f"Checkpoint file is missing or empty: {path}")
        checkpoint_files[str(path.relative_to(output_dir))] = path.stat().st_size
    if (checkpoint_root / "latest_checkpointed_iteration.txt").read_text().strip() != str(args.target_steps):
        raise RuntimeError("Checkpoint tracker does not match the target step")

    record = {
        "status": "passed",
        "phase": "postflight",
        "verified_at": datetime.now(timezone.utc).isoformat(),
        "run_id": args.run_id,
        "target_steps": args.target_steps,
        "train_batch_size": args.train_batch_size,
        "rollout_n": args.rollout_n,
        "tensorboard": {
            "event_files": [{"path": str(path), "bytes": path.stat().st_size} for path in event_files],
            "metrics": metric_summary,
        },
        "rollouts": rollout_summary,
        "checkpoint": {"path": str(checkpoint), "files": checkpoint_files},
    }
    atomic_json(output_dir / "verification.json", record)
    (output_dir / "SUCCESS").write_text(datetime.now(timezone.utc).isoformat() + "\n", encoding="utf-8")
    print(json.dumps(record, indent=2, ensure_ascii=False))
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="phase", required=True)
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--run-id", required=True)
    common.add_argument("--output-dir", type=Path, required=True)
    common.add_argument("--target-steps", type=int, required=True)
    common.add_argument("--train-batch-size", type=int, required=True)
    common.add_argument("--rollout-n", type=int, required=True)

    before = subparsers.add_parser("preflight", parents=[common])
    before.add_argument("--storage-root", type=Path, required=True)
    before.add_argument("--verl-root", type=Path, required=True)
    before.add_argument("--model-root", type=Path, required=True)
    before.add_argument("--data-root", type=Path, required=True)

    after = subparsers.add_parser("postflight", parents=[common])
    after.add_argument("--training-log", type=Path, required=True)
    return parser.parse_args()


if __name__ == "__main__":
    parsed = parse_args()
    raise SystemExit(preflight(parsed) if parsed.phase == "preflight" else postflight(parsed))