puppet_theater/show_bible.py

# LLM-assisted "show bible": title, setting, optional backdrop image, director + two puppets.
# Used at session creation when the director and/or actor backend is an LLM
# (hf_api, openbmb, local_lora, local_gguf); falls back to deterministic defaults if parsing fails.

from __future__ import annotations

import json
import re
from dataclasses import replace
from typing import Any
from urllib.parse import urlparse

from puppet_theater.backends import (
    HFAPIBackend,
    HF_API_BACKDROP_URL_MAX_TOKENS,
    HF_API_SHOW_BIBLE_MAX_TOKENS,
    HF_API_SUMMON_ACTOR_MAX_TOKENS,
    LocalGGUFActorBackend,
    LocalLoRAActorBackend,
    OpenBMBTransformersBackend,
    _run_with_timeout,
    format_chatml,
    get_backend,
)
from puppet_theater.models import Actor, TheaterSession
from puppet_theater.tools import ALLOWED_TOOL_NAMES

_SHOW_BIBLE_INSTRUCTIONS = """You are casting a three-character puppet improv for "AI Puppet Theater".
Return ONE compact JSON object only. No markdown. No keys beyond the schema.

Schema:
{
  "show_title": "short catchy title, max 8 words",
  "setting": "one vivid sentence describing the stage backdrop and mood",
  "backdrop_description": "two short sentences (see Rules)",
  "director": {
    "name": "director puppet name (include a playful title)",
    "avatar": "single emoji",
    "avatar_image_url": "https portrait URL or empty string",
    "goal": "one sentence",
    "secret": "one sentence, playful",
    "speaking_style": "short phrase",
    "tools": ["subset of: change_lighting, consult_stage_oracle, inspect_prop"]
  },
  "puppet_actors": [
    {
      "name": "...",
      "avatar": "emoji",
      "avatar_image_url": "https or empty",
      "goal": "...",
      "secret": "...",
      "speaking_style": "...",
      "tools": ["..."]
    },
    {
      "name": "...",
      "avatar": "emoji",
      "avatar_image_url": "https or empty",
      "goal": "...",
      "secret": "...",
      "speaking_style": "...",
      "tools": ["..."]
    }
  ]
}

Rules:
- director.tools must include change_lighting or consult_stage_oracle (at least one).
- puppet_actors must have **at least two** entries (exactly two preferred). Never put the director inside puppet_actors; never add a third puppet there if you already have a top-level director. Extra puppets beyond two are ignored by the app.
- Names must be unique across all three characters.
- Keep every string field under 200 characters (backdrop_description may be up to ~320 characters).
- backdrop_description: REQUIRED. Describe a **minimal** wide-stage photographic background for puppets: soft light, simple shapes or gentle gradients, lots of calm empty space in the **center** for characters, **low detail**, **no crowds or faces**, **no text or logos**, **no busy fine patterns**. Mood may echo the premise but stay restrained so puppets stay readable. Do not include URLs here.
- avatar_image_url may be https://api.dicebear.com/7.x/avataaars/svg?seed=ENCODED_NAME or empty.

Premise:
"""

SHOW_BIBLE_SYSTEM_MESSAGE = (
    "You are generating show metadata for AI Puppet Theater. "
    "Return only one valid JSON object. No markdown. No commentary."
)

BACKDROP_URL_SYSTEM_MESSAGE = (
    "You return one valid JSON object for AI Puppet Theater: only the key backdrop_image_url. "
    "No markdown. No commentary."
)

# Any backend that can run a one-shot text completion for casting (same family as actor engine).
_LLM_CAST_BACKENDS = frozenset({"hf_api", "openbmb", "local_lora", "local_gguf"})


def llm_backend_order(director_mode: str, backend_name: str) -> list[str]:
    """Prefer director LLM first, then actor backend, deduplicated."""
    order: list[str] = []
    for mode in (director_mode, backend_name):
        key = (mode or "").strip().lower()
        if key in _LLM_CAST_BACKENDS and key not in order:
            order.append(key)
    return order


# Image hosts where we upgrade http→https so model output still counts as LLM backdrop.
_HTTP_TO_HTTPS_IMAGE_HOSTS = frozenset(
    {
        "images.unsplash.com",
        "plus.unsplash.com",
        "cdn.pixabay.com",
        "upload.wikimedia.org",
    }
)


def _netloc_host(netloc: str) -> str:
    host = netloc.split("@")[-1]
    return host.split(":")[0].lower()


def _strip_trailing_junk(url: str) -> str:
    u = url.strip()
    junk = ').,];}\'" \n\t'
    while u and u[-1] in junk:
        u = u[:-1]
    return u


