#!/usr/bin/env python3
"""make_jellyfin_gallery.py — build a gallery.json spanning an entire Jellyfin library.

Queries the Jellyfin API for every Movie/Series, collects the unique cast
across the whole library, downloads each actor's headshot directly from
Jellyfin (no TMDB key needed), embeds them with the sae_embed module (SCRFD +
ArcFace, loaded once), and writes one global gallery.json.

Because identity_matcher scores every detected face against the whole
gallery, scene_analyze can then recognise any actor in your library in any
film — not just the cast TMDB lists for that one title.

For a single-title run, use scripts/filter_gallery.py afterwards to restrict
matching to that title's credited cast (faster, fewer look-alike mismatches).

Requirements:
    pip install requests Pillow

Usage:
    python scripts/make_jellyfin_gallery.py \\
        --jellyfin-url http://jellyfin.local:8096 \\
        --api-key YOUR_API_KEY \\
        --output gallery.json

    # Re-run later to pick up newly added titles without re-embedding
    # actors already in the gallery:
    python scripts/make_jellyfin_gallery.py \\
        --jellyfin-url http://jellyfin.local:8096 \\
        --api-key YOUR_API_KEY \\
        --output gallery.json --merge

    # Fall back to TMDB profile images for actors with no usable Jellyfin image:
    python scripts/make_jellyfin_gallery.py \\
        --jellyfin-url http://jellyfin.local:8096 \\
        --api-key YOUR_API_KEY \\
        --tmdb-key YOUR_TMDB_KEY \\
        --output gallery.json

Get a Jellyfin API key from Dashboard → Advanced → API Keys.
Get a free TMDB API key at: https://www.themoviedb.org/settings/api
"""

import argparse
import concurrent.futures
import os
import io
import json
import sys
from pathlib import Path

import requests
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sae_embed_loader import load_embedder

TMDB_BASE = "https://api.themoviedb.org/3"
TMDB_IMG = "https://image.tmdb.org/t/p/original"


# ── Jellyfin API helpers ────────────────────────────────────────────────────

def jf_get(base_url: str, api_key: str, path: str, **params) -> dict:
    url = base_url.rstrip("/") + path
    headers = {"X-Emby-Token": api_key, "Accept": "application/json"}
    r = requests.get(url, params=params, headers=headers, timeout=30)
    r.raise_for_status()
    return r.json()


def fetch_library_items(base_url: str, api_key: str, item_types: list[str]):
    """Yield every Movie/Series item dict (with its People field)."""
    start = 0
    limit = 100
    while True:
        data = jf_get(
            base_url, api_key, "/Items",
            Recursive="true",
            IncludeItemTypes=",".join(item_types),
            Fields="People",
            StartIndex=start,
            Limit=limit,
        )
        items = data.get("Items", [])
        for item in items:
            yield item
        start += limit
        if start >= data.get("TotalRecordCount", 0) or not items:
            break


def collect_actors(base_url: str, api_key: str, item_types: list[str]) -> dict:
    """Return {person_id: {name, primary_image_tag, appearances: [...]}}."""
    actors: dict[str, dict] = {}
    n_items = 0
    for item in fetch_library_items(base_url, api_key, item_types):
        n_items += 1
        for person in item.get("People", []):
            if person.get("Type") != "Actor":
                continue
            pid = person["Id"]
            entry = actors.setdefault(pid, {
                "name": person["Name"],
                "primary_image_tag": person.get("PrimaryImageTag"),
                "appearances": [],
            })
            if not entry["primary_image_tag"]:
                entry["primary_image_tag"] = person.get("PrimaryImageTag")
            entry["appearances"].append({
                "item_id": item["Id"],
                "title": item.get("Name", ""),
                "type": item.get("Type", ""),
            })
        if n_items % 50 == 0:
            print(f"  scanned {n_items} items, {len(actors)} unique actors so far…", file=sys.stderr)
    print(f"Scanned {n_items} items, found {len(actors)} unique actors", file=sys.stderr)
    return actors


