#!/usr/bin/python3
"""Recommend movies based on the whole Radarr library, via the local LLM stack.

  movie_recs [-n COUNT] [-m MODEL] [--json] [--prompt-only]

Talks to LiteLLM (OpenAI-compatible) rather than Ollama, which was removed
2026-08-28. Ollama and llama-server both wanted the whole 3060, and with
gpt-oss-20b resident at ~11.3 GB of 11.9 GB there was under 1 GB left, so
every Ollama request OOM'd. llama-swap now arbitrates a single GPU owner.

The model is the same weights as before (gemma3 12B instruct, Q4_K_M), just
served by llama.cpp instead, so recommendations keep their prior character.

Reads every movie in Radarr, asks the model for titles that are NOT already
in the library, and filters the answer against the library again (models
reliably recommend things you already own).
"""
import argparse
import json
import os
import re
from collections import Counter, defaultdict
import sys
import unicodedata
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

RADARR_CONFIG = "/home/connor/config/radarr/config.xml"
# Router DNS override points this at the LAN, so it never leaves the network.
RADARR_URL = "https://radarr.rcjohnstone.com"
# LiteLLM's loopback publish -- no reason to make it a network round trip.
# (Caddy at llm.rcjohnstone.com is the LAN/OpenVPN path; this runs on the host.)
LITELLM_URL = "http://127.0.0.1:4000"
# A LiteLLM virtual key scoped to just this model, NOT the master key: this
# runs unattended from a timer, and a scoped key also gives the spend log a
# per-consumer attribution instead of lumping every caller under "master".
# 0600 file, same pattern as movie_recs_notify's ntfy credential.
LITELLM_ENV = os.path.expanduser("~/.config/litellm/movie-recs.env")

SCHEMA = {
    "type": "object",
    "properties": {
        "recommendations": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "year": {"type": "integer"},
                    "anchors": {
                        "type": "array",
                        "items": {"type": "string"},
                        "minItems": 3,
                        "maxItems": 5,
                    },
                    "pattern": {"type": "string"},
                    "reason": {"type": "string"},
                },
                "required": ["title", "year", "anchors", "pattern", "reason"],
            },
        }
    },
    "required": ["recommendations"],
}


def norm(t):
    """Loose title key: fold accents, drop articles/punctuation/case."""
    t = unicodedata.normalize("NFKD", t)
    t = "".join(c for c in t if not unicodedata.combining(c)).lower()
    t = re.sub(r"[^a-z0-9 ]+", "", t)
    t = re.sub(r"^(the|a|an) ", "", t).strip()
    return re.sub(r"\s+", " ", t)


def get_json(url, data=None, timeout=600, headers=None):
    req = urllib.request.Request(
        url,
        data=json.dumps(data).encode() if data is not None else None,
        headers={"Content-Type": "application/json", **(headers or {})},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode())


def litellm_key():
    """Read the virtual key, env first so a manual run can override."""
    key = os.environ.get("LITELLM_KEY")
    if key:
        return key
    try:
        with open(LITELLM_ENV) as fh:
            for line in fh:
                line = line.strip()
                if line.startswith("#") or "=" not in line:
                    continue
                k, v = line.split("=", 1)
                if k.strip() == "LITELLM_KEY":
                    return v.strip()
    except OSError:
        pass
    sys.exit("No LiteLLM key: set LITELLM_KEY or write it to %s" % LITELLM_ENV)


def api_key():
    return ET.parse(RADARR_CONFIG).getroot().findtext("ApiKey")


def library(key):
    movies = get_json("%s/api/v3/movie?apikey=%s" % (RADARR_URL, key), timeout=60)
    out = []
    for m in movies:
        out.append({
            "title": m.get("title", "?"),
            "year": m.get("year"),
            "genres": m.get("genres") or [],
        })
    return sorted(out, key=lambda x: (x["year"] or 0))