def validate_https_url(raw: Any) -> str | None:
    if raw is None:
        return None
    url = _strip_trailing_junk(str(raw).strip())
    if not url:
        return None
    lower = url.lower()
    if lower.startswith("http://"):
        host = _netloc_host(urlparse(url).netloc)
        if host in _HTTP_TO_HTTPS_IMAGE_HOSTS:
            url = "https://" + url[7:]
    parsed = urlparse(url)
    if parsed.scheme != "https" or not parsed.netloc:
        return None
    if "@" in parsed.netloc:
        return None
    return url


def _salvage_show_bible_backdrop_url(raw_text: str) -> str | None:
    """
    Recover a backdrop URL when the JSON field is missing/invalid but the model
    still emitted a usable https image link in the completion.
    """
    patterns = (
        r"https://images\.unsplash\.com/photo-[\w-]+(?:\?[\w%.&=+-]*)?",
        r"https://plus\.unsplash\.com/premium_photo-[\w-]+(?:\?[\w%.&=+-]*)?",
        r"https://cdn\.pixabay\.com/photo/[\w/.-]+(?:\?[\w%.&=+-]*)?",
        r"https://upload\.wikimedia\.org/wikipedia/commons/[\w/.%-]+",
    )
    for pat in patterns:
        m = re.search(pat, raw_text)
        if m:
            got = validate_https_url(m.group(0))
            if got:
                return got
    return None


def _extract_json_object(text: str) -> dict[str, Any] | None:
    cleaned = text.strip()
    fence = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned, re.IGNORECASE)
    if fence:
        cleaned = fence.group(1).strip()
    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start == -1 or end <= start:
        return None
    try:
        data = json.loads(cleaned[start : end + 1])
    except json.JSONDecodeError:
        return None
    return data if isinstance(data, dict) else None


def _filter_tools(raw: Any) -> list[str]:
    if not isinstance(raw, list):
        return []
    out: list[str] = []
    for item in raw:
        name = str(item).strip()
        if name in ALLOWED_TOOL_NAMES and name not in out:
            out.append(name)
    return out


def _clip_avatar_emoji(s: str) -> str:
    t = " ".join(s.strip().split())
    if not t:
        return "🎭"
    return t[:16]


def _actor_from_dict(blob: Any, default_tools: list[str]) -> Actor | None:
    if not isinstance(blob, dict):
        return None
    name = " ".join(str(blob.get("name", "")).strip().split())
    if not name:
        return None
    goal = " ".join(str(blob.get("goal", "")).strip().split()) or "Stay in character."
    secret = " ".join(str(blob.get("secret", "")).strip().split()) or "Has a tiny backstage secret."
    style = " ".join(str(blob.get("speaking_style", "")).strip().split()) or "playful and theatrical"
    tools = _filter_tools(blob.get("tools"))
    if not tools:
        tools = list(default_tools)
    avatar = _clip_avatar_emoji(str(blob.get("avatar", "🎭")))
    img = validate_https_url(blob.get("avatar_image_url"))
    return Actor(
        name=name,
        avatar=avatar,
        goal=goal[:220],
        secret=secret[:220],
        speaking_style=style[:160],
        tools=tools,
        avatar_image_url=img,
    )


def build_backdrop_url_user_content(
    premise: str,
    show_title: str,
    setting: str,
    backdrop_description: str,
) -> str:
    prem = " ".join(premise.strip().split())[:900]
    tit = " ".join(show_title.strip().split())[:120]
    st = " ".join(setting.strip().split())[:400]
    desc = " ".join(backdrop_description.strip().split())[:500]
    return (
        "Choose one real https image URL for a puppet stage backdrop (wide landscape).\n\n"
        "Follow this art direction closely (minimal, readable behind puppets):\n"
        f"{desc}\n\n"
        "Context:\n"
        f"- Premise: {prem}\n"
        f"- Show title: {tit}\n"
        f"- Setting (narrative): {st}\n\n"
        "Rules for the image:\n"
        "- https only; prefer images.unsplash.com or plus.unsplash.com; also allowed: cdn.pixabay.com, upload.wikimedia.org.\n"
        "- Calm, uncluttered, generous empty space toward the center for puppets; no crowds, faces, signage, or busy textures.\n"
        "- Return ONE JSON object only, no markdown, no extra keys:\n"
        '  {"backdrop_image_url":"https://images.unsplash.com/photo-..."}\n'
        "- Use a genuine URL pattern; prefer well-known generic landscape/studio/sky photos if unsure."
    )


