puppet_theater/session.py

# This file builds a "puppet theater show session" from a given story premise.
# It converts a short idea (premise) into a structured theatrical session with actors,
# a setting, and metadata used to simulate a puppet show performance.

import logging

from puppet_theater.models import Actor, TheaterSession
from puppet_theater.backdrop_gen import (
    backdrop_url_for_trace,
    setting_backdrop_t2i_enabled,
    try_setting_backdrop_data_url,
)
from puppet_theater.show_bible import (
    invoke_show_bible_llm,
    llm_backend_order,
    parse_show_bible_response,
    resolve_backdrop_image_url_via_llm,
)
from puppet_theater.trace import add_trace_event

logger = logging.getLogger(__name__)

# These are predefined show length configurations.
# Each tuple represents: (minimum beats, target beats, maximum beats)
# "beats" can be thought of as story steps or scene turns in the puppet show.
SHOW_LENGTH_PRESETS: dict[str, tuple[int, int, int]] = {
    "short": (5, 7, 8),
    "standard": (7, 10, 12),
    "extended": (10, 14, 16),
}

# Default show length used when no preference is provided
DEFAULT_SHOW_LENGTH = "standard"


def resolve_show_length(show_length: str | None = None) -> tuple[str, int, int, int]:
    """
    Converts a show length label (like "short", "standard", "extended")
    into actual numeric constraints (min, target, max beats).

    If the input is invalid or None, it falls back to DEFAULT_SHOW_LENGTH.
    """
    normalized = (show_length or DEFAULT_SHOW_LENGTH).strip().lower()
    if normalized not in SHOW_LENGTH_PRESETS:
        normalized = DEFAULT_SHOW_LENGTH
    min_beats, target_beats, max_beats = SHOW_LENGTH_PRESETS[normalized]
    return normalized, min_beats, target_beats, max_beats


def _clean_premise(premise: str) -> str:
    """
    Cleans up the input premise by removing extra spaces and normalizing it.
    If the premise is empty after cleanup, returns a fallback default story idea.
    """
    cleaned = " ".join(premise.strip().split())
    return cleaned or "A mysterious puppet show with no premise"


def _title_from_premise(premise: str) -> str:
    """
    Generates a show title based on important words from the premise.

    - Splits the premise into words
    - Removes punctuation
    - Picks words longer than 3 characters
    - Capitalizes them and uses up to 4 words for the title

    If no good words exist, returns a default title.
    """
    words = [word.strip(".,!?;:()[]{}\"'") for word in premise.split()]
    keywords = [word.title() for word in words if len(word.strip(".,!?;:()[]{}\"'")) > 3]
    if not keywords:
        return "The Tiny Improv"
    return f"The {' '.join(keywords[:4])}"


def _setting_from_premise(premise: str) -> str:
    """
    Chooses a stage setting based on keywords in the premise.

    This is a simple rule-based system:
    - space/moon/star → space-themed stage
    - castle/dragon/wizard → fantasy castle stage
    - detective/mystery → noir detective setting
    - kitchen/chef/toaster → kitchen stage
    - otherwise → generic puppet stage
    """
    lowered = premise.lower()
    if "moon" in lowered or "space" in lowered or "star" in lowered:
        return "a cardboard moon base with glittery stars and a squeaky hatch"
    if "castle" in lowered or "dragon" in lowered or "wizard" in lowered:
        return "a shoebox castle with velvet curtains and a suspicious tower"
    if "detective" in lowered or "mystery" in lowered:
        return "a rainy cardboard alley lit by one dramatic desk lamp"
    if "kitchen" in lowered or "toaster" in lowered or "chef" in lowered:
        return "a tiny kitchen counter where every appliance has stage fright"
    return "a pocket-sized improv stage with painted flats and a wobbly spotlight"


# Wide Unsplash images (curated IDs) used only when the LLM omits or invalidates backdrop_image_url
# (see create_show_from_premise: llm_backdrop_url or keyword fallback).
_DEFAULT_BACKDROP_URL = (
    "https://images.unsplash.com/photo-1578662996442-48f60103fc96"
    "?auto=format&fit=crop&w=1600&q=80"
)

