#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import subprocess
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]


def require(path: Path, min_bytes: int = 1) -> None:
    if not path.is_file() or path.stat().st_size < min_bytes:
        raise FileNotFoundError(f"missing or truncated: {path}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checksums", action="store_true")
    parser.add_argument("--materialized", action="store_true")
    args = parser.parse_args()

    required = {
        ROOT / "data/covt/data.json": 100_000_000,
        ROOT / "checkpoints/stage1/checkpoint-32000.bin": 1_000_000_000,
        ROOT / "models/base/Qwen3-VL-4B-Instruct/model-00001-of-00002.safetensors": 4_000_000_000,
        ROOT / "models/base/Qwen3-VL-4B-Instruct/model-00002-of-00002.safetensors": 3_000_000_000,
        ROOT / "models/teachers/sam_vit_h_4b8939.pth": 2_000_000_000,
        ROOT / "models/teachers/VGGT-1B/model.pt": 4_000_000_000,
        ROOT / "models/teachers/table5_baseline.pth": 1_000_000,
        ROOT / "models/teachers/siglip2-large-patch16-256/model.safetensors": 3_000_000_000,
        ROOT / "models/teachers/dinov2-large-processor/preprocessor_config.json": 100,
        ROOT / "cache/torch/hub/checkpoints/dinov2_vitl14_pretrain.pth": 1_000_000_000,
        ROOT / "cache/torch/hub/facebookresearch_dinov2_main/hubconf.py": 1_000,
        ROOT / "environment/unvideo-conda.tar.gz": 1_000_000_000,
    }
    for path, size in required.items():
        require(path, size)

    stats = json.loads((ROOT / "data/covt/stats.json").read_text(encoding="utf-8"))
    shards = sorted((ROOT / "data/covt/shards").glob("images-*.tar"))
    if len(shards) != stats["shard_count"]:
        raise AssertionError(f"shard count {len(shards)} != {stats['shard_count']}")
    for shard in shards:
        require(Path(f"{shard}.sha256"), 64)
    if args.materialized:
        image_count = sum(1 for path in (ROOT / "data/covt/images").iterdir() if path.is_file())
        if image_count != stats["image_count"]:
            raise AssertionError(f"image count {image_count} != {stats['image_count']}")
    if args.checksums:
        subprocess.run(["sha256sum", "-c", "MANIFEST.sha256"], cwd=ROOT, check=True)

    config = ROOT / "configs/stage2_covt.yaml"
    if not config.is_file():
        subprocess.run([str(ROOT / ".env/unvideo/bin/python"), str(ROOT / "scripts/render_config.py")], check=True)
    env = os.environ.copy()
    env["PYTHONPATH"] = ":".join(
        [
            str(ROOT / "code/umm/runtime"),
            str(ROOT / "code/umm"),
            str(ROOT / "third_party/transformers/src"),
            str(ROOT / "third_party/diffusers/src"),
        ]
    )
    env["UMM_STAGE2_CONFIG"] = str(config)
    subprocess.run(
        [str(ROOT / ".env/unvideo/bin/python"), str(ROOT / "code/umm/stage2_covt/validate_setup.py")],
        cwd=ROOT,
        env=env,
        check=True,
    )
    print("UMM Stage-2 bundle validation passed")


if __name__ == "__main__":
    main()