# experiment/probing/build_dctrl.py
"""Build a 1k-image D_ctrl jsonl from CC3M, excluding bathroom/toilet keywords.

Outputs (one line per image):
    {"image_id": "<id>", "image_path": "", "caption": "..."}

Note: CC3M is stored as WebDataset tar shards under
  <root>/preproc/CC3M-Dataset/cc3m-wds-train/cc3m-train-NNNN.tar
Each tar contains triplets: <key>.jpg, <key>.json, <key>.txt
The .txt file holds the caption; .jpg bytes are embedded in the tar
(no separate on-disk image files).  image_path is left as "" because
calibrate_tau.py will open images directly from the tar; downstream
scripts that need actual image bytes should use _open_from_tar().

D_ctrl is consumed by calibrate_tau.py and identify_toilet_features.py
(Method C — differential activation needs a clean control distribution).
"""
from __future__ import annotations

import argparse
import json
import os
import random
import re
import tarfile

EXCLUDE_RE = re.compile(r"\b(bathroom|toilets?|lavatory|restroom)\b", re.IGNORECASE)

# Relative subpath inside the root where the wds shards live.
_WDS_SUBDIR = os.path.join("preproc", "CC3M-Dataset", "cc3m-wds-train")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--cc3m_root", required=True, help="Root of CC3M-Dataset")
    p.add_argument("--output", required=True, help="Output jsonl path")
    p.add_argument("--n", type=int, default=1000)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument(
        "--cache_dir",
        default="outputs/feature_ks/d_ctrl_images/",
        help="Directory where extracted JPEG images are cached (one file per image).",
    )
    args = p.parse_args()

    candidates = _scan_cc3m(args.cc3m_root)
    print(f"Scanned {len(candidates)} candidates")
    rng = random.Random(args.seed)
    rng.shuffle(candidates)

    kept: list[dict] = []
    for c in candidates:
        if EXCLUDE_RE.search(c.get("caption", "")):
            continue
        kept.append(c)
        if len(kept) >= args.n:
            break

    if len(kept) < args.n:
        raise RuntimeError(
            f"Only {len(kept)}/{args.n} samples pass the exclusion filter"
        )

    # Extract JPEG bytes for kept records whose image_path is a "tar::key" reference.
    os.makedirs(args.cache_dir, exist_ok=True)
    _extract_images(kept, args.cache_dir)

    out_dir = os.path.dirname(args.output)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    with open(args.output, "w") as f:
        for c in kept:
            f.write(json.dumps(c) + "\n")
    print(f"Wrote {len(kept)} → {args.output}")


def _scan_cc3m(root: str) -> list[dict]:
    """Scan CC3M WebDataset tar shards and return (image_id, image_path, caption) dicts.

    Real layout (differs from the JSON-sidecar convention in the plan):
      <root>/preproc/CC3M-Dataset/cc3m-wds-train/cc3m-train-NNNN.tar
    Each tar entry is one of: <key>.jpg | <key>.json | <key>.txt
    The caption is the text in <key>.txt (stripped).
    image_path encodes the shard+key as "<tar_path>::<key>" so callers
    can re-open the tar to extract the image bytes if needed.
    """
    wds_dir = os.path.join(root, _WDS_SUBDIR)
    if not os.path.isdir(wds_dir):
        raise FileNotFoundError(
            f"Expected WebDataset shard directory not found: {wds_dir}"
        )

    out: list[dict] = []
    for fname in sorted(os.listdir(wds_dir)):
        if not fname.endswith(".tar"):
            continue
        tar_path = os.path.join(wds_dir, fname)
        try:
            with tarfile.open(tar_path, "r") as tf:
                # Collect all .txt members; each gives us one (key, caption) pair.
                for member in tf.getmembers():
                    if not member.name.endswith(".txt"):
                        continue
                    key = member.name[:-4]  # strip ".txt"
                    fobj = tf.extractfile(member)
                    if fobj is None:
                        continue
                    caption = fobj.read().decode("utf-8", errors="replace").strip()
                    out.append(
                        {
                            "image_id": key,
                            # Encode location so callers can extract the JPEG.
                            "image_path": f"{tar_path}::{key}.jpg",
                            "caption": caption,
                        }
                    )
        except tarfile.TarError:
            continue
    return out


def _extract_images(records: list[dict], cache_dir: str) -> None:
    """Extract JPEG bytes from tar shards for each record in *records*.

    Records whose ``image_path`` looks like ``"<tar_path>::<entry>"`` are
    extracted to ``<cache_dir>/<image_id>.jpg`` and the dict is updated in
    place.  Records that already point to a real file are left untouched.
    Extraction is skipped when the destination file already exists (idempotent).

    Tars are opened one at a time in sorted order to avoid repeatedly
    reopening the same archive.
    """
    # Group by tar path so we open each shard at most once.
    from collections import defaultdict

    by_tar: dict[str, list[dict]] = defaultdict(list)
    for c in records:
        ip = c.get("image_path", "")
        if "::" in ip:
            tar_path, _entry = ip.split("::", 1)
            by_tar[tar_path].append(c)

    for tar_path in sorted(by_tar):
        group = by_tar[tar_path]
        # Build a lookup: entry_name -> record
        entry_map = {}
        for c in group:
            _tar_path, entry = c["image_path"].split("::", 1)
            dest = os.path.join(cache_dir, c["image_id"] + ".jpg")
            if os.path.exists(dest):
                c["image_path"] = dest  # already cached
            else:
                entry_map[entry] = (c, dest)

        if not entry_map:
            continue

        try:
            with tarfile.open(tar_path, "r") as tf:
                for member in tf.getmembers():
                    if member.name not in entry_map:
                        continue
                    c, dest = entry_map[member.name]
                    fobj = tf.extractfile(member)
                    if fobj is None:
                        continue
                    with open(dest, "wb") as out_f:
                        out_f.write(fobj.read())
                    c["image_path"] = dest
        except tarfile.TarError as exc:
            raise RuntimeError(f"Failed to extract from {tar_path}: {exc}") from exc


if __name__ == "__main__":
    main()