_BACKDROP_RULES: tuple[tuple[tuple[str, ...], str], ...] = (
    (
        (
            "moon",
            "space",
            "star",
            "orbit",
            "galaxy",
            "mars",
            "alien",
            "planet",
            "rocket",
            "astronaut",
            "comet",
            "lunar",
            "cosmos",
        ),
        "https://images.unsplash.com/photo-1516339901601-2e1b62dc0c45?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("kitchen", "cook", "chef", "toast", "oven", "recipe", "pan", "stove", "fridge", "cupcake", "tea"),
        "https://images.unsplash.com/photo-1556912173-3c541015bf3b?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("castle", "dragon", "wizard", "knight", "sword", "enchant", "fairy", "throne", "dungeon"),
        "https://images.unsplash.com/photo-1518173946689-a94480f4bd0e?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        (
            "detective",
            "mystery",
            "noir",
            "crime",
            "alley",
            "shadow",
            "clue",
            "murder",
            "case file",
            "interrogat",
        ),
        "https://images.unsplash.com/photo-1428908728789-d2de25dbd4e0?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("ocean", "sea", "beach", "wave", "sail", "island", "submarine", "whale", "harbor", "pirate"),
        "https://images.unsplash.com/photo-1505118380757-91f5f5632ce0?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("forest", "wood", "tree", "cabin", "hike", "camp", "bear", "owl", "moss"),
        "https://images.unsplash.com/photo-1448375260088-37575c2fbe04?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("school", "classroom", "student", "teacher", "homework", "blackboard", "campus"),
        "https://images.unsplash.com/photo-1580582932707-520aed937a7e?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("library", "book", "scroll", "archive", "museum", "gallery"),
        "https://images.unsplash.com/photo-1507842217343-303bb9a4036a?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("circus", "carnival", "tent", "acrobat", "clown", "trapeze"),
        "https://images.unsplash.com/photo-1508807526345-15e9b5f4eaff?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("desert", "cactus", "dune", "mirage", "oasis"),
        "https://images.unsplash.com/photo-1509316785289-025f5b846b35?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("winter", "snow", "blizzard", "frost", "igloo", "icicle", "snowfall"),
        "https://images.unsplash.com/photo-1519681393784-d120267933ba?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("hospital", "doctor", "nurse", "clinic", "surgery", "medic"),
        "https://images.unsplash.com/photo-1519494026892-80bbd2d6fd0d?auto=format&fit=crop&w=1600&q=80",
    ),
    (
        ("train", "station", "locomotive", "railway", "subway"),
        "https://images.unsplash.com/photo-1474487548417-781cb71445bb?auto=format&fit=crop&w=1600&q=80",
    ),
)


def _backdrop_url_from_premise_and_setting(premise: str, setting: str) -> str:
    """
    Fallback when the show bible did not yield a usable backdrop_image_url.

    Matches keywords in premise + setting and returns a curated wide image, else _DEFAULT_BACKDROP_URL.
    """
    hay = f"{premise} {setting}".lower()
    for keywords, url in _BACKDROP_RULES:
        if any(k in hay for k in keywords):
            return url
    return _DEFAULT_BACKDROP_URL


def _default_cast() -> list[Actor]:
    """Fallback puppets when no LLM cast is available."""
    return [
        Actor(
            name="Pip the Director",
            avatar="🎬",
            goal="Keep the scene moving toward a crisp finale.",
            secret="Has already misplaced the final cue card.",
            speaking_style="brisk, theatrical, and slightly overconfident",
            tools=["change_lighting", "consult_stage_oracle"],
        ),
        Actor(
            name="Mina Moonbutton",
            avatar="🌙",
            goal="Find the emotional truth hiding inside the premise.",
            secret="Believes every prop is personally judging her.",
            speaking_style="earnest, poetic, and prone to dramatic pauses",
            tools=["consult_stage_oracle", "change_lighting"],
        ),
        Actor(
            name="Bolt McJiggle",
            avatar="🧰",
            goal="Turn every problem into a practical stage gag.",
            secret="Is secretly building a confetti finale backstage.",
            speaking_style="punchy, practical, and full of suspicious confidence",
            tools=["inspect_prop", "change_lighting"],
        ),
    ]


_REMINDER_SUFFIX = (
    "\n\nReminder: respond with exactly one JSON object using only these top-level keys: "
    "show_title, setting, backdrop_description, director, puppet_actors. "
    "backdrop_description: two short sentences, minimal uncluttered background for puppets (see Rules). "
    "The director and each puppet entry must include name, avatar, goal, secret, speaking_style, tools."
)


