#!/usr/bin/env python3
"""
tmdb_imdb_map.py — build & consult a cached tmdb_person_id → imdb_id crosswalk.

The gallery (and pipeline output) key actors by TMDB person id but carry no IMDb
`nm…` id, while the X-Ray / MovieNet ground truth keys on IMDb. Rather than join on
fuzzy names, we resolve tmdb→imdb once via TMDB's authoritative
`/person/{id}/external_ids` endpoint and cache the result to JSON. The eval loaders
consult this table to add an exact `imdb:` key alongside each `tmdb:` key.

Table format (JSON): { "<tmdb_person_id>": "nm0000123" | null, ... }
A null means "looked up, TMDB has no IMDb id" — cached so we don't re-query.

Build / refresh the table:
  python scripts/validation/tmdb_imdb_map.py \
      --gallery gallery_arcface_w600k_r50.json \
      --out scripts/validation/tmdb_imdb.json
  # TMDB_API_KEY read from env / .env (via sae_env)

Consult it from code:
  from tmdb_imdb_map import CrosswalkTable
  tbl = CrosswalkTable.load("scripts/validation/tmdb_imdb.json")
  nm = tbl.imdb_for("35467")   # -> "nm..." or None
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

_HERE = Path(__file__).resolve().parent


class CrosswalkTable:
    """Read-only view over the cached tmdb→imdb JSON. Missing file → empty table."""

    def __init__(self, mapping: dict[str, str | None]):
        self._m = mapping

    @classmethod
    def load(cls, path: str | Path) -> "CrosswalkTable":
        p = Path(path)
        if not p.exists():
            return cls({})
        return cls(json.loads(p.read_text()))

    def imdb_for(self, tmdb_id) -> str | None:
        if tmdb_id is None:
            return None
        return self._m.get(str(tmdb_id))

    def __len__(self) -> int:
        return sum(1 for v in self._m.values() if v)


# ── builder ─────────────────────────────────────────────────────────────────

def _external_ids(tmdb_get, tmdb_person_id: str, token: str) -> str | None:
    data = tmdb_get(f"/person/{tmdb_person_id}/external_ids", token)
    imdb = data.get("imdb_id")
    return imdb or None  # normalize "" → None


def build(gallery_path: str, out_path: str, token: str, sleep: float = 0.0) -> None:
    # import the existing TMDB helper (scripts/ is the parent dir)
    sys.path.insert(0, str(_HERE.parent))
    from sae_tmdb import tmdb_get

    with open(gallery_path) as f:
        actors = json.load(f).get("actors", [])
    tmdb_ids = sorted({str(a["tmdb_id"]) for a in actors if a.get("tmdb_id")})
    print(f"[map] gallery tmdb ids: {len(tmdb_ids)}", file=sys.stderr)

    out = Path(out_path)
    existing: dict[str, str | None] = {}
    if out.exists():
        existing = json.loads(out.read_text())
        print(f"[map] resuming from {len(existing)} cached entries", file=sys.stderr)

    todo = [t for t in tmdb_ids if t not in existing]
    print(f"[map] to resolve: {len(todo)}", file=sys.stderr)

    n_ok = n_none = n_err = 0
    for i, tid in enumerate(todo, 1):
        try:
            nm = _external_ids(tmdb_get, tid, token)
            existing[tid] = nm
            n_ok += (nm is not None)
            n_none += (nm is None)
        except Exception as e:  # network/rate-limit/404 — record nothing, keep going
            n_err += 1
            print(f"\n[map] error on tmdb {tid}: {e}", file=sys.stderr)
        if i % 25 == 0 or i == len(todo):
            print(f"\r[map] {i}/{len(todo)} resolved "
                  f"(imdb={n_ok} none={n_none} err={n_err})", end="", file=sys.stderr)
            out.write_text(json.dumps(existing, indent=2))  # checkpoint
        if sleep:
            time.sleep(sleep)

    out.write_text(json.dumps(existing, indent=2))
    print(f"\n[map] wrote {out} — {sum(1 for v in existing.values() if v)} imdb ids",
          file=sys.stderr)


def main():
    # load .env → os.environ (same convention as the gallery builders)
    sys.path.insert(0, str(_HERE.parent))
    try:
        import sae_env  # noqa: F401  (side-effect import)
    except Exception:
        pass

    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--gallery", required=True, help="gallery.json (source of tmdb ids)")
    p.add_argument("--out", default=str(_HERE / "tmdb_imdb.json"),
                   help="output crosswalk JSON (default: scripts/validation/tmdb_imdb.json)")
    p.add_argument("--tmdb-key", default=os.environ.get("TMDB_API_KEY"),
                   help="TMDB v3 API key or v4 read token. Env: TMDB_API_KEY")
    p.add_argument("--sleep", type=float, default=0.0,
                   help="seconds between requests (TMDB has no hard limit; use if throttled)")
    args = p.parse_args()

    if not args.tmdb_key:
        sys.exit("[map] no TMDB key — set TMDB_API_KEY or pass --tmdb-key")

    build(args.gallery, args.out, args.tmdb_key, args.sleep)


if __name__ == "__main__":
    main()