#!/usr/bin/env python3
"""Validate the two-Worker GRPO tutorial run."""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import sys
from pathlib import Path
from typing import Any


def required_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"Missing environment variable: {name}")
    return value


def require_file(path: Path) -> None:
    if not path.is_file() or path.stat().st_size <= 0:
        raise RuntimeError(f"Required file is missing or empty: {path}")


def ordered_ray_nodes(ray: Any) -> list[dict[str, Any]]:
    nodes = [node for node in ray.nodes() if node.get("Alive")]
    current_node_id = str(ray.get_runtime_context().get_node_id())
    head = [node for node in nodes if str(node.get("NodeID")) == current_node_id]
    if len(head) != 1:
        raise RuntimeError(
            f"Cannot identify one Ray Head node: current={current_node_id} nodes={nodes}"
        )
    workers = sorted(
        (node for node in nodes if node is not head[0]),
        key=lambda node: str(node.get("NodeManagerAddress", "")),
    )
    return head + workers


def run_preflight(args: argparse.Namespace) -> None:
    import ray
    from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy

    rlinf_root = Path(required_env("RLINF_ROOT"))
    model_root = Path(required_env("MODEL_ROOT"))
    data_file = Path(required_env("DATA_FILE"))
    expected_commit = required_env("RLINF_COMMIT")

    ray.init(address="auto")
    nodes = ordered_ray_nodes(ray)
    if len(nodes) != args.expected_nodes:
        raise RuntimeError(
            f"Expected {args.expected_nodes} Ray nodes, found {len(nodes)}"
        )
    resources = ray.cluster_resources()
    expected_gpus = args.expected_nodes * args.expected_gpus_per_node
    if int(resources.get("GPU", 0)) != expected_gpus:
        raise RuntimeError(f"Expected {expected_gpus} Ray GPUs, found {resources}")

    @ray.remote(num_cpus=1, num_gpus=1)
    def probe_node(
        expected_rank: int,
        expected_gpus_per_node: int,
        rlinf_path: str,
        expected_rlinf_commit: str,
        model_root_path: str,
        data_file_path: str,
    ) -> dict[str, Any]:
        import hashlib
        import os
        import shutil
        import socket
        import subprocess
        import sys
        from importlib import import_module, metadata
        from pathlib import Path

        import ray
        import torch

        def file_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()

        rows = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,driver_version",
                "--format=csv,noheader,nounits",
            ],
            text=True,
        ).strip().splitlines()
        parsed = [row.rsplit(",", 1) for row in rows if row.strip()]
        gpu_names = [name.strip() for name, _ in parsed]
        drivers = sorted({driver.strip() for _, driver in parsed})
        if len(gpu_names) != expected_gpus_per_node:
            raise RuntimeError(
                f"rank={expected_rank}: expected {expected_gpus_per_node} GPUs, "
                f"found {gpu_names}"
            )
        if not all("A100-SXM4-80GB" in name for name in gpu_names):
            raise RuntimeError(f"rank={expected_rank}: unexpected GPUs={gpu_names}")
        if len(drivers) != 1:
            raise RuntimeError(f"rank={expected_rank}: driver versions={drivers}")
        if not sys.executable.startswith("/opt/venv/reason/bin/"):
            raise RuntimeError(f"rank={expected_rank}: Python={sys.executable}")
        if os.environ.get("NCCL_IB_DISABLE") != "1":
            raise RuntimeError(
                f"rank={expected_rank}: NCCL_IB_DISABLE="
                f"{os.environ.get('NCCL_IB_DISABLE')}"
            )

        configured_rank = os.environ.get("RLINF_NODE_RANK")
        if configured_rank is not None:
            if not configured_rank.isdigit() or int(configured_rank) != expected_rank:
                raise RuntimeError(
                    f"rank={expected_rank}: RLINF_NODE_RANK={configured_rank}"
                )

        commit = subprocess.check_output(
            ["git", "-C", rlinf_path, "rev-parse", "HEAD"], text=True
        ).strip()
        if commit != expected_rlinf_commit:
            raise RuntimeError(f"rank={expected_rank}: RLinf commit={commit}")

        model_root_local = Path(model_root_path)
        data_file_local = Path(data_file_path)
        model_manifest = model_root_local / "SHA256SUMS"
        data_manifest = data_file_local.parent / "SHA256SUMS"
        for path in (
            model_root_local / "config.json",
            model_root_local / "tokenizer_config.json",
            model_root_local / "model.safetensors",
            model_manifest,
            data_file_local,
            data_manifest,
        ):
            if not path.is_file() or path.stat().st_size <= 0:
                raise RuntimeError(f"rank={expected_rank}: missing input={path}")
            with path.open("rb") as stream:
                stream.read(1)

        try:
            distribution = metadata.distribution("flash-attn-4")
        except metadata.PackageNotFoundError:
            disabled_candidates = list(
                Path(sys.prefix).glob("lib/python*/site-packages/*.dist-info.a100-disabled")
            )
            fa4_action = (
                "metadata-already-disabled"
                if any("flash_attn_4" in path.name for path in disabled_candidates)
                else "metadata-absent"
            )
        else:
            dist_info = Path(distribution._path).resolve()
            venv_root = Path(sys.prefix).resolve()
            if venv_root not in dist_info.parents:
                raise RuntimeError(
                    f"rank={expected_rank}: unexpected flash-attn-4 path={dist_info}"
                )
            disabled = Path(f"{dist_info}.a100-disabled")
            if disabled.exists():
                raise RuntimeError(
                    f"rank={expected_rank}: both active and disabled FA4 metadata exist"
                )
            shutil.move(str(dist_info), str(disabled))
            fa4_action = "metadata-disabled"

        try:
            visible_fa4 = metadata.version("flash-attn-4")
        except metadata.PackageNotFoundError:
            pass
        else:
            raise RuntimeError(
                f"rank={expected_rank}: flash-attn-4 remains visible={visible_fa4}"
            )

        if not torch.cuda.is_available():
            raise RuntimeError(f"rank={expected_rank}: CUDA is unavailable")
        capability = ".".join(map(str, torch.cuda.get_device_capability(0)))
        if capability != "8.0":
            raise RuntimeError(
                f"rank={expected_rank}: compute capability={capability}"
            )
        import_module("transformer_engine.pytorch")
        import_module("sglang")

        context = ray.get_runtime_context()
        return {
            "rank": expected_rank,
            "node_id": str(context.get_node_id()),
            "hostname": socket.gethostname(),
            "driver_version": drivers[0],
            "gpu_count": len(gpu_names),
            "gpu_name": gpu_names[0],
            "compute_capability": capability,
            "fa4_action": fa4_action,
            "rlinf_commit": commit,
            "model_manifest": file_sha256(model_manifest),
            "data_manifest": file_sha256(data_manifest),
        }

    refs = []
    for expected_rank, node in enumerate(nodes):
        refs.append(
            probe_node.options(
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    node_id=node["NodeID"], soft=False
                )
            ).remote(
                expected_rank,
                args.expected_gpus_per_node,
                str(rlinf_root),
                expected_commit,
                str(model_root),
                str(data_file),
            )
        )
    workers = sorted(ray.get(refs), key=lambda item: item["rank"])
    if len({item["node_id"] for item in workers}) != args.expected_nodes:
        raise RuntimeError(f"Probes did not run on distinct Ray nodes: {workers}")
    if len({item["driver_version"] for item in workers}) != 1:
        raise RuntimeError(f"NVIDIA driver versions differ: {workers}")
    if len({item["model_manifest"] for item in workers}) != 1:
        raise RuntimeError(f"Model manifests differ between Workers: {workers}")
    if len({item["data_manifest"] for item in workers}) != 1:
        raise RuntimeError(f"Dataset manifests differ between Workers: {workers}")

    report = {
        "cluster_resources": resources,
        "workers": workers,
    }
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )

    print("RAY_CLUSTER_RESOURCES=" + json.dumps(resources, sort_keys=True))
    for item in workers:
        role = "Ray Head" if item["rank"] == 0 else "Ray Worker"
        print(
            "WORKER_TOPOLOGY="
            f"rank={item['rank']} role={role} host={item['hostname']} "
            f"gpu_count={item['gpu_count']} gpu={item['gpu_name']}"
        )
        print(
            f"NVIDIA_DRIVER_VERSION_NODE_{item['rank']}={item['driver_version']}"
        )
        print(
            "A100_RUNTIME="
            f"rank={item['rank']} "
            f"compute_capability={item['compute_capability']} "
            f"fa4_action={item['fa4_action']}"
        )
    print(f"RLINF_PREFLIGHT_REPORT={args.report}")
    print("RLINF_PREFLIGHT_STATUS=passed")