def fetch_imdb_id(base_url: str, api_key: str, person_id: str) -> str | None:
    try:
        data = jf_get(base_url, api_key, f"/Items/{person_id}", Fields="ProviderIds")
    except requests.RequestException:
        return None
    return data.get("ProviderIds", {}).get("Imdb")


# ── TMDB fallback (for actors with no usable Jellyfin image) ────────────────

def tmdb_get(path: str, key: str, **params) -> dict:
    url = TMDB_BASE + path
    if key.startswith("eyJ"):
        headers = {"Authorization": f"Bearer {key}", "Accept": "application/json"}
        r = requests.get(url, params=params, headers=headers, timeout=10)
    else:
        params["api_key"] = key
        r = requests.get(url, params=params, headers={"Accept": "application/json"}, timeout=10)
    r.raise_for_status()
    return r.json()


def tmdb_person_images(tmdb_person_id: str, tmdb_key: str) -> list[str]:
    images_data = tmdb_get(f"/person/{tmdb_person_id}/images", tmdb_key)
    return [TMDB_IMG + p["file_path"] for p in images_data.get("profiles", []) if p.get("file_path")]


def tmdb_person_for_imdb(imdb_id: str, tmdb_key: str) -> tuple[str | None, list[str]]:
    """Return (tmdb_person_id, profile_image_urls) for the TMDB person matching this IMDB person id."""
    data = tmdb_get(f"/find/{imdb_id}", tmdb_key, external_source="imdb_id")
    people = data.get("person_results", [])
    if not people:
        return None, []
    tmdb_person_id = str(people[0]["id"])
    return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key)


def tmdb_person_by_name(name: str, tmdb_key: str) -> tuple[str | None, list[str]]:
    """Return (tmdb_person_id, profile_image_urls) for the best name match on TMDB.

    Used when Jellyfin has no IMDB ProviderId for this person (the common case —
    Jellyfin rarely populates ProviderIds on Person items), so /find/{imdb_id}
    isn't an option. /search/person is sorted by popularity; take the top hit.
    """
    data = tmdb_get("/search/person", tmdb_key, query=name)
    people = data.get("results", [])
    if not people:
        return None, []
    tmdb_person_id = str(people[0]["id"])
    return tmdb_person_id, tmdb_person_images(tmdb_person_id, tmdb_key)


def download_urls(urls: list[str], dest_dir: Path, n: int, start_index: int = 0) -> list[Path]:
    """Download up to n images from urls into dest_dir, numbered from start_index."""
    dest_dir.mkdir(parents=True, exist_ok=True)
    paths = []
    for i, url in enumerate(urls[:n]):
        out = dest_dir / f"{start_index + i:02d}.jpg"
        if out.exists() and out.stat().st_size > 1024:
            paths.append(out)
            continue
        try:
            r = requests.get(url, timeout=15)
            r.raise_for_status()
            Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG")
            paths.append(out)
        except Exception as e:
            print(f"  [warn] TMDB image download failed: {url}: {e}", file=sys.stderr)
    return paths


# ── Image download ──────────────────────────────────────────────────────────

def download_person_images(base_url: str, api_key: str, person_id: str,
                            dest_dir: Path, n: int) -> list[Path]:
    """Download up to n images for a Jellyfin person item into dest_dir."""
    dest_dir.mkdir(parents=True, exist_ok=True)
    headers = {"X-Emby-Token": api_key}
    paths = []
    for i in range(n):
        out = dest_dir / f"{i:02d}.jpg"
        if out.exists() and out.stat().st_size > 1024:
            paths.append(out)
            continue
        url = base_url.rstrip("/") + f"/Items/{person_id}/Images/Primary/{i}"
        try:
            r = requests.get(url, headers=headers, params={"api_key": api_key}, timeout=15)
            if r.status_code == 404:
                break
            r.raise_for_status()
            ctype = r.headers.get("Content-Type", "")
            if not ctype.startswith("image/"):
                print(f"  [warn] unexpected response for {person_id} index {i}: "
                      f"status={r.status_code} content-type={ctype!r} "
                      f"len={len(r.content)} body={r.content[:200]!r}", file=sys.stderr)
                break
            Image.open(io.BytesIO(r.content)).convert("RGB").save(out, "JPEG")
            paths.append(out)
        except Exception as e:
            print(f"  [warn] image download failed for {person_id} index {i}: {e}", file=sys.stderr)
            break
    return paths


