#!/usr/bin/env python3
"""Fail-fast acceptance test for the isolated VBench scoring environment."""

from __future__ import annotations

import argparse
import importlib
import json
import os
import sys
from pathlib import Path


MAIN_DIMENSIONS = (
    "subject_consistency",
    "background_consistency",
    "motion_smoothness",
    "dynamic_degree",
    "aesthetic_quality",
    "imaging_quality",
    "object_class",
    "multiple_objects",
    "human_action",
    "color",
    "spatial_relationship",
    "scene",
    "temporal_style",
    "appearance_style",
    "overall_consistency",
)

REQUIRED_WEIGHTS = {
    "/root/.cache/vbench/amt_model/amt-s.pth": 10_000_000,
    "/root/.cache/vbench/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth": 500_000_000,
    "/root/.cache/vbench/grit_model/grit_b_densecap_objectdet.pth": 400_000_000,
    "/root/.cache/vbench/caption_model/tag2text_swin_14m.pth": 4_000_000_000,
    "/root/.cache/vbench/ViCLIP/ViClip-InternVid-10M-FLT.pth": 1_500_000_000,
    "/root/.cache/vbench/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth": 100_000_000,
    "/root/.cache/vbench/raft_model/models/raft-things.pth": 20_000_000,
    "/root/.cache/torch/hub/checkpoints/dino_vitbase16_pretrain.pth": 300_000_000,
    "/root/.cache/clip/ViT-B-32.pt": 300_000_000,
    "/root/.cache/clip/ViT-L-14.pt": 900_000_000,
}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vbench-repo", default=os.environ.get("VBENCH_REPO", "/data/wxy/VBench"))
    parser.add_argument("--skip-dino-load", action="store_true")
    args = parser.parse_args()

    repo = Path(args.vbench_repo).resolve()
    if not (repo / "vbench" / "VBench_full_info.json").is_file():
        raise RuntimeError(f"invalid VBench repository: {repo}")
    sys.path.insert(0, str(repo))
    os.chdir(repo)

    report: dict[str, object] = {
        "python": sys.version,
        "executable": sys.executable,
        "repo": str(repo),
        "imports": {},
        "weights": {},
    }

    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("torch.cuda.is_available() is false")
    x = torch.randn(512, 512, device="cuda")
    torch.cuda.synchronize()
    report["cuda"] = {
        "torch": torch.__version__,
        "runtime": torch.version.cuda,
        "capability": list(torch.cuda.get_device_capability()),
        "probe_shape": list((x @ x).shape),
    }

    for dimension in MAIN_DIMENSIONS + ("temporal_flickering",):
        importlib.import_module(f"vbench.{dimension}")
        report["imports"][dimension] = "ok"

    for path_string, minimum_size in REQUIRED_WEIGHTS.items():
        path = Path(path_string)
        size = path.stat().st_size if path.is_file() else 0
        if size < minimum_size:
            raise RuntimeError(f"missing or truncated weight: {path} ({size} bytes)")
        report["weights"][path_string] = size

    if not args.skip_dino_load:
        model = torch.hub.load(
            "facebookresearch/dino:main",
            "dino_vitb16",
            source="github",
            trust_repo=True,
            verbose=True,
        )
        report["dino_parameters"] = sum(parameter.numel() for parameter in model.parameters())

    print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())