def run_postflight(args: argparse.Namespace) -> None:
    from omegaconf import OmegaConf
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

    output_dir = args.output_dir
    driver_log = output_dir / "driver.log"
    main_log = output_dir / "log" / "main.log"
    preflight_path = output_dir / "preflight.json"
    hydra_config_path = output_dir / "hydra" / ".hydra" / "config.yaml"
    for path in (driver_log, main_log, preflight_path, hydra_config_path):
        require_file(path)

    driver_text = driver_log.read_text(encoding="utf-8", errors="replace")
    main_text = main_log.read_text(encoding="utf-8", errors="replace")
    combined_text = driver_text + "\n" + main_text
    fatal_patterns = {
        "python traceback": r"Traceback \(most recent call last\):",
        "Hydra execution error": r"Error executing job with overrides",
        "Ray task error": r"RayTaskError\(",
        "Ray actor failure": r"(?:ActorDiedError|ActorUnavailableError)",
        "CUDA out of memory": r"CUDA out of memory",
        "preflight failure": r"RLINF_PREFLIGHT_STATUS=failed",
    }
    detected = [
        label
        for label, pattern in fatal_patterns.items()
        if re.search(pattern, combined_text)
    ]
    if detected:
        raise RuntimeError(f"Fatal markers found in logs: {detected}")
    if "RLINF_PYTHON_EXIT_CODE=0" not in driver_text:
        raise RuntimeError("RLinf Python did not report exit code 0")
    completion_line = (
        f"Step limit given by max_steps={args.target_steps} reached. Stopping run"
    )
    if completion_line not in combined_text:
        raise RuntimeError(f"Missing RLinf step-limit marker: {completion_line}")
    preflight = json.loads(preflight_path.read_text(encoding="utf-8"))
    workers = preflight.get("workers", [])

    cfg = OmegaConf.load(hydra_config_path)
    if int(cfg.cluster.num_nodes) != args.expected_nodes:
        raise RuntimeError(f"Resolved cluster.num_nodes={cfg.cluster.num_nodes}")
    placement = OmegaConf.to_container(cfg.cluster.component_placement, resolve=True)
    expected_placement = {
        "rollout": 0,
        "inference": 1,
        "actor": "2-3",
        "reward": 0,
    }
    if placement != expected_placement:
        raise RuntimeError(f"Resolved component placement={placement}")
    if int(cfg.runner.max_steps) != args.target_steps:
        raise RuntimeError(f"Resolved runner.max_steps={cfg.runner.max_steps}")
    if str(cfg.runner.output_dir) != str(output_dir.parent):
        raise RuntimeError(f"Resolved runner.output_dir={cfg.runner.output_dir}")
    if str(cfg.runner.experiment_name) != output_dir.name:
        raise RuntimeError(
            f"Resolved runner.experiment_name={cfg.runner.experiment_name}"
        )

    tensorboard_dir = output_dir / "tensorboard"
    events = sorted(tensorboard_dir.rglob("events.out.tfevents.*"))
    if not events or not all(path.stat().st_size > 0 for path in events):
        raise RuntimeError(f"TensorBoard event file is missing: {tensorboard_dir}")
    required_tags = (
        "rollout/reward_scores",
        "actor/training/actor/final_loss",
        "actor/training/actor/grad_norm",
        "actor/training/actor/lr",
    )
    accumulator = None
    observed_tags = {}
    for event_dir in sorted({path.parent for path in events}):
        candidate = EventAccumulator(str(event_dir), size_guidance={"scalars": 0})
        candidate.Reload()
        scalar_tags = set(candidate.Tags().get("scalars", []))
        observed_tags[str(event_dir)] = sorted(scalar_tags)
        if set(required_tags).issubset(scalar_tags):
            accumulator = candidate
            break
    if accumulator is None:
        raise RuntimeError(
            f"No TensorBoard event directory contains the required scalar tags: "
            f"{observed_tags}"
        )
    metrics = {}
    for tag in required_tags:
        points = accumulator.Scalars(tag)
        if not points:
            raise RuntimeError(f"TensorBoard scalar has no values: {tag}")
        value = float(points[-1].value)
        if not math.isfinite(value):
            raise RuntimeError(f"TensorBoard scalar is not finite: {tag}={value}")
        metrics[tag] = {"step": int(points[-1].step), "value": value}
    if metrics["actor/training/actor/lr"]["value"] <= 0:
        raise RuntimeError(f"Actor learning rate is not positive: {metrics}")

    checkpoint_dir = output_dir / "checkpoints" / f"global_step_{args.target_steps}"
    actor_dcp = checkpoint_dir / "actor" / "dcp_checkpoint"
    data_state = checkpoint_dir / "data" / "data.pt"
    require_file(actor_dcp / ".metadata")
    require_file(data_state)
    shards = sorted(path for path in actor_dcp.glob("*.distcp") if path.stat().st_size > 0)
    if len(shards) < 2:
        raise RuntimeError(f"Expected at least two FSDP checkpoint shards: {shards}")

    report = {
        "checkpoint": {
            "directory": str(checkpoint_dir),
            "fsdp_shards": [
                {"name": path.name, "bytes": path.stat().st_size} for path in shards
            ],
            "data_state_bytes": data_state.stat().st_size,
        },
        "metrics": metrics,
        "workers": workers,
        "resolved_component_placement": placement,
        "target_steps": args.target_steps,
    }
    report_path = output_dir / "verification.json"
    report_path.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    for tag, point in metrics.items():
        print(f"GRPO_METRIC={tag} step={point['step']} value={point['value']}")
    print(f"ACTOR_FSDP_SHARDS={len(shards)}")
    print(f"CHECKPOINT_DIR={checkpoint_dir}")
    print(f"RLINF_VERIFICATION_REPORT={report_path}")
    print("RLINF_VERIFICATION_STATUS=passed")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    preflight = subparsers.add_parser("preflight")
    preflight.add_argument("--expected-nodes", type=int, required=True)
    preflight.add_argument("--expected-gpus-per-node", type=int, required=True)
    preflight.add_argument("--report", type=Path, required=True)

    postflight = subparsers.add_parser("postflight")
    postflight.add_argument("--output-dir", type=Path, required=True)
    postflight.add_argument("--target-steps", type=int, required=True)
    postflight.add_argument("--expected-nodes", type=int, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "preflight":
        run_preflight(args)
    else:
        run_postflight(args)


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(
            f"RLINF_{'PREFLIGHT' if 'preflight' in sys.argv else 'VERIFICATION'}_STATUS=failed "
            f"error={type(error).__name__}: {error}",
            file=sys.stderr,
        )
        raise