def verify(key, title, year):
    """Resolve a suggestion against TMDB via Radarr's own lookup.

    Models at this size invent plausible-sounding films, so a title the
    model asserts is only accepted when TMDB agrees on both the name and
    (within a year) the release date. Returns the real record or None.
    """
    # Models often glue the year onto the title ("Gone Girl (2014)"), which
    # would never compare equal to TMDB's bare title. Split it back out.
    m = re.search(r"\((\d{4})\)\s*$", title)
    if m:
        year = year or int(m.group(1))
        title = title[:m.start()].strip()

    url = "%s/api/v3/movie/lookup?apikey=%s&term=%s" % (
        RADARR_URL, key, urllib.parse.quote(title))
    try:
        hits = get_json(url, timeout=30)
    except Exception:
        return None
    want = norm(title)
    matches = [h for h in hits[:10] if norm(h.get("title", "")) == want]
    if not matches and year:
        # Films are often listed under a longer official title ("The French
        # Dispatch" vs "The French Dispatch of the Liberty, Kansas Evening
        # Sun"). Accept a prefix, but only with the year agreeing, so this
        # stays tight enough to keep rejecting invented titles.
        matches = [h for h in hits[:10]
                   if norm(h.get("title", "")).startswith(want + " ")
                   and h.get("year") and abs(int(h["year"]) - int(year)) <= 1]
    if not matches:
        return None
    # The title existing is the real signal; models get years wrong on films
    # that are perfectly real. Use the year only to disambiguate remakes, and
    # let TMDB's value win.
    if year:
        matches.sort(key=lambda h: abs((h.get("year") or 0) - int(year)))
    h = matches[0]
    imdb = (h.get("ratings") or {}).get("imdb") or {}
    return {
        "title": h.get("title"),
        "year": h.get("year"),
        "tmdbId": h.get("tmdbId"),
        "overview": (h.get("overview") or "").strip(),
        "genres": h.get("genres") or [],
        "runtime": h.get("runtime"),
        "votes": imdb.get("votes") or 0,
        "score": imdb.get("value") or 0,
    }


def decade_mix(lib):
    c = Counter((m["year"] // 10) * 10 for m in lib if m.get("year"))
    total = sum(c.values()) or 1
    return c, total


def decade_targets(lib, n):
    """How many of n picks each decade should get, mirroring the library."""
    c, total = decade_mix(lib)
    raw = {d: n * v / total for d, v in c.items()}
    base = {d: int(v) for d, v in raw.items()}
    for d, _ in sorted(raw.items(), key=lambda kv: -(kv[1] - int(kv[1]))):
        if sum(base.values()) >= n:
            break
        base[d] += 1
    return base


def build_prompt(lib, n):
    # Newest first: the collection is weighted toward recent films, and a
    # list that opens in the 1940s drags recommendations into the past.
    ordered = sorted(lib, key=lambda m: -(m["year"] or 0))
    lines = ["%s (%s) [%s]" % (m["title"], m["year"], ", ".join(m["genres"][:3]))
             for m in ordered]
    c, total = decade_mix(lib)
    hist = ", ".join("%ss %d%%" % (d, round(100 * c[d] / total))
                     for d in sorted(c, reverse=True))
    # Spell out the per-decade counts. Given only a percentage breakdown the
    # model swings to whichever end the wording emphasizes.
    tgt = decade_targets(lib, n)
    quota = "; ".join("%d from the %ss" % (v, d)
                      for d, v in sorted(tgt.items(), reverse=True) if v)
    return (
        "Here is my complete personal film collection (%d titles):\n\n%s\n\n"
        "By decade my collection breaks down as: %s.\n\n"
        "Recommend exactly %d films that are NOT in that list.\n"
        "Match my collection's era spread. Give me approximately: %s.\n"
        "Count as you go and respect those per-decade numbers -- a list that is "
        "mostly old films is wrong, and so is one that is all recent films.\n\n"
        "Base them on the collection as a whole: the recurring directors, eras, "
        "genres, tones and preoccupations it reveals. Infer taste from what is "
        "there.\n"
        "I am trying to grow this collection, so a film being popular is fine -- "
        "widely loved films are widely loved for a reason. Aim for a spread: "
        "roughly half well-known films most people would recognize, and half "
        "less obvious picks. What I want to avoid is a list chosen purely by "
        "popularity that ignores what my collection actually says about my taste.\n"
        "Work from CLUSTERS, not single films. For each recommendation, first "
        "find at least three films in my list that share something real -- a "
        "director, a mood, a era, a recurring theme, a kind of storytelling -- "
        "then recommend a film that belongs with that group. A pick that merely "
        "resembles one film I own is not useful; a pick that sits in the middle "
        "of several is.\n"
        "Rules:\n"
        "- Never recommend a film already in the list. Check carefully.\n"
        "- Real, released films only, with correct release years.\n"
        "- 'title' must be the bare title only. Put the year in 'year', never in 'title'.\n"
        "- 'anchors': 3 to 5 titles copied EXACTLY from my list that form the cluster.\n"
        "- 'pattern': what those anchor films share, in one phrase.\n"
        "- 'reason': one sentence on why the recommendation belongs with them.\n"
        "- Use a different cluster for each recommendation.\n"
        % (len(lines), "\n".join(lines), hist, n, quota)
    )


def ask(model, prompt, n, timeout=1800):
    # OpenAI-shaped, because the backend is LiteLLM -> llama-swap ->
    # llama-server now. Three things moved when Ollama went away:
    #
    #   /api/chat            -> /v1/chat/completions
    #   "format": <schema>   -> "response_format": {"type": "json_schema", ...}
    #                           llama.cpp compiles the schema to a GBNF grammar
    #                           and constrains sampling, same guarantee Ollama's
    #                           "format" gave.
    #   "options".num_ctx    -> gone. Context is a server-side flag now (-c
    #                           32768 in config/llama-swap/config.yaml); a
    #                           client cannot resize it per request.
    #
    # max_tokens is explicit because the OpenAI schema defaults it to a finite
    # value on some backends, and 60 recommendations of five fields each is a
    # few thousand tokens of JSON. Truncated JSON fails the parse below rather
    # than silently returning a short list.
    body = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "movie_recommendations", "schema": SCHEMA},
        },
        "temperature": 0.85,
        "max_tokens": 16384,
    }
    try:
        resp = get_json(LITELLM_URL + "/v1/chat/completions", body,
                        timeout=timeout,
                        headers={"Authorization": "Bearer " + litellm_key()})
    except urllib.error.HTTPError as e:
        sys.exit("LiteLLM error %s: %s" % (e.code, e.read().decode()[:400]))
    choice = (resp.get("choices") or [{}])[0]
    content = (choice.get("message") or {}).get("content") or ""
    # A truncated generation yields JSON that will not parse, and the parse
    # error alone does not say why. Name the real cause.
    if choice.get("finish_reason") == "length":
        sys.exit("Model hit max_tokens before closing the JSON -- lower "
                 "--oversample or raise max_tokens in ask().")
    content = re.sub(r"<think>.*?</think>", "", content, flags=re.S).strip()
    if not content:
        sys.exit("Model returned nothing (thinking-only response).")
    try:
        return json.loads(content).get("recommendations", [])
    except json.JSONDecodeError:
        m = re.search(r"\{.*\}", content, re.S)
        if not m:
            sys.exit("Could not parse model output:\n" + content[:600])
        return json.loads(m.group()).get("recommendations", [])


