#!/usr/bin/env python3
"""Fail fast if a downloaded release is incomplete or has drifted."""

from __future__ import annotations

import hashlib
import json
import os
from collections import Counter
from pathlib import Path

from datasets import load_from_disk


RECIPE = Path(__file__).resolve().parents[1]
REPO = RECIPE.parent


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def read_jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def main() -> None:
    manifest = json.loads((RECIPE / "manifest.json").read_text())
    checked = []
    for relative, expected in manifest["files"].items():
        path = (RECIPE / relative).resolve()
        if not path.exists():
            raise FileNotFoundError(path)
        if path.name == "model.safetensors" and os.environ.get("SKIP_LARGE_HASH") == "1":
            continue
        actual = sha256(path)
        if actual != expected:
            raise RuntimeError(f"SHA256 mismatch for {path}: {actual} != {expected}")
        checked.append(str(path.relative_to(REPO)))

    config = json.loads((REPO / "config.json").read_text())
    if config.get("model_type") != "qwen3":
        raise RuntimeError(f"unexpected tuned-vanilla model_type: {config.get('model_type')}")
    dataset = load_from_disk(str(RECIPE / "data/am_distilled_long_mix"))
    if len(dataset["train"]) != 1024:
        raise RuntimeError(f"expected 1024 training rows, found {len(dataset['train'])}")

    helmet = read_jsonl(RECIPE / "data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl")
    mrcr = read_jsonl(RECIPE / "data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl")
    helmet_counts = Counter(row["config"] for row in helmet)
    mrcr_counts = Counter(row["config"] for row in mrcr)
    if sorted(helmet_counts.values()) != [50] * 5:
        raise RuntimeError(f"unexpected HELMET counts: {helmet_counts}")
    if sorted(mrcr_counts.values()) != [10] * 3:
        raise RuntimeError(f"unexpected MRCR counts: {mrcr_counts}")

    print(json.dumps({
        "status": "ok",
        "checked_sha256": checked,
        "model_type": config["model_type"],
        "training_rows": len(dataset["train"]),
        "helmet_rows": len(helmet),
        "mrcr_rows": len(mrcr),
    }, indent=2))


if __name__ == "__main__":
    main()