pipeline/flow.py

"""Shared pipeline operations — the single implementation behind BOTH the CLI
and the web UI (SPEC_NEW §11.3 two-take flow).

Prompts are built from the avatar's start-frame template (SPEC_NEW §11.2): the
start-frame image defines the whole scene, and each segment only varies the
spoken line. See prompts.veo_segment.render_prompt.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable

from psycopg.types.json import Json

from . import db, duration, storage
from .config import get_settings
from .prompts.veo_segment import render_prompt


def insert_segment(conn, video_id, idx: int, take: int, kind: str, text: str,
                   avatar: dict, seed: int) -> str:
    """Insert one segment with its rendered prompt + chosen duration bucket."""
    s = get_settings()
    prompt = render_prompt(avatar, text, kind)
    dur = duration.decide_duration(text, s.words_per_second, s.tail_budget_s)
    row = conn.execute(
        """
        insert into segments (video_id, idx, take, kind, spoken_text, veo_prompt, seed, duration_s)
        values (%s, %s, %s, %s, %s, %s, %s, %s)
        returning id
        """,
        (video_id, idx, take, kind, text, prompt, int(seed), int(dur)),
    ).fetchone()
    return str(row["id"])


def validate_script(avatar: dict, hook: str, body_segments: list[str], cta: str) -> None:
    """Raise (PromptRuleViolation / SegmentTooLong) on any part that can't be made."""
    s = get_settings()
    parts = [("hook", hook)] + [("body", t) for t in body_segments]
    if cta:
        parts.append(("cta", cta))
    for kind, text in parts:
        render_prompt(avatar, text, kind)  # validates the line
        duration.decide_duration(text, s.words_per_second, s.tail_budget_s)


def create_video(
    avatar_id: str,
    hook: str,
    body_segments: list[str],
    cta: str,
    *,
    angle: str = "manual",
    setting: str | None = None,
    takes: int | None = None,
    single_take: bool = False,
    seed: int | None = None,
) -> str:
    """Create script + video. Default: insert only the hook takes (body+CTA are
    created by pick_seed once a seed is locked). single_take inserts everything on
    one seed. Returns the new video id. Validates all parts up front (raises)."""
    s = get_settings()
    avatar = db.fetch_one("select * from avatars where id = %s", (avatar_id,))
    if avatar is None:
        raise ValueError(f"no avatar {avatar_id}")
    validate_script(avatar, hook, body_segments, cta)

    tags: dict[str, Any] = {"source": "ui" if not single_take else "ui-single", "angle": angle}
    if setting:
        tags["setting"] = setting

    with db.connect() as conn:
        srow = conn.execute(
            "insert into scripts (hook, body_segments, cta) values (%s, %s, %s) returning id",
            (hook, Json(body_segments), cta or ""),
        ).fetchone()

        if single_take:
            vseed = seed if seed is not None else duration.random_seed()
            vrow = conn.execute(
                "insert into videos (script_id, avatar_id, seed, tags) values (%s,%s,%s,%s) returning id",
                (srow["id"], avatar_id, vseed, Json(tags)),
            ).fetchone()
            parts = [("hook", hook)] + [("body", t) for t in body_segments]
            if cta:
                parts.append(("cta", cta))
            for idx, (kind, text) in enumerate(parts):
                insert_segment(conn, vrow["id"], idx, 1, kind, text, avatar, vseed)
        else:
            vrow = conn.execute(
                "insert into videos (script_id, avatar_id, tags) values (%s,%s,%s) returning id",
                (srow["id"], avatar_id, Json(tags)),
            ).fetchone()
            n = max(1, takes if takes is not None else s.seed_takes)
            for take in range(1, n + 1):
                insert_segment(conn, vrow["id"], 0, take, "hook", hook, avatar, duration.random_seed())
        return str(vrow["id"])


def create_video_from_script(
    avatar_id: str, script_text: str, *, angle: str = "manual", setting: str | None = None,
) -> str:
    """Free-form script → AI chunking (the MVP donor's break_script skill) → video.
    Chunks into hook/body/cta sized to 4–8 s clips, then create_video (SPEC_NEW §7.5)."""
    from . import agents

    avatar = db.fetch_one("select * from avatars where id = %s", (avatar_id,))
    if avatar is None:
        raise ValueError(f"no avatar {avatar_id}")
    chunks = agents.break_script(script_text, avatar)
    return create_video(
        avatar_id, chunks["hook"], chunks["body_segments"], chunks["cta"],
        angle=angle, setting=setting,
    )