def tier(votes):
    """Rough 'how widely seen is this' bucket, by IMDb vote count."""
    if votes >= 200_000:
        return "widely seen"
    if votes >= 50_000:
        return "known"
    return "deeper cut"


def select(cands, n, balance, targets):
    """Choose n recommendations matching the library's era distribution.

    Models drift toward the pre-1980 canon regardless of what the collection
    looks like, so era is enforced rather than requested: each decade gets a
    quota proportional to its share of the library. Within a decade, picks
    alternate between better- and lesser-known films so the result is not all
    blockbusters. Short buckets backfill from the most recent decades, since
    that is where the collection's mass sits.
    """
    if not balance:
        return cands[:n]

    buckets = defaultdict(list)
    for c in cands:
        buckets[((c["tmdb"]["year"] or 0) // 10) * 10].append(c)

    for d, items in buckets.items():
        ranked = sorted(items, key=lambda c: -c["tmdb"]["votes"])
        half = max(len(ranked) // 2, 1)
        hi, lo = ranked[:half], ranked[half:]
        mixed = []
        while hi or lo:
            if hi:
                mixed.append(hi.pop(0))
            if lo:
                mixed.append(lo.pop(0))
        buckets[d] = mixed

    out = []
    for d in sorted(targets, reverse=True):
        q = targets[d]
        out += buckets[d][:q]
        buckets[d] = buckets[d][q:]
    if len(out) < n:
        # Round-robin the shortfall across decades rather than draining the
        # newest bucket, which would just re-create the skew the quota fixes.
        order = sorted(buckets, key=lambda d: (-targets.get(d, 0), -d))
        while len(out) < n and any(buckets[d] for d in order):
            for d in order:
                if len(out) >= n:
                    break
                if buckets[d]:
                    out.append(buckets[d].pop(0))
    return sorted(out[:n], key=lambda c: -(c["tmdb"]["year"] or 0))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("-n", "--count", type=int, default=15)
    p.add_argument("-m", "--model", default="gemma3-12b",
                   help="LiteLLM model name (see /v1/models), not an Ollama tag")
    p.add_argument("--json", action="store_true")
    p.add_argument("--prompt-only", action="store_true")
    p.add_argument("--min-votes", type=int, default=5000,
                   help="drop films below this many IMDb votes (0 disables)")
    p.add_argument("--no-balance", dest="balance", action="store_false",
                   help="skip the popularity spread and take the model's order")
    p.add_argument("--timeout", type=int, default=1800,
                   help="seconds to wait on LiteLLM; a cold model load into "
                        "llama-swap costs ~30 s on top of generation")
    p.add_argument("--oversample", type=int, default=4,
                   help="ask for COUNT*N candidates; higher fills thin decades")
    a = p.parse_args()

    key = api_key()
    lib = library(key)
    # Over-generate: a chunk of any small model's output is fabricated or
    # already owned, and both get filtered below. Extra headroom also helps
    # fill the thinly-populated decades.
    n_candidates = a.count * a.oversample
    prompt = build_prompt(lib, n_candidates)
    if a.prompt_only:
        print(prompt)
        return

    print("library: %d films | model: %s | asking for %d, keeping %d..."
          % (len(lib), a.model, n_candidates, a.count), file=sys.stderr)

    recs = ask(a.model, prompt, n_candidates, a.timeout)
    owned = {norm(m["title"]) for m in lib}

    by_norm = {norm(m["title"]): m["title"] for m in lib}
    cands, invented, already, obscure, thin = [], [], [], [], []
    seen = set()
    for r in recs:
        title = r.get("title", "")
        if norm(title) in owned:
            already.append(title)
            continue
        real = verify(key, title, r.get("year"))
        if not real:
            invented.append("%s (%s)" % (title, r.get("year")))
            continue
        if norm(real["title"]) in owned:
            already.append(real["title"])
            continue
        if real["tmdbId"] in seen:      # models repeat themselves
            continue
        # Anchors are only meaningful if they name films actually in the
        # library -- otherwise the "cluster" is invented.
        anchors = []
        for anc in r.get("anchors", []):
            hit = by_norm.get(norm(re.sub(r"\s*\(\d{4}\)\s*$", "", anc)))
            if hit and hit not in anchors:
                anchors.append(hit)
        if len(anchors) < 2:
            thin.append("%s [%d/%d anchors real]"
                        % (real["title"], len(anchors), len(r.get("anchors", []))))
            continue
        if a.min_votes and real["votes"] and real["votes"] < a.min_votes:
            obscure.append("%s (%s, %dk votes)"
                           % (real["title"], real["year"], real["votes"] // 1000))
            continue
        seen.add(real["tmdbId"])
        cands.append({**r, "tmdb": real, "anchors": anchors})

    kept = select(cands, a.count, a.balance, decade_targets(lib, a.count))

    if a.json:
        print(json.dumps({"recommendations": kept,
                          "rejected_not_in_tmdb": invented,
                          "rejected_already_owned": already,
                          "rejected_below_min_votes": obscure,
                          "rejected_invented_cluster": thin},
                         indent=2, ensure_ascii=False))
        return

    for r in kept:
        t = r["tmdb"]
        print("\n\033[1m%s\033[0m (%s)  \033[2m%s/10, %sk IMDb votes - %s\033[0m"
              % (t["title"], t["year"], t["score"], t["votes"] // 1000, tier(t["votes"])))
        print("  %s" % r.get("reason", "").strip())
        print("  \033[2mcluster (%s): %s\033[0m"
              % (r.get("pattern", "").strip(), ", ".join(r["anchors"])))

    spread = Counter(tier(r["tmdb"]["votes"]) for r in kept)
    got = Counter(((r["tmdb"]["year"] or 0) // 10) * 10 for r in kept)
    want = decade_targets(lib, a.count)
    lib_years = sorted(m["year"] for m in lib if m.get("year"))
    rec_years = sorted(r["tmdb"]["year"] for r in kept if r["tmdb"]["year"])
    print("\npopularity: %s"
          % (", ".join("%s %d" % (k, v) for k, v in spread.items()) or "-"),
          file=sys.stderr)
    print("eras (got/target): %s"
          % ", ".join("%ss %d/%d" % (d, got.get(d, 0), want.get(d, 0))
                      for d in sorted(set(got) | set(want), reverse=True)),
          file=sys.stderr)
    if rec_years:
        print("median year: recommendations %d vs library %d"
              % (rec_years[len(rec_years) // 2], lib_years[len(lib_years) // 2]),
              file=sys.stderr)
    print("%d shown | %d not real (%s) | %d owned | %d too obscure (%s) | %d invented cluster (%s)"
          % (len(kept),
             len(invented), ", ".join(invented[:3]) or "-",
             len(already),
             len(obscure), ", ".join(obscure[:3]) or "-",
             len(thin), ", ".join(thin[:3]) or "-"),
          file=sys.stderr)


if __name__ == "__main__":
    main()