# ── Per-actor pipeline ───────────────────────────────────────────────────────

def fetch_actor_images(base_url: str, api_key: str, pid: str, info: dict,
                        images_per_actor: int, actor_dir: Path,
                        fetch_imdb_ids: bool, tmdb_key: str | None
                        ) -> tuple[list[Path], str | None, str | None]:
    """Network-bound: download Jellyfin image(s), then fall back to TMDB if short."""
    name = info["name"]
    print(f"{name} ({pid}) — in {len(info['appearances'])} title(s)", file=sys.stderr)
    image_paths = download_person_images(base_url, api_key, pid, actor_dir, images_per_actor)

    imdb_id = None
    if fetch_imdb_ids or tmdb_key:
        imdb_id = fetch_imdb_id(base_url, api_key, pid)

    # Resolve the TMDB person id whenever possible — independent of whether
    # Jellyfin already gave us enough images, so tmdb_id is always populated
    # when --tmdb-key is set. Jellyfin Person items rarely have an IMDB
    # ProviderId, so fall back to a name search when imdb_id is unknown.
    tmdb_id = None
    tmdb_urls: list[str] = []
    if tmdb_key:
        try:
            if imdb_id:
                tmdb_id, tmdb_urls = tmdb_person_for_imdb(imdb_id, tmdb_key)
            if tmdb_id is None:
                tmdb_id, tmdb_urls = tmdb_person_by_name(name, tmdb_key)
        except requests.RequestException as e:
            print(f"  [warn] {name}: TMDB lookup failed: {e}", file=sys.stderr)

    if len(image_paths) < images_per_actor and tmdb_urls:
        needed = images_per_actor - len(image_paths)
        print(f"  {name}: Jellyfin image missing/incomplete, falling back to TMDB "
              f"({len(tmdb_urls)} image(s) available)…", file=sys.stderr)
        image_paths += download_urls(tmdb_urls, actor_dir, needed, start_index=len(image_paths))

    return image_paths, imdb_id, tmdb_id


def embed_actor(pid: str, info: dict, image_paths: list[Path],
                 imdb_id: str | None, tmdb_id: str | None, embedder, fetch_imdb_ids: bool,
                 embed_executor: concurrent.futures.ThreadPoolExecutor) -> tuple[dict | None, str | None]:
    """GPU-bound: run sae_embed, always on embed_executor's single dedicated thread.

    onnxruntime's CUDA EP / cudnn_frontend execution plans are not safe to run
    from arbitrary threads — calling Run() from a different OS thread than the
    one that last used the session corrupts the cudnn graph (CUDNN_FE failure
    11 / CUDNN_BACKEND_API_FAILED). Routing every embed() call through one
    persistent thread avoids that regardless of how many worker threads are
    fetching images concurrently.
    """
    name = info["name"]
    if not image_paths:
        print(f"  {name}: no image available, skipping", file=sys.stderr)
        return None, "no image available"

    embeddings = []
    source_images = []
    print(f"  {name}: embedding {len(image_paths)} image(s)…", file=sys.stderr)
    for path in image_paths:
        res = embed_executor.submit(embedder.embed, str(path)).result()
        if not res.ok:
            print(f"  [skip] {name}/{path.name}: {res.error}", file=sys.stderr)
            continue
        embeddings.append(res.embedding)
        source_images.append(path.name)

    if not embeddings:
        print(f"  {name}: no valid embeddings, skipping actor", file=sys.stderr)
        return None, "no valid embeddings"

    print(f"  {name}: → {len(embeddings)} embedding(s) stored", file=sys.stderr)
    return {
        "imdb_id": imdb_id if (fetch_imdb_ids and imdb_id) else "",
        "tmdb_id": tmdb_id or "",
        "jellyfin_id": pid,
        "name": name,
        "source_images": source_images,
        "embeddings": embeddings,
        "appearances": info["appearances"],
    }, None


