#!/usr/bin/env python3
"""make_jellyfin_gallery.py — build a gallery.h5 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.h5.

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.h5

    # 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.h5 --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.h5

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 json
import sys
from pathlib import Path

import requests

sys.path.insert(0, str(Path(__file__).resolve().parent))
import sae_env  # noqa: F401  — loads .env into os.environ on import
from sae_embed_loader import load_embedder
from sae_gallery import (download_image, download_images, load_gallery_hdf5,
                        save_gallery, wikidata_image_urls)
from sae_jellyfin import actor_jellyfin_id, jf_get, normalize_jellyfin_url
from sae_tmdb import (
    tmdb_person_by_name,
    tmdb_person_for_imdb,
)


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

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")


# ── 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.

    Jellyfin's per-index Primary image endpoint 404s once you run past the
    number of images it has, so stop at the first index that fails rather than
    probing all n.
    """
    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"
        url = base_url.rstrip("/") + f"/Items/{person_id}/Images/Primary/{i}"
        path = download_image(url, out, headers=headers, params={"api_key": api_key})
        if path is None:
            break
        paths.append(path)
    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,
                        fetch_overfetch: float = 1.0
                        ) -> 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)
    # Over-fetch target: see the TMDB block below.
    tmdb_budget = int(images_per_actor * fetch_overfetch)
    image_paths = download_person_images(base_url, api_key, pid, actor_dir, images_per_actor)

    # Need the IMDB id for the TMDB /find lookup, to persist it (--fetch-imdb-ids),
    # and for the Wikidata fallback below. Fetch it once if Jellyfin's own image
    # came up short (so Wikidata can fire) or whenever it's otherwise wanted.
    short = len(image_paths) < images_per_actor
    imdb_id = None
    if fetch_imdb_ids or tmdb_key or short:
        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)

    # Over-fetch from TMDB: near-duplicate stills (the same photo at different
    # crops/resolutions) are dropped after embedding, so downloading exactly
    # images_per_actor would leave the actor short of that many *distinct*
    # embeddings. Pulling extra candidates lets the dedup filter discard
    # duplicates while still reaching the target.
    if len(image_paths) < tmdb_budget and tmdb_urls:
        needed = tmdb_budget - 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_images(tmdb_urls, actor_dir, needed, start_index=len(image_paths))

    # Last resort: a CC-licensed Commons headshot via Wikidata, keyed by the
    # actor's IMDB id. Catches actors TMDB has no usable image for (or that the
    # name search missed entirely).
    if len(image_paths) < tmdb_budget and imdb_id:
        wiki_urls = wikidata_image_urls(imdb_id)
        if wiki_urls:
            needed = tmdb_budget - len(image_paths)
            print(f"  {name}: still short, falling back to Wikidata/Commons "
                  f"({len(wiki_urls)} image(s) available)…", file=sys.stderr)
            image_paths += download_images(wiki_urls, actor_dir, needed,
                                           start_index=len(image_paths))

    return image_paths, imdb_id, tmdb_id


# Default cosine-distance tolerance below which two embeddings of the same
# actor are treated as the same image. Embeddings are L2-normalised by the
# backend, so cosine similarity is a plain dot product and the distance is
# 1 - dot. Expanding an actor's photo set via TMDB frequently returns the same
# still at different crops/resolutions; those embed to nearly identical vectors
# and add gallery size and match cost without adding information.
DEDUP_TOL = 1e-3


def _cosine(a, b) -> float:
    """Cosine similarity of two L2-normalised embeddings."""
    return float(sum(x * y for x, y in zip(a, b)))