def _resolve_show_content_from_llm_or_defaults(
    cleaned_premise: str,
    backend_name: str,
    director_mode: str,
    backend_max_new_tokens: int,
    backend_temperature: float,
) -> tuple[str, str, str | None, list[Actor], str | None, bool, list[dict[str, object]], str | None]:
    """
    Returns show_title, setting, backdrop_image_url, actors, llm_source_or_none, cast_fallback_used,
    cast_attempt_log, and backdrop_description (LLM minimal art direction, or None).

    cast_fallback_used is True only when at least one LLM-capable backend was tried and none
    returned a valid show bible (so the built-in cast and heuristic title/setting are used).
    """
    fallback_title = _title_from_premise(cleaned_premise)
    fallback_setting = _setting_from_premise(cleaned_premise)
    default_actors = _default_cast()
    candidates = llm_backend_order(director_mode, backend_name)
    attempt_log: list[dict[str, object]] = []

    if not candidates:
        logger.info(
            "premise_cast: no LLM backends in order (director_mode=%r backend_name=%r); "
            "using heuristic title=%r heuristic_setting_chars=%s",
            director_mode,
            backend_name,
            fallback_title,
            len(fallback_setting),
        )
        return fallback_title, fallback_setting, None, default_actors, None, False, attempt_log, None

    logger.info(
        "premise_cast: trying backends %s (premise_len=%s max_new_tokens=%s temp=%s)",
        candidates,
        len(cleaned_premise),
        backend_max_new_tokens,
        backend_temperature,
    )

    for mode in candidates:
        suffixes: tuple[str, ...] = ("", _REMINDER_SUFFIX) if mode in {"local_lora", "local_gguf"} else ("",)
        for suffix in suffixes:
            entry: dict[str, object] = {
                "backend": mode,
                "with_schema_reminder": bool(suffix),
            }
            try:
                raw = invoke_show_bible_llm(
                    mode,
                    cleaned_premise,
                    max_new_tokens=backend_max_new_tokens,
                    temperature=backend_temperature,
                    extra_user_suffix=suffix,
                )
                entry["raw_char_len"] = len(raw)
                entry["raw_preview"] = raw[:450]
                parsed = parse_show_bible_response(raw)
                entry["parsed_ok"] = parsed is not None
                if parsed is None:
                    attempt_log.append(entry)
                    logger.info(
                        "premise_cast: parse_miss backend=%r reminder=%s raw_len=%s raw_head=%r",
                        mode,
                        bool(suffix),
                        len(raw),
                        raw[:240],
                    )
                    continue
                title, setting, backdrop_desc, actors = parsed
                entry["resolved_show_title"] = title
                entry["resolved_actor_names"] = [a.name for a in actors]
                entry["backdrop_description"] = (backdrop_desc or "")[:320]
                if setting_backdrop_t2i_enabled():
                    entry["backdrop_url_resolution"] = {"deferred": "hf_setting_text_to_image"}
                    entry["has_backdrop_url"] = False
                    attempt_log.append(entry)
                    logger.info(
                        "premise_cast: ok backend=%r title=%r setting_len=%s actors=%s backdrop_url=deferred_t2i",
                        mode,
                        title,
                        len(setting),
                        [a.name for a in actors],
                    )
                    return title, setting, None, actors, mode, False, attempt_log, backdrop_desc
                backdrop, bmeta = resolve_backdrop_image_url_via_llm(
                    mode,
                    cleaned_premise,
                    title,
                    setting,
                    backdrop_desc or "",
                    max_new_tokens=backend_max_new_tokens,
                    temperature=backend_temperature,
                )
                entry["backdrop_url_resolution"] = bmeta
                entry["has_backdrop_url"] = bool(backdrop)
                attempt_log.append(entry)
                logger.info(
                    "premise_cast: ok backend=%r title=%r setting_len=%s actors=%s backdrop_url=%s",
                    mode,
                    title,
                    len(setting),
                    [a.name for a in actors],
                    "yes" if backdrop else "no",
                )
                return title, setting, backdrop, actors, mode, False, attempt_log, backdrop_desc
            except Exception as exc:
                entry["error"] = str(exc)[:500]
                attempt_log.append(entry)
                logger.warning(
                    "premise_cast: exception backend=%r reminder=%s err=%r",
                    mode,
                    bool(suffix),
                    str(exc)[:400],
                )
                continue

    logger.warning(
        "premise_cast: all attempts failed; fallback title=%r actors=%s (see trace premise_cast_resolved)",
        fallback_title,
        [a.name for a in default_actors],
    )
    return fallback_title, fallback_setting, None, default_actors, None, True, attempt_log, None