def parse_backdrop_url_response(raw_text: str) -> str | None:
    data = _extract_json_object(raw_text)
    if data:
        u = validate_https_url(data.get("backdrop_image_url"))
        if u:
            return u
        for alt in ("backdrop_url", "background_image_url"):
            u = validate_https_url(data.get(alt))
            if u:
                return u
    return _salvage_show_bible_backdrop_url(raw_text)


def resolve_backdrop_image_url_via_llm(
    backend_name: str,
    premise: str,
    show_title: str,
    setting: str,
    backdrop_description: str,
    *,
    max_new_tokens: int | None = None,
    temperature: float | None = None,
) -> tuple[str | None, dict[str, object]]:
    """
    Second LLM step: turn minimal backdrop_description (+ premise/setting) into a validated image URL.
    """
    meta: dict[str, object] = {}
    desc = " ".join(backdrop_description.strip().split())
    if not desc:
        desc = " ".join(setting.strip().split())
        meta["description_fallback"] = "setting"
    else:
        meta["description_fallback"] = None
    try:
        raw = invoke_backdrop_image_url_llm(
            backend_name,
            premise=premise,
            show_title=show_title,
            setting=setting,
            backdrop_description=desc,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
        )
        meta["raw_char_len"] = len(raw)
        meta["raw_preview"] = raw[:400]
        url = parse_backdrop_url_response(raw)
        meta["parsed_ok"] = bool(url)
        return url, meta
    except Exception as exc:
        meta["error"] = str(exc)[:500]
        return None, meta


def parse_show_bible_response(raw_text: str) -> tuple[str, str, str | None, list[Actor]] | None:
    data = _extract_json_object(raw_text)
    if not data:
        return None
    title = " ".join(str(data.get("show_title", "")).strip().split())
    setting = " ".join(str(data.get("setting", "")).strip().split())
    if not title or not setting:
        return None
    raw_desc: Any = data.get("backdrop_description") or data.get("stage_backdrop_description")
    backdrop_description = " ".join(str(raw_desc or "").strip().split()) or None
    if backdrop_description and len(backdrop_description) > 360:
        backdrop_description = backdrop_description[:360].rsplit(" ", 1)[0].rstrip(",;:") or backdrop_description[:360]
    director_blob: Any
    puppets: Any
    raw_actors = data.get("actors")
    if isinstance(raw_actors, list) and len(raw_actors) == 3 and all(isinstance(x, dict) for x in raw_actors):
        director_blob = raw_actors[0]
        puppets = raw_actors[1:3]
    else:
        director_blob = data.get("director")
        puppets = data.get("puppet_actors")
    if not isinstance(director_blob, dict) or not isinstance(puppets, list) or len(puppets) < 2:
        return None
    # Models often emit three puppets plus a separate director; we only stage director + two puppets.
    if len(puppets) > 2:
        puppets = puppets[:2]
    director = _actor_from_dict(director_blob, ["change_lighting", "consult_stage_oracle"])
    p0 = _actor_from_dict(puppets[0], ["consult_stage_oracle", "change_lighting"])
    p1 = _actor_from_dict(puppets[1], ["inspect_prop", "change_lighting"])
    if director is None or p0 is None or p1 is None:
        return None
    names = {director.name.lower(), p0.name.lower(), p1.name.lower()}
    if len(names) != 3:
        return None
    actors = [director, p0, p1]
    return title[:120], setting[:400], backdrop_description, actors