def process_actor(pid: str, info: dict, base_url: str, api_key: str,
                   embedder, images_per_actor: int, image_root: Path,
                   fetch_imdb_ids: bool, tmdb_key: str | None,
                   embed_executor: concurrent.futures.ThreadPoolExecutor) -> tuple[dict | None, str | None]:
    safe_name = info["name"].replace(" ", "_")
    actor_dir = image_root / f"{pid}_{safe_name}"
    image_paths, imdb_id, tmdb_id = fetch_actor_images(
        base_url, api_key, pid, info, images_per_actor, actor_dir, fetch_imdb_ids, tmdb_key)
    return embed_actor(pid, info, image_paths, imdb_id, tmdb_id, embedder, fetch_imdb_ids, embed_executor)


# ── Gallery assembly ─────────────────────────────────────────────────────────

def build_gallery(base_url: str, api_key: str, embedder, item_types: list[str],
                   images_per_actor: int, image_root: Path,
                   fetch_imdb_ids: bool, existing_actors: dict,
                   tmdb_key: str | None = None, workers: int = 8) -> tuple[dict, list[dict]]:
    actors = collect_actors(base_url, api_key, item_types)

    gallery_actors = []
    todo = []
    n_retry = 0
    for pid, info in actors.items():
        existing = existing_actors.get(pid)
        if existing is not None and existing.get("embeddings"):
            gallery_actors.append(existing)
        else:
            if existing is not None:
                n_retry += 1
            todo.append((pid, info))

    if len(gallery_actors):
        print(f"Skipping {len(gallery_actors)} actor(s) already present in existing gallery", file=sys.stderr)
    if n_retry:
        print(f"Retrying {n_retry} actor(s) with no embeddings in existing gallery", file=sys.stderr)
    print(f"Processing {len(todo)} new actor(s) with {workers} worker(s)…", file=sys.stderr)

    missing = []
    n_done = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor, \
         concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="embed") as embed_executor:
        futures = {
            executor.submit(process_actor, pid, info, base_url, api_key, embedder,
                             images_per_actor, image_root,
                             fetch_imdb_ids, tmdb_key, embed_executor): (pid, info["name"])
            for pid, info in todo
        }
        for future in concurrent.futures.as_completed(futures):
            n_done += 1
            pid, name = futures[future]
            try:
                actor, reason = future.result()
            except Exception as e:
                print(f"  [error] {name}: {e}", file=sys.stderr)
                missing.append({"jellyfin_id": pid, "name": name, "reason": str(e)})
                continue
            if actor:
                gallery_actors.append(actor)
            else:
                missing.append({"jellyfin_id": pid, "name": name, "reason": reason})
            if n_done % 25 == 0:
                print(f"  progress: {n_done}/{len(todo)} actors processed", file=sys.stderr)

    return {"actors": gallery_actors}, missing


