pipeline/casting.py

"""Phase 9 — audience-matched startframe selection (SPEC_NEW §16.5).

`select_startframe(audience_target, product_id, script_context)` picks the best
startframe (an `avatars` row, §16.1) for a product given a free-form audience
description (e.g. the §16.4 winning audience: "kvinnor 50+, intresse hälsa").

The contract (§16.5):
- **SQL pre-filter on `product_id` ONLY** — the creative anchor. NO age/gender
  buckets: the owner decided Claude reasons about age itself. Retired startframes
  (`status = 'retired'`) are excluded; everything else (incl. NULL status) is a
  candidate.
- **Claude ranks** ALL candidates by how well their `audience` / `selection_text`
  fit `audience_target` + `script_context` — it decides which age/look wins — and
  returns the best `startframe_id`, a `rationale`, and an **adapted prompt**
  (that startframe's `prompt_template` tailored to the script/variation). The JSON
  is parsed defensively and the call retries ONCE with the parse error appended.
- **Deterministic fallback** (no ANTHROPIC_API_KEY): score every candidate by
  case-insensitive token overlap between the stringified `audience_target` and the
  candidate's `selection_text` + `audience` values; pick the highest score,
  tie-break on the most-recent `created_at`. The prompt is the chosen
  `prompt_template` unchanged; the rationale is "matched on: <overlapping tokens>".
  This makes audience matching demonstrable WITHOUT Claude (DRY_RUN-testable).

Seed-orthogonality holds (§12.4): the startframe / audience are creative levers in
`videos.tags`, never the render seed. This module never sees a seed.
"""

from __future__ import annotations

import json
import re
from typing import Any

from .config import get_settings

# ── Audience → startframe ranking system prompt (§16.5) ──────────────────────

CASTING_SYSTEM = """You cast the best STARTFRAME (a UGC avatar in a locked scene) \
for a Swedish video ad, given a target AUDIENCE and the SCRIPT it will speak.

You are given a list of candidate startframes for ONE product. Each candidate has \
an id, a one-line `selection_text` descriptor, and a structured `audience` object \
(gender, age as FREE text, appearance, setting, vibe, ...). There are NO fixed age \
buckets — YOU decide which age and look best fit the target audience by reasoning \
over the descriptors. A "kvinnor 50+" audience should pick a candidate who reads as \
an older woman; a younger-skewing audience should pick a younger one, and so on.

Pick the SINGLE best-fitting candidate. Then ADAPT that candidate's \
`prompt_template` (its scene-locking prompt) to THIS specific script/variation — \
keep the person, scene, framing and identity intact, only tailor the wording so it \
fits what the avatar is about to say. Do not invent a new scene.

Output ONLY this JSON object — no prose, no markdown fences:
{"startframe_id": "the chosen candidate id, exactly as given",
 "prompt": "the adapted scene-locking prompt for this script",
 "rationale": "one short sentence on why this startframe fits the audience"}"""


# ── Token helpers (deterministic fallback, §16.5) ────────────────────────────

# Swedish-friendly tokenizer: letters/digits incl. å ä ö; everything else splits.
_TOKEN_RE = re.compile(r"[0-9a-zåäö]+", re.IGNORECASE)
# Tiny stopword set so generic glue words don't fake an "overlap" match.
_STOPWORDS = frozenset(
    {"och", "i", "med", "som", "för", "av", "en", "ett", "den", "det", "på", "till",
     "ca", "the", "a", "an", "of", "and", "with"}
)


def _tokens(value: Any) -> set[str]:
    """Lowercased content tokens from any value (str / dict / list), stopwords removed."""
    return {t for t in _TOKEN_RE.findall(_stringify(value).lower()) if t not in _STOPWORDS}


def _stringify(value: Any) -> str:
    """Flatten audience_target / audience metadata to a single searchable string.
    A dict/list is rendered via its values (keys like "gender" are not content)."""
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return " ".join(_stringify(v) for v in value.values())
    if isinstance(value, (list, tuple, set)):
        return " ".join(_stringify(v) for v in value)
    return str(value)


def _candidate_text(cand: dict) -> str:
    """The matchable text for a candidate: its selection_text + audience values."""
    return f"{cand.get('selection_text') or ''} {_stringify(cand.get('audience'))}"


# Stem length for soft token overlap so Swedish inflections of the same word match
# (e.g. "kvinnor" ↔ "kvinna" share the 5-char stem "kvinn"). Short tokens (gender
# words, ages) must still match exactly, so the stem is only applied to longer tokens.
_STEM_LEN = 4
_NUM_RE = re.compile(r"^\d+$")


def _stem(token: str) -> str:
    return token[:_STEM_LEN] if len(token) > _STEM_LEN else token


def _ages(tokens: set[str]) -> set[int]:
    return {int(t) for t in tokens if _NUM_RE.match(t)}


def _score(cand: dict, target_tokens: set[str]) -> tuple[float, set[str]]:
    """Audience-fit score for one candidate (case-insensitive TOKEN OVERLAP, §16.5).

    A shared word counts when its stems match (so "kvinnor"↔"kvinna" overlap). On
    top of the overlap count we add a small age-proximity bonus: a target age (e.g.
    "över 50") pulls toward the nearest candidate age ("ca 55") rather than a distant
    one ("ca 28"), without ever outweighing a real word match. Returns (score, the
    overlapping target tokens) for the "matched on:" rationale."""
    cand_tokens = _tokens(_candidate_text(cand))
    cand_stems = {_stem(t) for t in cand_tokens}
    overlap = {t for t in target_tokens if _stem(t) in cand_stems}
    score = float(len(overlap))

    target_ages, cand_ages = _ages(target_tokens), _ages(cand_tokens)
    if target_ages and cand_ages:
        nearest = min(abs(ta - ca) for ta in target_ages for ca in cand_ages)
        # < 1.0 so it only ever breaks ties between equal word-overlap candidates.
        score += max(0.0, 1.0 - nearest / 100.0)
    return score, overlap