def invoke_show_bible_llm(
    backend_name: str,
    premise: str,
    *,
    max_new_tokens: int | None = None,
    temperature: float | None = None,
    extra_user_suffix: str = "",
) -> str:
    """Run one completion; raises on failure."""
    user_content = _SHOW_BIBLE_INSTRUCTIONS + premise.strip() + extra_user_suffix
    budget = max_new_tokens if max_new_tokens is not None else 256
    temp = float(temperature) if temperature is not None else 0.55

    backend = get_backend(backend_name)
    if isinstance(backend, HFAPIBackend):
        return backend._generate_text(
            user_content,
            max_tokens=HF_API_SHOW_BIBLE_MAX_TOKENS,
            temperature=0.35 if temperature is None else temp,
        )

    if isinstance(backend, OpenBMBTransformersBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        ot = max(96, min(budget, 160))
        backend.configure(max_new_tokens=ot, temperature=temp)
        try:
            return backend._generate_text(user_content)
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalLoRAActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(200, min(budget, 320))
        backend.configure(max_new_tokens=mt, temperature=temp)
        messages = [
            {"role": "system", "content": SHOW_BIBLE_SYSTEM_MESSAGE},
            {"role": "user", "content": user_content},
        ]
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(messages),
                backend.timeout_seconds,
                "Local LoRA show bible generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalGGUFActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(200, min(budget, 320))
        backend.configure(max_new_tokens=mt, temperature=temp)
        messages = [
            {"role": "system", "content": SHOW_BIBLE_SYSTEM_MESSAGE},
            {"role": "user", "content": user_content},
        ]
        prompt = format_chatml(messages)
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(prompt),
                backend.timeout_seconds,
                "Local GGUF show bible generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    raise RuntimeError(f"Backend {backend_name!r} cannot run show bible LLM")


def invoke_backdrop_image_url_llm(
    backend_name: str,
    *,
    premise: str,
    show_title: str,
    setting: str,
    backdrop_description: str,
    max_new_tokens: int | None = None,
    temperature: float | None = None,
) -> str:
    """One completion: JSON with backdrop_image_url only. Raises on failure."""
    user_content = build_backdrop_url_user_content(premise, show_title, setting, backdrop_description)
    budget = max_new_tokens if max_new_tokens is not None else 256
    temp = float(temperature) if temperature is not None else 0.45

    backend = get_backend(backend_name)

    if isinstance(backend, HFAPIBackend):
        return backend._generate_text(
            user_content,
            max_tokens=HF_API_BACKDROP_URL_MAX_TOKENS,
            system_message=BACKDROP_URL_SYSTEM_MESSAGE,
            temperature=0.25 if temperature is None else min(float(temp), 0.55),
        )

    if isinstance(backend, OpenBMBTransformersBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        ot = max(64, min(budget, 140))
        backend.configure(max_new_tokens=ot, temperature=min(temp, 0.45))
        try:
            return backend._generate_text(
                f"{BACKDROP_URL_SYSTEM_MESSAGE}\n\n{user_content}",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalLoRAActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(96, min(budget, 220))
        backend.configure(max_new_tokens=mt, temperature=min(temp, 0.45))
        messages = [
            {"role": "system", "content": BACKDROP_URL_SYSTEM_MESSAGE},
            {"role": "user", "content": user_content},
        ]
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(messages),
                backend.timeout_seconds,
                "Local LoRA backdrop URL generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalGGUFActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(96, min(budget, 220))
        backend.configure(max_new_tokens=mt, temperature=min(temp, 0.45))
        messages = [
            {"role": "system", "content": BACKDROP_URL_SYSTEM_MESSAGE},
            {"role": "user", "content": user_content},
        ]
        prompt = format_chatml(messages)
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(prompt),
                backend.timeout_seconds,
                "Local GGUF backdrop URL generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    raise RuntimeError(f"Backend {backend_name!r} cannot run backdrop URL LLM")


SUMMON_ACTOR_SYSTEM_MESSAGE = (
    "You are casting one new puppet for AI Puppet Theater. "
    "Return only one valid JSON object. No markdown. No commentary."
)

_SUMMON_ACTOR_SCHEMA = """
Return ONE compact JSON object only. No markdown. Allowed top-level keys only:
{
  "name": "unique puppet stage name (must not match any name already on stage)",
  "avatar": "single emoji",
  "avatar_image_url": "https portrait URL or empty string",
  "goal": "one sentence",
  "secret": "one sentence, playful",
  "speaking_style": "short phrase",
  "tools": ["subset of: change_lighting, consult_stage_oracle, inspect_prop"]
}

Rules:
- tools must be non-empty and only use allowed tool names.
- name must be different from every name in "Names already on stage".
- Incorporate the audience's suggested name or spirit, but you may refine it for the stage.
- Keep every string field under 200 characters.
- avatar_image_url may be https://api.dicebear.com/7.x/avataaars/svg?seed=ENCODED_NAME or empty.

Show context:
"""

_SUMMON_REMINDER_SUFFIX = (
    "\n\nReminder: respond with one JSON object only, keys "
    "name, avatar, avatar_image_url, goal, secret, speaking_style, tools — no intent/line/emotion keys."
)


def build_summon_actor_user_content(session: TheaterSession, audience_suggested_name: str) -> str:
    cast = ", ".join(a.name for a in session.actors) or "(none)"
    label = audience_suggested_name.strip() or "Mystery Guest"
    return (
        f"{_SUMMON_ACTOR_SCHEMA}"
        f"premise: {session.premise}\n"
        f"show_title: {session.show_title}\n"
        f"setting: {session.setting}\n"
        f"Names already on stage: {cast}\n"
        f"Audience asked to summon (suggested name): {label}\n"
    )


def parse_summoned_actor_response(raw_text: str) -> Actor | None:
    data = _extract_json_object(raw_text)
    if not data:
        return None
    blob: Any = data.get("summoned_actor")
    if blob is None:
        blob = data.get("actor")
    if blob is None and isinstance(data.get("name"), str):
        blob = data
    if not isinstance(blob, dict):
        return None
    return _actor_from_dict(blob, ["consult_stage_oracle", "inspect_prop", "change_lighting"])


def _unique_actor_name(actor: Actor, taken_lower: set[str], audience_label: str) -> Actor:
    if actor.name.lower() not in taken_lower:
        return actor
    base = " ".join(audience_label.split()) or actor.name
    for i in range(2, 14):
        candidate = f"{base} {i}"
        if candidate.lower() not in taken_lower:
            return replace(actor, name=candidate)
    return replace(actor, name=f"{base} the Wanderer")


def invoke_summon_actor_llm(
    backend_name: str,
    user_content: str,
    *,
    max_new_tokens: int | None = None,
    temperature: float | None = None,
    extra_user_suffix: str = "",
) -> str:
    """One completion for summoned-actor JSON; raises on failure."""
    text = user_content + extra_user_suffix
    budget = max_new_tokens if max_new_tokens is not None else 256
    temp = float(temperature) if temperature is not None else 0.55
    backend = get_backend(backend_name)

    if isinstance(backend, HFAPIBackend):
        return backend._generate_text(
            text,
            max_tokens=HF_API_SUMMON_ACTOR_MAX_TOKENS,
            temperature=0.35 if temperature is None else temp,
        )

    if isinstance(backend, OpenBMBTransformersBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        ot = max(96, min(budget, 160))
        backend.configure(max_new_tokens=ot, temperature=temp)
        try:
            return backend._generate_text(text)
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalLoRAActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(180, min(budget, 320))
        backend.configure(max_new_tokens=mt, temperature=temp)
        messages = [
            {"role": "system", "content": SUMMON_ACTOR_SYSTEM_MESSAGE},
            {"role": "user", "content": text},
        ]
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(messages),
                backend.timeout_seconds,
                "Local LoRA summon actor generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    if isinstance(backend, LocalGGUFActorBackend):
        prev_tokens = backend.max_new_tokens
        prev_temp = backend.temperature
        mt = max(180, min(budget, 320))
        backend.configure(max_new_tokens=mt, temperature=temp)
        messages = [
            {"role": "system", "content": SUMMON_ACTOR_SYSTEM_MESSAGE},
            {"role": "user", "content": text},
        ]
        prompt = format_chatml(messages)
        try:
            backend._load()
            return _run_with_timeout(
                lambda: backend._generate_text(prompt),
                backend.timeout_seconds,
                "Local GGUF summon actor generation timed out",
            )
        finally:
            backend.configure(max_new_tokens=prev_tokens, temperature=prev_temp)

    raise RuntimeError(f"Backend {backend_name!r} cannot run summon actor LLM")


def resolve_summoned_actor_via_llm_or_default(
    session: TheaterSession,
    audience_suggested_name: str,
) -> tuple[Actor, str | None, bool]:
    """
    Returns (actor, llm_backend_used_or_none, summon_llm_fallback_used).

    summon_llm_fallback_used is True when an LLM backend was tried and parsing failed for all attempts.
    """
    label = " ".join(audience_suggested_name.strip().split()) or "Mystery Guest"
    taken_lower = {a.name.lower() for a in session.actors}
    default = Actor(
        name=label,
        avatar="✨",
        goal="Make the scene stranger without derailing the finale.",
        secret="Arrived with one completely unexplained cue.",
        speaking_style="fresh, eager, and just a little too dramatic",
        tools=["consult_stage_oracle"],
    )

    candidates = llm_backend_order(session.director_mode, session.backend_name)
    if not candidates:
        return default, None, False

    user = build_summon_actor_user_content(session, label)
    for mode in candidates:
        suffixes: tuple[str, ...] = (
            ("", _SUMMON_REMINDER_SUFFIX) if mode in {"local_lora", "local_gguf"} else ("",)
        )
        for suffix in suffixes:
            try:
                raw = invoke_summon_actor_llm(
                    mode,
                    user,
                    max_new_tokens=session.backend_max_new_tokens,
                    temperature=session.backend_temperature,
                    extra_user_suffix=suffix,
                )
                actor = parse_summoned_actor_response(raw)
                if actor is None:
                    continue
                actor = _unique_actor_name(actor, taken_lower, label)
                if actor.name.lower() in taken_lower:
                    continue
                return actor, mode, False
            except Exception:
                continue

    return default, None, True