#!/usr/bin/env python3
"""Build stats.json and SHA256SUMS for a TLT12 release directory.

Run this once, after dropping the shards into `shards/`, before publishing.
It is the last step before upload: everything it writes is part of the release.

    python make_release_manifest.py            # scans ./shards
    python make_release_manifest.py --root .   # explicit release root

Writes:
    stats.json   per-split sequence counts, per-category counts, the objects
                 in each split, and whether the splits are object-disjoint
    SHA256SUMS   one `<sha256>  <relative path>` line for EVERY file in the
                 release, in `sha256sum -c` format

SHA256SUMS is written last and covers the whole release, not just the shards,
so it doubles as the download manifest: `download.sh` fetches it first and
uses it as the list of files to pull.

The per-category numbers require reading each shard once (only the small
meta.json members are parsed). Pass --no-stats to skip that.
"""
import argparse
import hashlib
import json
import os
import sys
import tarfile
from collections import Counter, defaultdict

CHUNK = 1 << 20

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

# Never checksummed: SHA256SUMS cannot contain its own hash, and the rest is
# noise that should not ship.
MANIFEST_NAME = "SHA256SUMS"
SKIP_DIRS = {"__pycache__", ".ipynb_checkpoints", ".git"}
SKIP_SUFFIXES = (".pyc", ".pyo")


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 release_files(root):
    """Every file that should ship, as paths relative to `root`, sorted."""
    out = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = sorted(d for d in dirnames
                             if d not in SKIP_DIRS and not d.startswith("."))
        for name in filenames:
            if name == MANIFEST_NAME or name.startswith("."):
                continue
            if name.endswith(SKIP_SUFFIXES):
                continue
            out.append(os.path.relpath(os.path.join(dirpath, name), root))
    return sorted(out)


def find_splits(shards_dir):
    """Return {split_name: dir}. A flat shards/ dir counts as one unnamed split."""
    splits = {}
    for name in sorted(os.listdir(shards_dir)):
        sub = os.path.join(shards_dir, name)
        if os.path.isdir(sub) and any(f.endswith(SHARD_EXT) for f in os.listdir(sub)):
            splits[name] = sub
    if not splits and any(f.endswith(SHARD_EXT) for f in os.listdir(shards_dir)):
        splits["all"] = shards_dir
    return splits


def scan_shard(path):
    """Return (n_sequences, Counter(synset_id), set((synset_id, obj_id)))."""
    n = 0
    per_synset = Counter()
    objects = set()
    with tarfile.open(path, "r|*") as tar:          # streaming, no seeking
        for member in tar:
            if not member.name.endswith("meta.json"):
                continue
            meta = json.loads(tar.extractfile(member).read().decode())
            n += 1
            per_synset[meta["synset_id"]] += 1
            objects.add((meta["synset_id"], meta["obj_id"]))
    return n, per_synset, objects


