tree.py

from __future__ import annotations
import base64
import io
import uuid
from dataclasses import dataclass, field


@dataclass
class SceneNode:
    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    parent_id: str | None = None
    prompt: str = ""
    raw_input: str = ""
    reference_image_b64: str | None = None
    generated_images_b64: list[str] = field(default_factory=list)
    selected_image_idx: int | None = None
    video_prompts: list[str] = field(default_factory=list)  # one per image
    video_prompt: str = ""                                   # selected + edited
    video_b64: str | None = None
    last_frame_b64: str | None = None
    narration_text: str = ""
    narration_b64: str | None = None       # raw TTS wav
    narration_offset_s: float = 0.0
    narration_volume: float = 1.0
    narrated_video_b64: str | None = None  # video remixed with narration
    status: str = "draft"


def _thumb_data_uri(n: SceneNode, size=(168, 95)) -> str | None:
    """Small jpeg data-URI for the canvas card. Full PNGs are ~3MB each and the
    canvas HTML is rebuilt on every interaction — keep thumbs at ~5KB."""
    from PIL import Image
    src = None
    if n.generated_images_b64:
        idx = n.selected_image_idx if n.selected_image_idx is not None else 0
        src = n.generated_images_b64[idx]
    if not src:
        return None
    img = Image.open(io.BytesIO(base64.b64decode(src)))
    img.thumbnail(size)
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=70)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


@dataclass
class StoryTree:
    nodes: dict[str, SceneNode] = field(default_factory=dict)
    root_id: str | None = None

    def add_node(self, parent_id: str | None = None, raw_input: str = "") -> SceneNode:
        node = SceneNode(parent_id=parent_id, raw_input=raw_input)
        self.nodes[node.id] = node
        if parent_id is None:
            self.root_id = node.id
        return node

    def children_of(self, node_id: str) -> list[SceneNode]:
        return [n for n in self.nodes.values() if n.parent_id == node_id]

    def path_to_root(self, node_id: str) -> list[str]:
        path = []
        current = node_id
        while current is not None:
            path.append(current)
            current = self.nodes[current].parent_id
        return list(reversed(path))

    def selected_image_b64(self, node_id: str) -> str | None:
        node = self.nodes.get(node_id)
        if node and node.selected_image_idx is not None and node.generated_images_b64:
            return node.generated_images_b64[node.selected_image_idx]
        return None

    def story_so_far(self, node_id: str) -> str:
        """Chronological motion prompts along the path to root — the 'story so
        far' context for narration drafting and continuation suggestions."""
        parts = []
        for nid in self.path_to_root(node_id):
            n = self.nodes[nid]
            text = n.video_prompt or n.prompt or n.raw_input
            if text:
                parts.append(text)
        return "\n\n".join(parts)

    def to_dict(self) -> dict:
        """Serialize tree to JSON-safe dict for canvas rendering."""
        return {
            "nodes": {
                nid: {
                    "id": n.id,
                    "parent_id": n.parent_id,
                    "prompt": n.prompt,
                    "raw_input": n.raw_input,
                    "status": n.status,
                    "has_images": len(n.generated_images_b64) > 0,
                    "selected_image_idx": n.selected_image_idx,
                    "has_video": n.video_b64 is not None,
                    "has_narration": n.narrated_video_b64 is not None,
                    "thumb": _thumb_data_uri(n),
                }
                for nid, n in self.nodes.items()
            },
            "root_id": self.root_id,
        }

    def to_state_dict(self) -> dict:
        """Serialize full tree state for server_functions (includes all data)."""
        return {
            "nodes": {
                nid: {
                    "id": n.id,
                    "parent_id": n.parent_id,
                    "prompt": n.prompt,
                    "raw_input": n.raw_input,
                    "reference_image_b64": n.reference_image_b64,
                    "generated_images_b64": n.generated_images_b64,
                    "selected_image_idx": n.selected_image_idx,
                    "video_prompts": n.video_prompts,
                    "video_prompt": n.video_prompt,
                    "video_b64": n.video_b64,
                    "last_frame_b64": n.last_frame_b64,
                    "narration_text": n.narration_text,
                    "narration_b64": n.narration_b64,
                    "narration_offset_s": n.narration_offset_s,
                    "narration_volume": n.narration_volume,
                    "narrated_video_b64": n.narrated_video_b64,
                    "status": n.status,
                }
                for nid, n in self.nodes.items()
            },
            "root_id": self.root_id,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "StoryTree":
        """Deserialize tree from state dict."""
        tree = cls()
        tree.root_id = d.get("root_id")
        for nid, node_dict in d.get("nodes", {}).items():
            tree.nodes[nid] = SceneNode(**node_dict)
        return tree