cinegen/models.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional


@dataclass
class CharacterSpec:
    identifier: str
    name: str
    role: str
    description: str
    traits: List[str] = field(default_factory=list)
    reference_image: Optional[str] = None

    def to_row(self) -> dict:
        traits = ", ".join(self.traits)
        return {
            "ID": self.identifier,
            "Name": self.name,
            "Role": self.role,
            "Traits": traits or "—",
        }


@dataclass
class SceneBeat:
    scene_id: str
    title: str
    visuals: str
    action: str
    characters: List[str] = field(default_factory=list)
    duration: int = 6
    mood: str = ""
    camera: str = ""

    def to_row(self) -> dict:
        return {
            "Scene": self.scene_id,
            "Title": self.title,
            "Action": self.action,
            "Visuals": self.visuals,
            "Characters": ", ".join(self.characters) or "—",
            "Duration (s)": self.duration,
        }


@dataclass
class Storyboard:
    title: str
    synopsis: str
    style: str
    inspiration_hint: Optional[str]
    characters: List[CharacterSpec] = field(default_factory=list)
    scenes: List[SceneBeat] = field(default_factory=list)

    def characters_table(self) -> List[dict]:
        return [char.to_row() for char in self.characters]

    def scenes_table(self) -> List[dict]:
        return [scene.to_row() for scene in self.scenes]