def generate_pending(video_id: str, kinds: tuple[str, ...] | None = None,
                     on_event: Callable[[str], None] | None = None) -> list[str]:
    """Generate every not-yet-generated, non-superseded segment (optionally limited
    to `kinds`). Raises on the first failure (fail loud). Returns generated ids."""
    from . import veo

    rows = db.fetch_all(
        """
        select id, kind, take, status, file_path from segments
         where video_id = %s and coalesce(status,'') <> 'superseded'
         order by idx, take
        """,
        (video_id,),
    )
    if kinds:
        rows = [r for r in rows if r["kind"] in kinds]
    if not rows:
        raise RuntimeError(f"video {video_id} has no segments to generate")

    generated: list[str] = []
    db.execute("update videos set status = 'generating', error = null, gen_note = null where id = %s", (video_id,))
    for seg in rows:
        done = (seg["file_path"] and seg["status"] != "qc_fail"
                and (storage.is_remote(seg["file_path"]) or Path(seg["file_path"]).exists()))
        if done:
            if on_event:
                on_event(f"segment {seg['id']} ({seg['kind']}#{seg['take']}) already generated, skipping")
            continue
        path = veo.generate_and_download(str(seg["id"]))
        generated.append(str(seg["id"]))
        if on_event:
            on_event(f"segment {seg['id']} ({seg['kind']}#{seg['take']}) -> {path}")
    return generated


def generate_broll(video_id: str) -> list[str]:
    """Plan + generate B-roll for a video's body segments and link each clip to its
    segment (SPEC_NEW §15). No-op if the video has no product. Runs before assembly,
    which then overlays the B-roll over the talking audio (§15.6)."""
    from . import agents, seedance

    video = db.fetch_one("select * from videos where id = %s", (video_id,))
    if not video or not video.get("product_id"):
        return []
    product = db.fetch_one("select * from products where id = %s", (video["product_id"],))
    if not product:
        return []
    segs = db.fetch_all(
        "select * from segments where video_id = %s and coalesce(status,'') <> 'superseded' order by idx, take",
        (video_id,),
    )
    by_idx = {s["idx"]: s for s in segs if s["kind"] == "body"}

    made: list[str] = []
    for item in agents.plan_broll(product, segs):
        if not item.get("overlay"):
            continue
        seg = by_idx.get(item["segment_idx"])
        prompt = (item.get("seedance_prompt") or "").strip()
        if seg is None or not prompt:
            continue
        with db.connect() as conn:
            row = conn.execute(
                "insert into broll (product_id, prompt, duration_s, status) values (%s,%s,%s,'pending') returning id",
                (product["id"], prompt, seg.get("duration_s") or 6),
            ).fetchone()
        broll_id = str(row["id"])
        seedance.generate_and_download(broll_id)          # → file_path, status='ready'
        db.execute("update segments set broll_id = %s where id = %s", (broll_id, seg["id"]))
        made.append(broll_id)
    return made


def pick_seed(video_id: str, take: int | None = None) -> dict[str, Any]:
    """Lock the hook seed (explicit `take`, or auto by QC), supersede the losers,
    and create the body+CTA segments on the locked seed (SPEC_NEW §11.3)."""
    video = db.fetch_one("select * from videos where id = %s", (video_id,))
    if video is None:
        raise ValueError(f"no video {video_id}")
    takes = db.fetch_all(
        "select * from segments where video_id = %s and kind = 'hook' order by take", (video_id,)
    )
    if not takes:
        raise RuntimeError("no hook takes — create + generate first")
    for t in takes:
        if not t["file_path"]:
            raise RuntimeError(f"hook take #{t['take']} not generated yet")

    if take is not None:
        winner = next((t for t in takes if t["take"] == take), None)
        if winner is None:
            raise ValueError(f"no hook take #{take}")
    else:
        from . import qc
        for t in takes:
            if not t["qc"]:
                qc.run_qc(str(t["id"]))
        takes = db.fetch_all(
            "select * from segments where video_id = %s and kind = 'hook' order by take", (video_id,)
        )
        winner = max(takes, key=lambda t: qc.take_score(t["qc"]))

    losers = [t for t in takes if t["id"] != winner["id"]]
    avatar = db.fetch_one("select * from avatars where id = %s", (video["avatar_id"],))
    script = db.fetch_one("select * from scripts where id = %s", (video["script_id"],))
    if avatar is None or script is None:
        raise RuntimeError("video is missing its avatar or script")

    body_count = 0
    with db.connect() as conn:
        conn.execute("update videos set seed = %s, status = 'seed_picked', error = null where id = %s",
                     (winner["seed"], video_id))
        for t in losers:
            conn.execute("update segments set status = 'superseded' where id = %s", (t["id"],))
        idx = 1
        for text in (script["body_segments"] or []):
            insert_segment(conn, video_id, idx, 1, "body", text, avatar, winner["seed"])
            idx += 1
            body_count += 1
        if script["cta"]:
            insert_segment(conn, video_id, idx, 1, "cta", script["cta"], avatar, winner["seed"])
    return {"seed": winner["seed"], "take": winner["take"], "body_count": body_count,
            "superseded": len(losers)}