#!/usr/bin/env python3
"""Validate transferred RoboTwin assets before a smoke or Slurm launch."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from train_frameskip_robotwin_continuation import load_cache


def manifest_support(path: Path) -> set[tuple[str, str]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    support = set()
    for dataset in payload["datasets"]:
        name = dataset.get("name", dataset.get("dataset_name"))
        for episode in dataset.get("episodes", dataset.get("trajectories", [])):
            trajectory_id = episode.get("trajectory_id", episode.get("id"))
            support.add((name, str(trajectory_id)))
    if not support:
        raise RuntimeError("valid-anchor manifest has no trajectory support")
    return support


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True, type=Path)
    parser.add_argument("--manifest", required=True, type=Path)
    parser.add_argument("--cache", required=True, type=Path)
    args = parser.parse_args()

    commit_path = args.checkpoint / "checkpoint.json"
    commit = json.loads(commit_path.read_text(encoding="utf-8"))
    if commit.get("schema_version") != 1 or int(commit.get("completed_steps", -1)) != 10000:
        raise RuntimeError("checkpoint commit is not a complete step-10000 state")
    if int(commit.get("world_size", -1)) != 2:
        raise RuntimeError("10K checkpoint world size must be 2")
    inventory = commit.get("files", [])
    missing = [
        entry["name"]
        for entry in inventory
        if not (args.checkpoint / entry["name"]).is_file()
        or (args.checkpoint / entry["name"]).stat().st_size != int(entry["bytes"])
    ]
    if not inventory or missing:
        raise RuntimeError(f"incomplete 10K checkpoint inventory: missing={missing}")
    parent_config = args.checkpoint.parent.parent / "config.full.yaml"
    if not parent_config.is_file():
        raise RuntimeError(f"missing parent config required by exact resume: {parent_config}")

    cache, payload = load_cache(str(args.cache))
    support = manifest_support(args.manifest)
    cache_support = set(cache)
    if support != cache_support:
        raise RuntimeError(
            "FrameSkip cache support differs from manifest support: "
            f"missing={len(support - cache_support)}, extra={len(cache_support - support)}"
        )

    print("FRAMESKIP_ROBOTWIN_ASSET_PREFLIGHT=PASS")
    print(f"checkpoint_files={len(inventory)}")
    print(f"trajectory_support={len(support)}")
    print(f"cache_records={len(payload['records'])}")


if __name__ == "__main__":
    main()