#!/usr/bin/env python3
"""Build a binary-label manifest for DAD clips → drop-in input for
make_belief_cache_v2.py in binary pretraining mode (Stage K0).

Output (matches make_policy_labels.py schema):
{
  "samples": [
    {
      "video_id":      "dad_positive_training_000001",
      "source":        "dad",
      "category":      "ego_positive" | "safe_neg",
      "source_dir":    "data/pretrain_v2/dad_frames/positive/000001",
      "frame_indices": [last N frame ids, tail-biased],
      "tta_raw":       -1.0,
      "action_label":  2 (positive) | 0 (negative),
      "ce_weight":     1.0,
      "metadata":      {"fps": 20.0, "n_frames": 100, "split": "training"|"testing"|"single"}
    }, ...
  ],
  "label_counts": {"ALERT": n_pos, "SILENT": n_neg},
  "excluded":     {"missing_frames": k}
}

DAD has no TTA; tta_raw=-1.0 and action_label is a 0/2 binary mapping that
downstream binary heads reduce to {0, 1}. 3-class PolicyHead should NOT
be trained on this manifest.
"""
from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path
from typing import List, Tuple


def _build_frame_indices(n_frames: int, window: int) -> List[int]:
    last = n_frames - 1
    first = max(0, last - window + 1)
    return list(range(first, last + 1))


def _discover_splits(root: Path) -> List[Tuple[str, Path]]:
    """Return [(tag, dir)] for each split dir. Falls back to a no-split layout."""
    split_names = ["training", "testing"]
    found = [(s, root / s) for s in split_names if (root / s).exists()]
    if found:
        return found
    if (root / "positive").exists() or (root / "negative").exists():
        return [("single", root)]
    return []


def _scan(split_tag: str, split_dir: Path, frame_window: int,
          excluded: Counter) -> List[dict]:
    out: List[dict] = []
    for cls_name, label in [("positive", 2), ("negative", 0)]:
        cls_dir = split_dir / cls_name
        if not cls_dir.exists():
            continue
        for clip_dir in sorted(cls_dir.iterdir()):
            if not clip_dir.is_dir():
                continue
            ann_path = clip_dir / "annotation.json"
            if ann_path.exists():
                ann = json.load(open(ann_path))
                n_frames = int(ann.get("n_frames", 0))
                fps = float(ann.get("fps", 20.0))
            else:
                n_frames = len(list(clip_dir.glob("*.jpg")))
                fps = 20.0
            if n_frames <= 0:
                excluded["missing_frames"] += 1
                continue
            frame_idx = _build_frame_indices(n_frames, frame_window)
            tail = frame_idx[-1]
            if not (clip_dir / f"{tail:03d}.jpg").exists():
                excluded["missing_frames"] += 1
                continue
            out.append({
                "video_id":      f"dad_{cls_name}_{split_tag}_{clip_dir.name}",
                "source":        "dad",
                "category":      "ego_positive" if label == 2 else "safe_neg",
                "source_dir":    str(clip_dir),
                "frame_indices": frame_idx,
                "tta_raw":       -1.0,
                "action_label":  label,
                "ce_weight":     1.0,
                "metadata": {
                    "fps":      fps,
                    "n_frames": n_frames,
                    "split":    split_tag,
                },
            })
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames_root",
                    default="data/pretrain_v2/dad_frames")
    ap.add_argument("--out",
                    default="data/policy_labels/dad_binary.json")
    ap.add_argument("--frame_window", type=int, default=60,
                    help="tail-biased window of frame ids baked into manifest "
                         "(must be ≥ n_frames used in cache build)")
    ap.add_argument("--exclude_testing", action="store_true",
                    help="keep only DAD training split (default: include both)")
    args = ap.parse_args()

    root = Path(args.frames_root)
    if not root.exists():
        raise SystemExit(f"frames_root {root} does not exist")

    splits = _discover_splits(root)
    if not splits:
        raise SystemExit(f"no split/class directories under {root}")
    if args.exclude_testing:
        splits = [s for s in splits if s[0] != "testing"]

    samples: List[dict] = []
    excluded: Counter = Counter()
    for tag, d in splits:
        samples.extend(_scan(tag, d, args.frame_window, excluded))

    label_counts = Counter("ALERT" if s["action_label"] == 2 else "SILENT"
                           for s in samples)
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as f:
        json.dump({
            "samples":      samples,
            "label_counts": dict(label_counts),
            "excluded":     dict(excluded),
        }, f)

    by_split = Counter(s["metadata"]["split"] for s in samples)
    print(f"[dad_manifest] {len(samples)} clips  "
          f"labels={dict(label_counts)}  splits={dict(by_split)}  "
          f"excluded={dict(excluded)}  → {out_path}")


if __name__ == "__main__":
    main()