def _fallback_select(candidates: list[dict], audience_target: Any) -> dict:
    """Deterministic audience match (no Claude): highest token-overlap score wins,
    tie-break on most-recent created_at. Candidates arrive pre-sorted created_at DESC,
    so the first-seen candidate keeps a tie (§16.5)."""
    target_tokens = _tokens(audience_target)
    best: dict | None = None
    best_score = -1.0
    best_overlap: set[str] = set()
    for cand in candidates:  # created_at DESC → first seen wins a tie
        score, overlap = _score(cand, target_tokens)
        if score > best_score:
            best, best_score, best_overlap = cand, score, overlap
    assert best is not None  # candidates is non-empty (checked by the caller)
    matched = ", ".join(sorted(best_overlap)) if best_overlap else "none (newest active fallback)"
    return {
        "startframe_id": str(best["id"]),
        "prompt": best.get("prompt_template") or "",
        "rationale": f"matched on: {matched}",
    }


# ── Claude defensive JSON parse (§16.5) ──────────────────────────────────────

def _parse_choice(raw: str, valid_ids: set[str]) -> dict:
    """Parse the CASTING_SYSTEM reply; validate the chosen id is a real candidate."""
    s = raw.strip()
    a, b = s.find("{"), s.rfind("}")
    if a == -1 or b == -1:
        raise ValueError("no JSON object in select_startframe output")
    obj = json.loads(s[a:b + 1])
    sid = str(obj.get("startframe_id") or "").strip()
    if sid not in valid_ids:
        raise ValueError(f"chosen startframe_id {sid!r} is not a candidate for this product")
    prompt = str(obj.get("prompt") or "").strip()
    rationale = str(obj.get("rationale") or "").strip()
    if not prompt:
        raise ValueError("select_startframe returned an empty adapted prompt")
    return {"startframe_id": sid, "prompt": prompt, "rationale": rationale}


def _build_user_prompt(audience_target: Any, candidates: list[dict], script_context: Any) -> str:
    """The user turn: the target audience, the script context, and every candidate
    (id + selection_text + audience + its base prompt_template) for Claude to rank."""
    cand_view = [
        {
            "id": str(c["id"]),
            "selection_text": c.get("selection_text") or "",
            "audience": c.get("audience") or {},
            "prompt_template": c.get("prompt_template") or "",
        }
        for c in candidates
    ]
    return (
        "AUDIENCE (free-form, from the winning ad-set — §16.4):\n"
        f"{json.dumps(audience_target, ensure_ascii=False)}\n\n"
        "SCRIPT CONTEXT for this variation (may be null):\n"
        f"{json.dumps(script_context, ensure_ascii=False)}\n\n"
        "CANDIDATE STARTFRAMES (pick exactly one id):\n"
        f"{json.dumps(cand_view, ensure_ascii=False, indent=2)}"
    )


# ── Public API (§16.5) ───────────────────────────────────────────────────────

def select_startframe(
    audience_target: Any,
    product_id: str,
    script_context: Any | None = None,
) -> dict:
    """Choose the best startframe for `product_id` given a free-form `audience_target`.

    Returns {startframe_id (str), prompt, rationale}. `prompt` is the chosen
    startframe's scene-locking prompt — adapted to `script_context` by Claude, or the
    base `prompt_template` unchanged in the deterministic fallback (SPEC_NEW §16.5).

    SQL pre-filters `avatars` on `product_id` ONLY (no age/gender buckets) and drops
    retired startframes. Raises if the product has zero startframes.
    """
    from . import db

    s = get_settings()

    # SQL pre-filter: this product's non-retired startframes. coalesce(status,'active')
    # so a NULL status counts as active. created_at DESC for the fallback tie-break.
    candidates = db.fetch_all(
        """
        select id, name, persona, prompt_template, product_id, audience,
               selection_text, status, created_at
          from avatars
         where product_id = %s
           and coalesce(status, 'active') <> 'retired'
         order by created_at desc
        """,
        (product_id,),
    )
    if not candidates:
        raise RuntimeError(
            f"select_startframe: product {product_id} has no startframes — "
            f"generate or upload at least one (§16.3) before casting"
        )

    # Claude ranks ALL candidates (§16.5); deterministic fallback when there is no key.
    if s.anthropic_api_key:
        import anthropic

        valid_ids = {str(c["id"]) for c in candidates}
        client = anthropic.Anthropic(api_key=s.anthropic_api_key)
        user = _build_user_prompt(audience_target, candidates, script_context)
        last_err: str | None = None
        for attempt in range(2):
            content = user if attempt == 0 else (
                f"{user}\n\nYour previous reply could not be used ({last_err}). "
                f"Reply with ONLY valid JSON in the required format."
            )
            resp = client.messages.create(
                model=s.anthropic_model, max_tokens=1024,
                system=CASTING_SYSTEM, messages=[{"role": "user", "content": content}],
            )
            text = "".join(b.text for b in resp.content if getattr(b, "type", None) == "text")
            try:
                return _parse_choice(text, valid_ids)
            except Exception as exc:  # noqa: BLE001 — retry once, then fall back
                last_err = str(exc)
        # Claude failed twice → deterministic match still returns a usable startframe.
        return _fallback_select(candidates, audience_target)

    return _fallback_select(candidates, audience_target)