def build_stats(root, splits, names):
    """Scan every shard and return the stats dict."""
    stats = {"splits": {}}
    all_objects = set()
    objects_per_split = {}

    for split, sdir in sorted(splits.items()):
        tars = sorted(f for f in os.listdir(sdir) if f.endswith(SHARD_EXT))
        n_total = 0
        per_synset = Counter()
        objects = set()
        for t in tars:
            n, c, o = scan_shard(os.path.join(sdir, t))
            n_total += n
            per_synset.update(c)
            objects |= o
            print(f"  scanned {split}/{t}: {n} sequences")
        all_objects |= objects
        objects_per_split[split] = objects

        info_path = os.path.join(sdir, "dataset_info.json")
        info = {}
        if os.path.isfile(info_path):
            with open(info_path) as fh:
                info = json.load(fh)
            if info.get("n_sequences") not in (None, n_total):
                print(f"  WARNING: {split}/dataset_info.json says "
                      f"{info['n_sequences']} sequences, shards contain {n_total}")

        # Which objects landed in this split -- the thing you need to report
        # when the splits are object-disjoint.
        objs_by_synset = defaultdict(list)
        for sid, oid in sorted(objects):
            objs_by_synset[sid].append(oid)

        stats["splits"][split] = {
            "n_shards": len(tars),
            "n_sequences": n_total,
            "n_objects": len(objects),
            "n_frames": info.get("n_frames"),
            "image_size": info.get("image_size"),
            "split_by": info.get("split_by"),
            "bytes": sum(os.path.getsize(os.path.join(sdir, t)) for t in tars),
            "sequences_per_category": {
                sid: {"name": names.get(sid, "?"), "n_sequences": cnt,
                      "n_objects": len(objs_by_synset[sid])}
                for sid, cnt in sorted(per_synset.items())
            },
            "objects": dict(objs_by_synset),
        }

    by_synset = defaultdict(list)
    for sid, oid in sorted(all_objects):
        by_synset[sid].append(oid)
    stats["n_objects_total"] = len(all_objects)
    stats["objects"] = dict(by_synset)

    # Measured from the shards themselves, not taken on trust from the
    # metadata: do any objects appear in more than one split?
    overlap = {}
    split_names = sorted(objects_per_split)
    for i, a in enumerate(split_names):
        for b in split_names[i + 1:]:
            overlap[f"{a}|{b}"] = len(objects_per_split[a] & objects_per_split[b])
    stats["object_overlap_between_splits"] = overlap
    stats["splits_are_object_disjoint"] = all(v == 0 for v in overlap.values())
    return stats


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("--no-stats", action="store_true",
                    help="Skip the per-category shard scan; write SHA256SUMS only")
    args = ap.parse_args()

    root = args.root
    shards_dir = os.path.join(root, "shards")
    if not os.path.isdir(shards_dir):
        sys.exit(f"ERROR: no shards directory at {shards_dir}")

    splits = find_splits(shards_dir)
    if not splits:
        sys.exit(f"ERROR: no shards found under {shards_dir}")
    print(f"Found splits: {', '.join(sorted(splits))}")

    # ── stats first, so SHA256SUMS can cover stats.json ──────────────────────
    stats = None
    if not args.no_stats:
        labels_path = os.path.join(root, "synset_labels.json")
        names = {}
        if os.path.isfile(labels_path):
            with open(labels_path) as fh:
                names = {k: v["name"] for k, v in json.load(fh)["synsets"].items()}

        stats = build_stats(root, splits, names)
        stats_path = os.path.join(root, "stats.json")
        with open(stats_path, "w") as fh:
            json.dump(stats, fh, indent=2)
        print(f"\nWrote {stats_path}")

    # ── checksums for the whole release, written last ────────────────────────
    files = release_files(root)
    lines = []
    total_bytes = 0
    for rel in files:
        p = os.path.join(root, rel)
        total_bytes += os.path.getsize(p)
        lines.append(f"{sha256_file(p)}  {rel}")

    sums_path = os.path.join(root, MANIFEST_NAME)
    with open(sums_path, "w") as fh:
        fh.write("\n".join(lines) + "\n")
    print(f"\nWrote {sums_path}  ({len(lines)} files, "
          f"{total_bytes / 1e9:.2f} GB)")

    # ── summary ──────────────────────────────────────────────────────────────
    if stats:
        for split, s in stats["splits"].items():
            print(f"  {split:<6s} {s['n_sequences']:>7d} sequences  "
                  f"{s['n_shards']:>4d} shards  {s['n_objects']:>4d} objects  "
                  f"{s['bytes'] / 1e9:6.2f} GB")
        if stats["splits_are_object_disjoint"]:
            print("  splits are object-disjoint (no object appears in two splits)")
        else:
            shared = {k: v for k, v in
                      stats["object_overlap_between_splits"].items() if v}
            print("  splits SHARE objects: " +
                  ", ".join(f"{k} {v}" for k, v in shared.items()))

    print("\nNext: upload the release directory, then point download.sh's "
          "BASE_URL at it.")


if __name__ == "__main__":
    main()