cinegen/story_engine.py

from __future__ import annotations

import json
import os
from typing import Any, Dict, Optional

from .models import Storyboard, CharacterSpec, SceneBeat
from .placeholders import (
    build_stub_storyboard,
    describe_image_reference,
    normalize_scene_count,
)

DEFAULT_STORY_MODEL = os.environ.get("CINEGEN_STORY_MODEL", "gemini-2.5-flash")


def _load_google_client(api_key: Optional[str]):
    if not api_key:
        return None, "Missing API key"

    try:
        from google import genai

        client = genai.Client(api_key=api_key)
        return client, None
    except Exception as exc:  # pragma: no cover - depends on optional deps
        return None, str(exc)


class StoryGenerator:
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.environ.get("GOOGLE_API_KEY")
        self.client, self.client_error = _load_google_client(self.api_key)

    def generate(
        self,
        idea: str,
        style: str,
        scene_count: int | float | str,
        inspiration_path: Optional[str] = None,
    ) -> Storyboard:
        scene_total = normalize_scene_count(scene_count)
        if not self.client:
            return build_stub_storyboard(
                idea=idea,
                style=style,
                scene_count=scene_total,
                inspiration_hint=describe_image_reference(inspiration_path),
            )

        prompt = self._build_prompt(idea, style, scene_total)
        contents = [prompt]
        parts = self._maybe_add_image_part(inspiration_path)
        contents = parts + contents if parts else contents

        try:  # pragma: no cover - relies on remote API
            response = self.client.models.generate_content(
                model=DEFAULT_STORY_MODEL,
                contents=contents,
                config={"response_mime_type": "application/json"},
            )
            payload = json.loads(response.text)
            return self._parse_payload(
                payload,
                style=style,
                inspiration_hint=describe_image_reference(inspiration_path),
            )
        except Exception:
            return build_stub_storyboard(
                idea=idea,
                style=style,
                scene_count=scene_total,
                inspiration_hint=describe_image_reference(inspiration_path),
            )

    @staticmethod
    def _build_prompt(idea: str, style: str, scene_count: int) -> str:
        return (
            "You are CineGen, an AI film director. Convert the provided idea into a "
            "structured storyboard JSON with the following keys:\n"
            "{\n"
            '  "title": str,\n'
            '  "synopsis": str,\n'
            '  "characters": [\n'
            '     {"id": "CHAR-1", "name": str, "role": str, "description": str, "traits": [str, ...]}\n'
            "  ],\n"
            '  "scenes": [\n'
            '     {"id": "SCENE-1", "title": str, "visuals": str, "action": str, "characters": [str], "duration": int, "mood": str, "camera": str}\n'
            "  ]\n"
            "}\n"
            f"Idea: {idea or 'Use the inspiration image only.'}\n"
            f"Visual Style: {style}\n"
            f"Scene Count: {scene_count}\n"
            "Ensure every scene references at least one character ID."
        )

    def _maybe_add_image_part(self, inspiration_path: Optional[str]):
        if not inspiration_path or not os.path.exists(inspiration_path):
            return None
        try:
            from google.genai import types  # pragma: no cover - optional dependency

            with open(inspiration_path, "rb") as handle:
                data = handle.read()
            mime = "image/png" if inspiration_path.endswith(".png") else "image/jpeg"
            return [types.Part.from_bytes(data=data, mime_type=mime)]
        except Exception:
            return None

    @staticmethod
    def _parse_payload(
        payload: Dict[str, Any],
        style: str,
        inspiration_hint: Optional[str],
    ) -> Storyboard:
        characters = [
            CharacterSpec(
                identifier=item.get("id", f"CHAR-{idx+1}"),
                name=item.get("name", f"Character {idx+1}"),
                role=item.get("role", "Supporting"),
                description=item.get("description", ""),
                traits=item.get("traits", []),
            )
            for idx, item in enumerate(payload.get("characters", []))
        ]
        scenes = [
            SceneBeat(
                scene_id=item.get("id", f"SCENE-{idx+1}"),
                title=item.get("title", f"Scene {idx+1}"),
                visuals=item.get("visuals", ""),
                action=item.get("action", ""),
                characters=item.get("characters", []),
                duration=int(item.get("duration", 6)),
                mood=item.get("mood", ""),
                camera=item.get("camera", ""),
            )
            for idx, item in enumerate(payload.get("scenes", []))
        ]
        if not characters or not scenes:
            raise ValueError("Incomplete payload")
        return Storyboard(
            title=payload.get("title", "Untitled Short"),
            synopsis=payload.get("synopsis", ""),
            style=style,
            inspiration_hint=inspiration_hint,
            characters=characters,
            scenes=scenes,
        )