"""cc12m_data.py — CC12M-resident data path for dist_bed (pod rung).

Trains straight from the pixparse/cc12m-wds tars via a one-time offset
index (no extraction: 1.18TB stays as 2,176 tars). Row order = the
feature-bank concatenation order (features_0000..2175), so targets[idx]
indexing works exactly like the COCO path.

    python tools/cc12m_data.py --build-index     # one-time, parallel scan
"""
import glob
import json
import os
import tarfile

import numpy as np
import torch
from PIL import Image, ImageFile
from torch.utils.data import Dataset

ImageFile.LOAD_TRUNCATED_IMAGES = True

DATA_ROOT = os.environ.get("DIST_DATA_ROOT", "./data")
CC12M_TARS = os.environ.get("CC12M_TARS",
                            os.path.join(DATA_ROOT, "cc12m"))
CC12M_FEATS = os.environ.get("CC12M_FEATS",
                             os.path.join(DATA_ROOT, "cc12m_feats",
                                          "clip_b16_laion2b"))
INDEX = os.path.join(CC12M_TARS, "jpg_offset_index.npz")


def _scan_shard(path):
    """-> (shard_name, {key: (offset, size)}) for .jpg members."""
    out = {}
    with tarfile.open(path, "r") as tar:
        for m in tar:
            if m.isfile() and m.name.endswith(".jpg"):
                out[m.name[:-4]] = (m.offset_data, m.size)
    return os.path.basename(path), out


def build_index(workers=32):
    from multiprocessing import Pool
    tars = sorted(glob.glob(os.path.join(CC12M_TARS, "*.tar")))
    assert tars, f"no tars under {CC12M_TARS}"
    print(f"[index] scanning {len(tars)} tars with {workers} workers")
    shard_names, keys, shard_idx, offs, sizes = [], [], [], [], []
    with Pool(workers) as pool:
        for name, d in pool.imap(_scan_shard, tars, chunksize=4):
            si = len(shard_names)
            shard_names.append(name)
            for k, (o, s) in d.items():
                keys.append(k)
                shard_idx.append(si)
                offs.append(o)
                sizes.append(s)
            if si % 200 == 0:
                print(f"  [{si}] {sum(map(len, [keys]))} keys", flush=True)
    np.savez(INDEX, shards=np.array(shard_names),
             keys=np.array(keys), shard_idx=np.array(shard_idx, np.int32),
             off=np.array(offs, np.int64), size=np.array(sizes, np.int64))
    print(f"[index] {len(keys)} jpgs -> {INDEX}")


def load_cc12m_tower(cfg):
    """-> (keys, fp16 UNNORMALIZED) for an extra tower bank
    (cc12m_feats_{cfg} under CC12M_FEATS_ROOT), concat-cached."""
    root = os.path.join(os.environ.get("CC12M_FEATS_ROOT", "/data"),
                        f"cc12m_feats_{cfg}")
    cache = os.path.join(root, "all_concat.pt")
    if os.path.isfile(cache):
        d = torch.load(cache, map_location="cpu", weights_only=True)
        return d["keys"], d["emb"]
    files = sorted(glob.glob(os.path.join(root, "features_*.pt")))
    assert len(files) == 2176, f"{cfg} bank incomplete: {len(files)}/2176"
    keys, embs = [], []
    for f in files:
        try:
            d = torch.load(f, map_location="cpu", weights_only=True)
        except Exception:
            # pod tower shards carry numpy.str_ keys (offset-index array
            # slices) which weights_only refuses; these are SELF-AUTHORED
            # files on our own disk — trusted load, keys coerced below.
            d = torch.load(f, map_location="cpu", weights_only=False)
        keys += [str(k) for k in d["keys"]]
        embs.append(d["emb"])
    emb = torch.cat(embs)
    torch.save({"keys": keys, "emb": emb}, cache)
    return keys, emb


def load_cc12m_gpa():
    """-> (keys, fp16 GPA targets) built by cc12m_gpa.py."""
    p = os.path.join(os.environ.get("CC12M_FEATS_ROOT", "/data"),
                     "cc12m_gpa.pt")
    d = torch.load(p, map_location="cpu", weights_only=True)
    return d["keys"], d["emb"]


def load_cc12m_targets():
    """-> (keys list, fp16 targets (N,512) UNNORMALIZED) in bank file order.
    Caller normalizes (the bed normalizes per-batch)."""
    cache = os.path.join(os.path.dirname(CC12M_FEATS), "all_concat.pt")
    if os.path.isfile(cache):
        d = torch.load(cache, map_location="cpu", weights_only=True)
        return d["keys"], d["emb"]
    files = sorted(glob.glob(os.path.join(CC12M_FEATS, "features_*.pt")))
    assert len(files) == 2176, f"bank incomplete: {len(files)}/2176"
    keys, embs = [], []
    for f in files:
        d = torch.load(f, map_location="cpu", weights_only=True)
        keys += d["keys"]
        embs.append(d["emb"])
    emb = torch.cat(embs)
    torch.save({"keys": keys, "emb": emb}, cache)
    return keys, emb


class Cc12mImages(Dataset):
    """(row i) -> (STUDENT_TF tensor, i), rows in bank order. Per-worker
    persistent file handles; reads jpg bytes by offset from the tars."""
    def __init__(self, keys, tf):
        z = np.load(INDEX, allow_pickle=False)
        pos = {k: i for i, k in enumerate(z["keys"].tolist())}
        sel = np.array([pos[k] for k in keys], dtype=np.int64)
        self.shards = [os.path.join(CC12M_TARS, s)
                       for s in z["shards"].tolist()]
        self.shard_idx = z["shard_idx"][sel]
        self.off = z["off"][sel]
        self.size = z["size"][sel]
        self.tf = tf

    def __len__(self):
        return len(self.off)

    def _read(self, i):
        # open/seek/read/close per item: with 2,176 shards x N workers,
        # persistent handles exhaust the fd limit (learned the hard way,
        # wave 1). Syscall cost is noise next to jpg decode.
        with open(self.shards[int(self.shard_idx[i])], "rb") as fh:
            fh.seek(int(self.off[i]))
            return fh.read(int(self.size[i]))

    def __getitem__(self, i):
        import io
        im = Image.open(io.BytesIO(self._read(i))).convert("RGB")
        return self.tf(im), i


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--build-index", action="store_true")
    ap.add_argument("--workers", type=int, default=32)
    a = ap.parse_args()
    if a.build_index:
        build_index(a.workers)