#!/usr/bin/env python3
"""Check that a downloaded copy of TLT12 is complete and intact.

    python verify_release.py             # full check, verifies every checksum
    python verify_release.py --quick     # skip checksums, structure only

Standard library only -- run this before installing anything else. A truncated
or partially-synced download is by far the most common problem with a
multi-gigabyte dataset, and it usually shows up much later as a confusing
decode error in the middle of training.

Exit status is 0 when everything checks out, 1 otherwise.
"""
import argparse
import hashlib
import json
import os
import sys
import tarfile

CHUNK = 1 << 20

# Shards may be plain or gzipped; webdataset reads both.
SHARD_EXT = (".tar", ".tar.gz")


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(CHUNK), b""):
            h.update(block)
    return h.hexdigest()


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=os.path.dirname(os.path.abspath(__file__)),
                    help="Release root (default: this script's directory)")
    ap.add_argument("--quick", action="store_true",
                    help="Skip checksum verification (structure checks only)")
    args = ap.parse_args()

    problems = []
    root = args.root

    def fail(msg):
        problems.append(msg)
        print(f"  FAIL  {msg}")

    def ok(msg):
        print(f"  ok    {msg}")

    # ── 1. structure ─────────────────────────────────────────────────────────
    print("Structure:")
    shards_dir = os.path.join(root, "shards")
    if not os.path.isdir(shards_dir):
        fail("shards/ directory is missing")
        print("\nCannot continue without shards/.")
        return 1

    splits = {}
    for name in sorted(os.listdir(shards_dir)):
        sub = os.path.join(shards_dir, name)
        if os.path.isdir(sub):
            tars = sorted(f for f in os.listdir(sub) if f.endswith(SHARD_EXT))
            if tars:
                splits[name] = (sub, tars)
    if not splits:
        tars = sorted(f for f in os.listdir(shards_dir) if f.endswith(SHARD_EXT))
        if tars:
            splits["all"] = (shards_dir, tars)

    if not splits:
        fail("no .tar shards found under shards/")
        return 1

    for split, (sdir, tars) in splits.items():
        ok(f"split '{split}': {len(tars)} shards")
        if not os.path.isfile(os.path.join(sdir, "dataset_info.json")):
            fail(f"split '{split}': dataset_info.json is missing")

    # ── 2. checksums ─────────────────────────────────────────────────────────
    sums_path = os.path.join(root, "SHA256SUMS")
    if args.quick:
        print("\nChecksums: skipped (--quick)")
    elif not os.path.isfile(sums_path):
        print("\nChecksums: SHA256SUMS not found -- skipping")
    else:
        print("\nChecksums:")
        with open(sums_path) as fh:
            entries = [ln.strip().split("  ", 1) for ln in fh if ln.strip()]
        n_good = 0
        for i, (expect, rel) in enumerate(entries, 1):
            p = os.path.join(root, rel)
            if not os.path.isfile(p):
                fail(f"{rel}: missing")
                continue
            if sha256_file(p) != expect:
                fail(f"{rel}: checksum mismatch (re-download this file)")
                continue
            n_good += 1
            if i % 10 == 0 or i == len(entries):
                print(f"  ok    {n_good}/{len(entries)} files verified")

    # ── 3. sequence counts ───────────────────────────────────────────────────
    print("\nContents:")
    stats_path = os.path.join(root, "stats.json")
    stats = {}
    if os.path.isfile(stats_path):
        with open(stats_path) as fh:
            stats = json.load(fh).get("splits", {})

    for split, (sdir, tars) in splits.items():
        n = 0
        for t in tars:
            try:
                with tarfile.open(os.path.join(sdir, t), "r|*") as tar:
                    n += sum(1 for m in tar if m.name.endswith("meta.json"))
            except tarfile.TarError as e:
                fail(f"{split}/{t}: unreadable tar ({e})")
        expected = None
        info_path = os.path.join(sdir, "dataset_info.json")
        if os.path.isfile(info_path):
            with open(info_path) as fh:
                expected = json.load(fh).get("n_sequences")
        if expected is None:
            expected = stats.get(split, {}).get("n_sequences")

        if expected is None:
            ok(f"split '{split}': {n} sequences (no reference count to compare)")
        elif n == expected:
            ok(f"split '{split}': {n} sequences, matches metadata")
        else:
            fail(f"split '{split}': {n} sequences, metadata expects {expected}")

    # ── 4. one sample decodes ────────────────────────────────────────────────
    print("\nSample decode:")
    split, (sdir, tars) = next(iter(splits.items()))
    try:
        with tarfile.open(os.path.join(sdir, tars[0]), "r|*") as tar:
            keys = []
            first_key = None
            for m in tar:
                key = m.name.rsplit(".", 1)[0].split(".")[0]
                if first_key is None:
                    first_key = key
                if key != first_key:
                    break
                keys.append(m.name)
        n_jpg = sum(1 for k in keys if k.endswith(".jpg"))
        have = {os.path.basename(k).split(".", 1)[1] for k in keys}
        ok(f"first sample carries {n_jpg} frames")
        for need in ("latents.npy", "base_latent.npy", "velocities.npy", "meta.json"):
            if need not in have:
                fail(f"first sample is missing {need}")
    except tarfile.TarError as e:
        fail(f"could not read {split}/{tars[0]}: {e}")

    # ── verdict ──────────────────────────────────────────────────────────────
    if problems:
        print(f"\n{len(problems)} problem(s) found:")
        for p in problems:
            print(f"  - {p}")
        return 1
    print("\nAll checks passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())