scenario_engine.py

"""
Scenario Engine — Generates complete theatrical scripts via OpenRouter (qwen/qwen3.6-27b).
Returns a structured JSON scenario with characters, scenes, and dialogue.
"""

import os
import re
import json
import requests
from dotenv import load_dotenv

load_dotenv()

OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL_ID = "qwen/qwen3.6-27b"

# ── Preset Themes ──────────────────────────────────────────────
PRESETS = {
    "british_comedy": (
        "A hilarious British comedy of manners set in a 1920s country estate. "
        "Witty banter, class clashes between upstairs and downstairs, a bumbling lord, "
        "a sharp-tongued dowager, and a mysterious guest who is definitely not who they claim to be. "
        "Think P.G. Wodehouse meets Oscar Wilde. 3–4 characters, 3 scenes, comedic misunderstandings."
    ),
    "victorian_drama": (
        "A dark, atmospheric Victorian drama set during a stormy night in a crumbling manor. "
        "Secrets, betrayal, and a family inheritance that someone is willing to kill for. "
        "Gothic tension, candlelit rooms, whispered threats. 3–4 characters, 3 scenes, building dread."
    ),
    "infidelity": (
        "A tense modern drama about infidelity and its aftermath. "
        "A dinner party where the truth slowly unravels. Passive-aggressive dialogue, "
        "uncomfortable silences, and a devastating revelation. "
        "3–4 characters, 3 scenes, escalating tension."
    ),
}

SYSTEM_PROMPT = """/no_think
You are a brilliant playwright. Generate a COMPLETE theatrical script as valid JSON.

RULES:
- Create exactly the number of characters and scenes appropriate for the theme.
- Each character MUST have a gender field ("male" or "female").
- Write 4-6 dialogue lines per scene. Each line should be 1-2 sentences MAX. Keep lines SHORT.
- Stage directions should be brief and evocative.
- The dialogue must feel NATURAL, DRAMATIC, and THEATRICAL.
- image_prompt must be a vivid, detailed visual description for an AI image generator (no character names, just the setting/mood).
- Do NOT include any text outside the JSON block.
- Do NOT wrap in markdown code fences.
- Output ONLY the raw JSON object, nothing else.

OUTPUT FORMAT (strict JSON):
{
  "title": "The Play Title",
  "characters": [
    {"name": "Character Name", "gender": "male", "description": "Brief role description"}
  ],
  "scenes": [
    {
      "number": 1,
      "title": "Scene Title",
      "setting": "Brief stage direction for the scene opening",
      "image_prompt": "Detailed visual description for AI image generation",
      "dialogues": [
        {"character": "Character Name", "line": "What they say", "direction": "stage direction or null"}
      ]
    }
  ]
}"""


def _repair_json(raw: str) -> str:
    """Attempt to repair truncated JSON by closing open structures."""
    # Count open/close braces and brackets
    open_braces = raw.count("{") - raw.count("}")
    open_brackets = raw.count("[") - raw.count("]")

    # If we're inside a string (odd number of unescaped quotes after last structure)
    # try to close it
    stripped = raw.rstrip()
    if stripped and stripped[-1] not in ['}', ']', '"', 'l', 'e']:
        # Likely mid-string or mid-value
        raw = raw.rstrip()
        # Close any open string
        if raw.count('"') % 2 != 0:
            raw += '"'

    # Remove trailing comma if present
    raw = raw.rstrip().rstrip(",")

    # Close brackets and braces
    raw += "]" * max(0, open_brackets)
    raw += "}" * max(0, open_braces)

    return raw


def generate_scenario(theme: str) -> dict:
    """
    Generate a complete theatrical scenario.
    `theme` can be a preset key or custom text.
    Returns parsed JSON dict or raises. Retries up to 3 times.
    """
    if theme in PRESETS:
        user_prompt = f"Write a theatrical play based on this theme:\n\n{PRESETS[theme]}"
    else:
        user_prompt = f"Write a theatrical play based on this theme:\n\n{theme}"

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": MODEL_ID,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        "max_tokens": 16384,
        "temperature": 0.85,
        "provider": {"require_parameters": False},
    }

    last_error = None
    for attempt in range(3):
        try:
            print(f"[ScenarioEngine] Generating scenario (attempt {attempt+1})...")
            resp = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=180)
            resp.raise_for_status()

            raw = resp.json()["choices"][0]["message"]["content"]

            # Strip <think> blocks
            raw = re.sub(r"<think>.*?(?:</think>|$)", "", raw, flags=re.DOTALL).strip()

            # Extract JSON from markdown fences
            json_match = re.search(r"```(?:json)?\s*(.*?)```", raw, re.DOTALL)
            if json_match:
                raw = json_match.group(1).strip()

            # Find the first {
            brace_start = raw.find("{")
            if brace_start > 0:
                raw = raw[brace_start:]

            # Try parsing as-is
            try:
                scenario = json.loads(raw)
            except json.JSONDecodeError:
                print("[ScenarioEngine] JSON parse failed, attempting repair...")
                repaired = _repair_json(raw)
                scenario = json.loads(repaired)

            # Validate
            assert "title" in scenario, "Missing 'title'"
            assert "characters" in scenario, "Missing 'characters'"
            assert "scenes" in scenario, "Missing 'scenes'"
            assert len(scenario["scenes"]) > 0, "No scenes"

            print(f"[ScenarioEngine] ✓ Generated: '{scenario['title']}' "
                  f"({len(scenario['characters'])} characters, {len(scenario['scenes'])} scenes)")
            return scenario

        except Exception as e:
            last_error = e
            print(f"[ScenarioEngine] Attempt {attempt+1} failed: {e}")
            if attempt < 2:
                print("[ScenarioEngine] Retrying...")

    raise last_error