playwright.py

"""
playwright.py — Script generation and director/character rewrite logic via Qwen on OpenRouter.
"""

import json
import requests
from voice_pool import VOICE_POOL

_MODEL = "google/gemma-4-26b-a4b-it"
_OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"

# Keys the playwright may assign
_VOICE_LABELS = list(VOICE_POOL.keys())

_SYSTEM_PROMPT_GENERATE = f"""You are a master theater playwright. Your task is to write a complete, high-quality theatrical script in strict JSON format.

VOICE LABELS available (you must assign one per character, chosen to fit personality):
{json.dumps(_VOICE_LABELS, indent=2)}

Output ONLY a single valid JSON object matching this schema exactly:
{{
  "title": "string",
  "genre": "string",
  "characters": [
    {{
      "name": "string",
      "personality": "string (2-3 sentences describing the character)",
      "voice_label": "string (must be one of the voice labels listed above)"
    }}
  ],
  "scenes": [
    {{
      "scene_number": 1,
      "theme": "string",
      "image_prompt": "string (vivid, specific, stageable, widescreen theater scene)",
      "image_path": "string (relative path like generated_images/scene_1.png)",
      "dialogues": [
        {{
          "character": "string (must match a character name exactly)",
          "action": "string (stage direction, e.g. 'raising his hand slowly')",
          "line": "string (the spoken line, natural and theatrical)"
        }}
      ]
    }}
  ]
}}

Rules:
- Generate exactly 4 scenes minimum
- Each scene must have at least 4 dialogue entries
- Make dialogues feel natural, witty, and theatrical
- Image prompts must be vivid, specific, and visually stageable
- Assign voice labels thoughtfully based on character personality
- Return ONLY valid JSON — no markdown fences, no preamble, no explanation. Do not output anything outside of the JSON block.
"""

_SYSTEM_PROMPT_REWRITE = """You are a master theater playwright assisting with a live script rewrite.
Return ONLY a valid JSON array of scene objects — the same schema as before, no markdown, no preamble.
Each scene object: {{ "scene_number": int, "theme": str, "image_prompt": str, "image_path": str, "dialogues": [{{"character": str, "action": str, "line": str}}] }}
The dialogues array MUST contain objects with 'character', 'action', and 'line' keys. DO NOT output strings for dialogues.
"""


def _call_model(api_key: str, system: str, user: str) -> str:
    """Call the LLM via OpenRouter and return the raw response text."""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": _MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "max_tokens": 8192,
        "temperature": 0.8,
        "provider": {
            "order": ["Cloudflare"]
        }
    }
    response = requests.post(_OPENROUTER_URL, headers=headers, json=payload, timeout=120)
    if response.status_code != 200:
        raise RuntimeError(
            f"OpenRouter API error {response.status_code}: {response.text[:300]}"
        )
    return response.json()["choices"][0]["message"]["content"].strip()


def _strip_fences(raw: str) -> str:
    """Strip markdown code fences if the model wrapped its JSON in them."""
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
        raw = raw.strip()
    if raw.endswith("```"):
        raw = raw[: raw.rfind("```")].strip()
    return raw


def _parse_json_with_retry(api_key: str, system: str, user: str) -> dict | list:
    """Call the model and parse JSON; retry once with stricter instruction on failure."""
    raw = _strip_fences(_call_model(api_key, system, user))

    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        retry_user = (
            user + "\n\nIMPORTANT: Your previous response was not valid JSON. "
            "Return ONLY valid JSON, no text before or after."
        )
        raw2 = _strip_fences(_call_model(api_key, system, retry_user))
        try:
            return json.loads(raw2)
        except json.JSONDecodeError as e:
            raise ValueError(
                f"Playwright returned malformed JSON after retry: {e}\n\nRaw output:\n{raw2[:500]}"
            )


