#!/usr/bin/env python3
"""Validate an HF weight export and its terminal Prime-RL metrics."""

from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from pathlib import Path


def merged_metrics(path: Path) -> dict[int, dict]:
    by_step: dict[int, dict] = defaultdict(dict)
    for line in path.read_text().splitlines():
        if line.strip():
            row = json.loads(line)
            by_step[int(row["step"])].update(row)
    return dict(by_step)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("model", type=Path)
    parser.add_argument("metrics", type=Path)
    parser.add_argument("--expected-step", type=int, required=True)
    parser.add_argument(
        "--scan-finite",
        action="store_true",
        help="load every floating tensor and reject NaN/Inf (slower)",
    )
    args = parser.parse_args()

    model = args.model.resolve()
    problems: list[str] = []
    if not (model / "STABLE").is_file():
        problems.append("missing STABLE marker")
    index_path = model / "model.safetensors.index.json"
    config_path = model / "config.json"
    if not index_path.is_file():
        problems.append("missing model.safetensors.index.json")
    if not config_path.is_file():
        problems.append("missing config.json")
    if problems:
        print(json.dumps({"healthy": False, "problems": problems}, indent=2))
        raise SystemExit(1)

    index = json.loads(index_path.read_text())
    weight_map = index.get("weight_map") or {}
    if not weight_map:
        problems.append("empty weight_map")
    expected_by_shard: dict[str, set[str]] = defaultdict(set)
    for key, shard in weight_map.items():
        expected_by_shard[str(shard)].add(str(key))

    from safetensors import safe_open

    tensor_count = 0
    parameter_elements = 0
    scanned_elements = 0
    dtypes: dict[str, int] = defaultdict(int)
    for shard_name, expected_keys in sorted(expected_by_shard.items()):
        shard_path = model / shard_name
        if not shard_path.is_file() or shard_path.stat().st_size == 0:
            problems.append(f"missing or empty shard: {shard_name}")
            continue
        try:
            with safe_open(shard_path, framework="pt", device="cpu") as handle:
                actual_keys = set(handle.keys())
                if actual_keys != expected_keys:
                    missing = len(expected_keys - actual_keys)
                    extra = len(actual_keys - expected_keys)
                    problems.append(
                        f"index/header mismatch in {shard_name}: missing={missing}, extra={extra}"
                    )
                for key in sorted(actual_keys):
                    view = handle.get_slice(key)
                    shape = tuple(view.get_shape())
                    elements = math.prod(shape)
                    tensor_count += 1
                    parameter_elements += elements
                    dtypes[str(view.get_dtype())] += 1
                    if args.scan_finite:
                        tensor = handle.get_tensor(key)
                        if tensor.is_floating_point() or tensor.is_complex():
                            scanned_elements += tensor.numel()
                            if not tensor.isfinite().all().item():
                                problems.append(f"non-finite tensor: {key}")
        except Exception as error:  # a corrupt header/shard must fail the gate
            problems.append(f"cannot load {shard_name}: {type(error).__name__}: {error}")

    by_step = merged_metrics(args.metrics)
    if args.expected_step not in by_step:
        problems.append(f"metrics do not contain expected step {args.expected_step}")
        terminal = {}
    else:
        terminal = by_step[args.expected_step]
        for field in ("loss/mean", "optim/grad_norm", "optim/lr"):
            value = terminal.get(field)
            if not isinstance(value, (int, float)) or not math.isfinite(value):
                problems.append(f"terminal metric {field} is not finite: {value!r}")
        if terminal.get("loss/nan_count", 0) != 0:
            problems.append(
                f"terminal loss/nan_count is {terminal.get('loss/nan_count')!r}"
            )

    result = {
        "healthy": not problems,
        "model": str(model),
        "expected_step": args.expected_step,
        "stable_marker": (model / "STABLE").is_file(),
        "shards": len(expected_by_shard),
        "tensors": tensor_count,
        "tensor_elements": parameter_elements,
        "dtype_tensor_counts": dict(sorted(dtypes.items())),
        "finite_scan_enabled": args.scan_finite,
        "finite_scanned_elements": scanned_elements,
        "terminal_metrics": {
            key: terminal.get(key)
            for key in (
                "loss/mean",
                "loss/nan_count",
                "optim/grad_norm",
                "optim/lr",
                "perf/throughput",
                "perf/peak_memory",
                "progress/num_tokens",
                "progress/num_samples",
            )
        },
        "problems": problems,
    }
    print(json.dumps(result, indent=2))
    if problems:
        raise SystemExit(1)


if __name__ == "__main__":
    main()