def _near_duplicate(emb, existing, tol: float) -> int | None:
    """Index of the first embedding within `tol` cosine distance of `emb`.

    Returns None when `emb` is sufficiently distinct from everything in
    `existing`. tol <= 0 disables the check.
    """
    if tol <= 0:
        return None
    for i, prev in enumerate(existing):
        if 1.0 - _cosine(emb, prev) < tol:
            return i
    return None


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,
                 dedup_tol: float = DEDUP_TOL,
                 max_embeddings: int = 0) -> 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 = []
    n_dup = 0
    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
        emb = res.embedding
        dup = _near_duplicate(emb, embeddings, dedup_tol)
        if dup is not None:
            n_dup += 1
            print(f"  [dup]  {name}/{path.name}: matches {source_images[dup]} "
                  f"(cos={_cosine(emb, embeddings[dup]):.6f}), not stored",
                  file=sys.stderr)
            continue
        embeddings.append(emb)
        source_images.append(path.name)
        # Stop once we have the requested number of *distinct* embeddings; the
        # extra candidates were only fetched to absorb duplicates.
        if max_embeddings and len(embeddings) >= max_embeddings:
            break

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

    dup_note = f" ({n_dup} near-duplicate(s) dropped)" if n_dup else ""
    print(f"  {name}: → {len(embeddings)} embedding(s) stored{dup_note}",
          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,
                   dedup_tol: float = DEDUP_TOL,
                   fetch_overfetch: float = 1.0) -> 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, fetch_overfetch)
    return embed_actor(pid, info, image_paths, imdb_id, tmdb_id, embedder, fetch_imdb_ids,
                       embed_executor, dedup_tol, images_per_actor)


# ── 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,
                   dedup_tol: float = DEDUP_TOL,
                   fetch_overfetch: float = 1.0) -> tuple[dict, list[dict]]:
    actors = collect_actors(base_url, api_key, item_types)

    gallery_actors = []
    todo = []
    n_retry = 0
    n_topup = 0
    for pid, info in actors.items():
        existing = existing_actors.get(pid)
        n_existing = len(existing.get("embeddings", [])) if existing is not None else 0
        # Keep an existing actor as-is only once they're at the image target.
        # Fewer than that (including zero) → re-process to top up / retry, so a
        # merge run can upgrade 1-image actors to images_per_actor without
        # re-embedding actors already at the target.
        if n_existing >= images_per_actor:
            gallery_actors.append(existing)
        else:
            if n_existing > 0:
                n_topup += 1
            elif existing is not None:
                n_retry += 1
            todo.append((pid, info))

    if len(gallery_actors):
        print(f"Skipping {len(gallery_actors)} actor(s) already at {images_per_actor} "
              f"image(s) in existing gallery", file=sys.stderr)
    if n_topup:
        print(f"Topping up {n_topup} actor(s) with fewer than {images_per_actor} "
              f"embedding(s)", 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,
                             dedup_tol, fetch_overfetch): (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.h5 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.h5 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("--dedup-tol", type=float, default=DEDUP_TOL,
                        help=f"Cosine-distance threshold below which a new embedding is treated "
                             f"as a duplicate of one already stored for that actor and dropped "
                             f"(default: {DEDUP_TOL}). TMDB often returns the same still at "
                             f"different crops. Set 0 to keep every embedding.")
    parser.add_argument("--overfetch", type=float, default=2.0,
                        help="Download this multiple of --images-per-actor as candidates, then "
                             "keep the first N that survive dedup (default: 2.0). Raise it for "
                             "actors whose TMDB galleries are mostly duplicates; 1.0 disables "
                             "over-fetching.")
    parser.add_argument("--images-per-actor", type=int, default=10,
                        help="Images to download per actor (default: 10). Jellyfin usually has "
                             "only 1, so the rest come from the TMDB/Wikidata fallbacks; more "
                             "images per actor means more embeddings and more robust matching.")
    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()

    jellyfin_url = normalize_jellyfin_url(args.jellyfin_url)

    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 = load_gallery_hdf5(output)
        for actor in existing.get("actors", []):
            pid = actor_jellyfin_id(actor)
            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,
        dedup_tol=args.dedup_tol,
        fetch_overfetch=args.overfetch,
        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.")

    save_gallery(gallery, missing, output)


if __name__ == "__main__":
    main()