# ── Entry point ───────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Build a gallery.json spanning an entire Jellyfin library",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--jellyfin-url", default=os.environ.get("JELLYFIN_URL"),
                        required=not os.environ.get("JELLYFIN_URL"),
                        help="Jellyfin base server URL only, e.g. http://jellyfin.local:8096 "
                             "(no /Items or other API path). Env: JELLYFIN_URL")
    parser.add_argument("--api-key", default=os.environ.get("JELLYFIN_API_KEY"),
                        required=not os.environ.get("JELLYFIN_API_KEY"),
                        help="Jellyfin API key (Dashboard → Advanced → API Keys). Env: JELLYFIN_API_KEY")
    parser.add_argument("--output", required=True, help="Output gallery.json path")
    parser.add_argument("--item-types", default="Movie,Series",
                        help="Comma-separated Jellyfin item types to scan (default: Movie,Series)")
    parser.add_argument("--build-dir", default="build",
                        help="Build directory containing the sae_embed module (default: build)")
    parser.add_argument("--models-dir", default="models",
                        help="Directory containing ONNX models (default: models/)")
    parser.add_argument("--arcface", default=None,
                        help="Path to ArcFace ONNX model (overrides --models-dir selection)")
    parser.add_argument("--images-per-actor", type=int, default=1,
                        help="Images to download per actor (default: 1 — Jellyfin usually caches one)")
    parser.add_argument("--image-dir", default=None,
                        help="Where to store downloaded images (default: <output_dir>/images)")
    parser.add_argument("--fetch-imdb-ids", action="store_true",
                        help="Resolve each actor's real IMDB id via Jellyfin ProviderIds "
                             "(one extra API call per new actor; otherwise imdb_id is left empty)")
    parser.add_argument("--tmdb-key", default=os.environ.get("TMDB_API_KEY"),
                        help="TMDB API key/bearer token. If set, actors with no usable "
                             "Jellyfin image fall back to TMDB profile images (looked up "
                             "via the actor's IMDB id, requires one extra Jellyfin call per actor). "
                             "Env: TMDB_API_KEY")
    parser.add_argument("--merge", action="store_true",
                        help="If --output already exists, keep its actors and only embed "
                             "actors not already present (matched by jellyfin_id)")
    parser.add_argument("--workers", type=int, default=8,
                        help="Concurrent worker threads for Jellyfin/TMDB lookups and "
                             "image downloads (default: 8). Embedding itself always runs "
                             "on a single dedicated thread regardless of this value.")
    args = parser.parse_args()

    # Defend against a base URL that accidentally includes an API path,
    # e.g. "https://host/Items?" — strip any query string and trailing
    # /Items so we don't build double-nested, query-mangled URLs.
    jellyfin_url = args.jellyfin_url.split("?", 1)[0].rstrip("/")
    if jellyfin_url.endswith("/Items"):
        jellyfin_url = jellyfin_url[: -len("/Items")]

    output = Path(args.output)
    image_root = Path(args.image_dir) if args.image_dir else output.parent / "images"
    item_types = [t.strip() for t in args.item_types.split(",") if t.strip()]

    embedder = load_embedder(args.build_dir, args.models_dir, args.arcface)

    existing_actors = {}
    if args.merge and output.is_file():
        existing = json.loads(output.read_text())
        for actor in existing.get("actors", []):
            # older galleries used "jellyfin_person_id"
            pid = actor.get("jellyfin_id") or actor.get("jellyfin_person_id")
            if pid:
                existing_actors[pid] = actor
        print(f"Loaded {len(existing_actors)} actor(s) from existing gallery for merge", file=sys.stderr)

    gallery, missing = build_gallery(
        base_url=jellyfin_url,
        api_key=args.api_key,
        embedder=embedder,
        item_types=item_types,
        images_per_actor=args.images_per_actor,
        image_root=image_root,
        fetch_imdb_ids=args.fetch_imdb_ids,
        existing_actors=existing_actors,
        tmdb_key=args.tmdb_key,
        workers=args.workers,
    )

    n_actors = len(gallery["actors"])
    n_embeddings = sum(len(a["embeddings"]) for a in gallery["actors"])
    print(f"\nGallery: {n_actors} actors, {n_embeddings} total embeddings", file=sys.stderr)

    if n_actors == 0:
        sys.exit("No actors could be processed — check Jellyfin URL/API key and models.")

    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(gallery, indent=2) + "\n")
    print(f"Saved: {output}", file=sys.stderr)

    if missing:
        missing_path = output.with_name(output.stem + ".missing_images.json")
        missing_path.write_text(json.dumps(missing, indent=2) + "\n")
        print(f"{len(missing)} actor(s) need images — see {missing_path}", file=sys.stderr)


if __name__ == "__main__":
    main()