def generate_script(topic: str, api_key: str) -> dict:
    """
    Generate a complete theater script JSON dict for the given topic.

    Args:
        topic: The user-provided scenario description.
        api_key: OpenRouter API key.

    Returns:
        Parsed script dict matching the schema.
    """
    user_prompt = (
        f"Write a complete theatrical script about the following topic:\n\n{topic}\n\n"
        "Remember: return ONLY valid JSON, no preamble."
    )
    result = _parse_json_with_retry(api_key, _SYSTEM_PROMPT_GENERATE, user_prompt)
    if not isinstance(result, dict):
        raise ValueError("Playwright response was not a JSON object.")
    return result


def rewrite_from_director(script: dict, current_scene_index: int, event: str, api_key: str) -> dict:
    """
    Rewrite all scenes from current_scene_index onward based on a director event.

    Args:
        script: The full current script dict.
        current_scene_index: Index (0-based) of the scene currently playing.
        event: The director's injected event description.
        api_key: OpenRouter API key.

    Returns:
        Updated full script dict with rewritten scenes merged in.
    """
    completed = script["scenes"][:current_scene_index]
    completed_summary = json.dumps(
        [{"scene_number": s["scene_number"], "theme": s["theme"]} for s in completed],
        indent=2,
    )

    user_prompt = f"""You are rewriting a theatrical script mid-performance.

TITLE: {script['title']}
GENRE: {script['genre']}
CHARACTERS: {json.dumps(script['characters'], indent=2)}

COMPLETED SCENES (do not rewrite these):
{completed_summary}

DIRECTOR EVENT (a sudden in-world occurrence that must be integrated):
"{event}"

Rewrite ALL remaining scenes (starting from scene number {current_scene_index + 1}) to naturally incorporate and follow from this director event.
Keep the same characters, title, and genre.
Return ONLY a JSON array of the rewritten scene objects (from scene {current_scene_index + 1} onward).
Each scene: {{"scene_number": int, "theme": str, "image_prompt": str, "image_path": str, "dialogues": [{{"character": "string", "action": "string", "line": "string"}}]}}
No preamble, no markdown.
"""

    rewritten = _parse_json_with_retry(api_key, _SYSTEM_PROMPT_REWRITE, user_prompt)
    if not isinstance(rewritten, list):
        raise ValueError("Rewrite response was not a JSON array.")

    updated_script = dict(script)
    updated_script["scenes"] = completed + rewritten
    return updated_script


def rewrite_from_character(script: dict, current_scene_index: int, character_input: str, api_key: str) -> dict:
    """
    Rewrite scenes from current_scene_index onward based on a character action or new character entry.

    Args:
        script: The full current script dict.
        current_scene_index: Index (0-based) of the scene currently playing.
        character_input: Description of what the character does or says unexpectedly.
        api_key: OpenRouter API key.

    Returns:
        Updated full script dict with rewritten scenes merged in.
    """
    completed = script["scenes"][:current_scene_index]
    completed_summary = json.dumps(
        [{"scene_number": s["scene_number"], "theme": s["theme"]} for s in completed],
        indent=2,
    )

    user_prompt = f"""You are rewriting a theatrical script mid-performance.

TITLE: {script['title']}
GENRE: {script['genre']}
CHARACTERS: {json.dumps(script['characters'], indent=2)}

COMPLETED SCENES (do not rewrite these):
{completed_summary}

CHARACTER DISRUPTION (a character does something unexpected, or a new character enters):
"{character_input}"

Rewrite ALL remaining scenes (starting from scene number {current_scene_index + 1}) to naturally integrate this character disruption.
If a new character is introduced, add them to the ensemble and give them a fitting voice_label from: {json.dumps(_VOICE_LABELS)}.
Keep the same title and genre.
Return ONLY a JSON array of the rewritten scene objects (from scene {current_scene_index + 1} onward).
Each scene: {{"scene_number": int, "theme": str, "image_prompt": str, "image_path": str, "dialogues": [{{"character": "string", "action": "string", "line": "string"}}]}}
No preamble, no markdown.
"""

    rewritten = _parse_json_with_retry(api_key, _SYSTEM_PROMPT_REWRITE, user_prompt)
    if not isinstance(rewritten, list):
        raise ValueError("Rewrite response was not a JSON array.")

    updated_script = dict(script)
    updated_script["scenes"] = completed + rewritten
    return updated_script