def create_show_from_premise(
    premise: str,
    backend_name: str = "deterministic",
    backend_model_id: str | None = None,
    backend_max_new_tokens: int = 120,
    backend_temperature: float = 0.75,
    director_mode: str = "deterministic",
    show_length: str = DEFAULT_SHOW_LENGTH,
) -> TheaterSession:
    """
    Main function that creates a full TheaterSession from a simple premise.

    Steps it performs:
    1. Cleans the input premise
    2. Resolves show length into beats (story structure size)
    3. When director or actor backend is an LLM (hf_api, openbmb, local_lora, local_gguf), asks it for
       show title, setting, backdrop_description (minimal art direction), and three roles (director + two
       puppets) with optional portrait URLs. When HF text-to-image is enabled, the backdrop is generated
       from the setting sentence next; otherwise a second LLM call picks a stock https URL from that
       description. Deterministic cast uses heuristic title/setting only.
    4. Backdrop priority: HF text-to-image from `setting` (when enabled + token), else LLM stock URL (if
       cast succeeded without T2I deferral), else keyword stock images from premise + setting.
    5. Builds a TheaterSession object with all metadata
    6. Adds trace events for debugging/analytics

    Returns:
        TheaterSession: A fully initialized puppet theater session
    """

    cleaned_premise = _clean_premise(premise)
    active_show_length, min_beats, target_beats, max_beats = resolve_show_length(show_length)

    supported_generation_modes = {"deterministic", "openbmb", "hf_api", "local_lora", "local_gguf"}
    active_backend = backend_name if backend_name in supported_generation_modes else "deterministic"
    active_director_mode = director_mode if director_mode in supported_generation_modes else "deterministic"

    (
        show_title,
        setting,
        llm_backdrop_url,
        actors,
        cast_llm_source,
        cast_fallback_used,
        cast_attempt_log,
        backdrop_description,
    ) = _resolve_show_content_from_llm_or_defaults(
        cleaned_premise,
        active_backend,
        active_director_mode,
        backend_max_new_tokens,
        backend_temperature,
    )

    backdrop_image_url: str | None = None
    backdrop_image_source = "premise_keyword_fallback"
    setting_t2i_meta: dict[str, object] | None = None

    if setting_backdrop_t2i_enabled():
        backdrop_image_url, setting_t2i_meta = try_setting_backdrop_data_url(setting)
        if backdrop_image_url:
            backdrop_image_source = "hf_setting_text_to_image"

    if not backdrop_image_url:
        llm_url = llm_backdrop_url
        if llm_url is None and cast_llm_source and setting_backdrop_t2i_enabled():
            llm_url, bmeta = resolve_backdrop_image_url_via_llm(
                cast_llm_source,
                cleaned_premise,
                show_title,
                setting,
                backdrop_description or "",
                max_new_tokens=backend_max_new_tokens,
                temperature=backend_temperature,
            )
            if cast_attempt_log:
                cast_attempt_log[-1]["backdrop_url_resolution_after_t2i_fail"] = bmeta
                cast_attempt_log[-1]["has_backdrop_url_after_t2i_fail"] = bool(llm_url)
        if llm_url:
            backdrop_image_url = llm_url
            backdrop_image_source = "llm_from_description"
        else:
            backdrop_image_url = _backdrop_url_from_premise_and_setting(cleaned_premise, setting)
            backdrop_image_source = "premise_keyword_fallback"

    model_note = f" ({backend_model_id})" if backend_model_id else ""

    cast_note = (
        f"Cast and title from {cast_llm_source} show bible."
        if cast_llm_source
        else (
            "Cast and title use built-in defaults (no LLM-capable engine selected for casting)."
            if not llm_backend_order(active_director_mode, active_backend)
            else "Cast and title use built-in defaults (LLM show bible failed or was invalid)."
        )
    )
    if backdrop_image_source == "hf_setting_text_to_image":
        backdrop_note = "Backdrop: HF text-to-image from the setting sentence (wide stage, puppet-safe prompt)."
    elif backdrop_image_source == "llm_from_description":
        backdrop_note = "Backdrop: LLM stock image URL from minimal backdrop_description (two-step flow)."
    else:
        backdrop_note = "Backdrop: keyword stock image fallback (T2I/URL unavailable or disabled)."

    cast_roll = " · ".join(
        f"{a.name} ({a.avatar})" + (f" [img]" if a.avatar_image_url else "")
        for a in actors
    )

    # Internal log for debugging how the director created the session
    director_log = [
        f"Director created a {active_show_length} show plan.",
        f"Active backend: {active_backend}{model_note}.",
        f"Director mode: {active_director_mode}.",
        f"Show length: min {min_beats}, target {target_beats}, max {max_beats} beats.",
        cast_note,
        f"Resolved show_title: {show_title}",
        f"Resolved setting: {setting}",
        f"Resolved cast: {cast_roll}",
        f"Premise cast source: {cast_llm_source or 'deterministic/heuristic'}; LLM cast fallback used: {str(cast_fallback_used).lower()}.",
        (
            f"Backdrop description (LLM): {backdrop_description}"
            if backdrop_description
            else "Backdrop description (LLM): (omitted; URL step used setting text as art direction)."
        ),
        backdrop_note,
        (
            "Backdrop image: inline JPEG from setting (HF text-to-image)."
            if (backdrop_image_url or "").startswith("data:image")
            else f"Backdrop image URL: {backdrop_image_url or '(none)'}"
        ),
        "Three puppet actors are waiting for the first beat.",
    ]

    # Create the main theater session object that holds the entire show state
    session = TheaterSession(
        show_title=show_title,
        premise=cleaned_premise,
        setting=setting,
        actors=actors,
        backdrop_image_url=backdrop_image_url,
        backdrop_description=backdrop_description,
        beat_index=0,
        min_beats=min_beats,
        target_beats=target_beats,
        max_beats=max_beats,
        show_length_mode=active_show_length,
        transcript=[],
        props=[],
        latest_prop=None,
        latest_audience_action=None,
        director_log=director_log,
        trace_events=[],
        finale_requested=False,
        backend_name=active_backend,
        backend_model_id=backend_model_id,
        backend_max_new_tokens=backend_max_new_tokens,
        backend_temperature=backend_temperature,
        director_mode=active_director_mode,
        play_opening_curtain=True,
    )

    # Logging trace events for debugging/analytics pipeline
    add_trace_event(
        session,
        "premise_cast_resolved",
        reason_summary=(
            f"title={show_title!r}; source={cast_llm_source or 'deterministic'}; "
            f"fallback={cast_fallback_used}; attempts={len(cast_attempt_log)}"
        ),
        resolved_show_title=show_title,
        resolved_setting=setting,
        resolved_actor_names=[a.name for a in actors],
        resolved_actor_avatars=[a.avatar for a in actors],
        backdrop_image_url=backdrop_url_for_trace(backdrop_image_url),
        backdrop_image_source=backdrop_image_source,
        backdrop_description=(backdrop_description or "")[:400],
        setting_backdrop_t2i=setting_t2i_meta,
        premise_cast_source=cast_llm_source or "deterministic",
        cast_fallback_used=cast_fallback_used,
        cast_attempts=cast_attempt_log,
        llm_candidate_order=llm_backend_order(active_director_mode, active_backend),
    )

    add_trace_event(
        session,
        "show_created",
        backend_name=active_backend,
        model_id=backend_model_id,
        director_mode=active_director_mode,
        show_length=active_show_length,
        min_beats=min_beats,
        target_beats=target_beats,
        max_beats=max_beats,
        actor_count=len(actors),
        validation_status="valid",
        fallback_used=cast_fallback_used,
        premise_cast_source=cast_llm_source or "deterministic",
        backdrop_image_configured=bool(backdrop_image_url),
    )

    add_trace_event(session, "actors_created", actor_count=len(actors))

    add_trace_event(
        session,
        "director_plan_created",
        director_mode=active_director_mode,
        story_phase="opening",
        reason_summary=f"{active_show_length.title()} progress-based show plan created.",
    )

    return session