pipeline/variations.py

"""Phase 11 — variation testing: one body, many publishable hook×CTA videos
(SPEC_NEW §16.6).

Decided with the owner (§16.6): each hook×CTA combination is its **own
publishable video**, and the **body clip is generated ONCE and reused** across
every variation by reference (never regenerated).

The model:
- A `bodies` row = `{avatar_id (startframe), product_id, seed, script}` plus its
  generated body segment clip(s). The body's segments live in `segments` with
  `video_id = null`, `role = 'body'`, `body_id` set, all on the body's **locked
  seed** (identity continuity, §11.2). Generated ONCE.
- Around a body we generate a **pool of hooks** and a **pool of CTAs** — each a
  short spoken line → its own candidate clip, on the **same seed + startframe**
  as the body so the face matches. These live in `segments` with `role` of
  `hook_candidate` / `cta_candidate` and `video_id = null`.
- A `variations` row = `(body_id, hook_segment_id, cta_segment_id)` assembled
  into its **own `videos` row** (publishable), whose hook/body/cta segments are
  COPIES that point at the shared candidate + body clips by reference (no
  regeneration). Assembly (§11.4) + the approval/publish gate (§7.7) are
  unchanged — a variation is just a normal `videos` row.

KEY INVARIANT: the body clip is generated once. Every variation's body segments
carry the SAME `file_path` as the shared body segment clip — proven by the test.

§12.4 invariant: the render seed lives only in `bodies.seed` / `segments.seed`;
a variation's `videos.tags` carries `{"source":"variation","body_id":...}` and
**never** a `seed` key.
"""

from __future__ import annotations

import itertools
from typing import Any

from psycopg.types.json import Json

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


# ── Body + pool construction ─────────────────────────────────────────────────

def _insert_pool_segment(conn, body_id: str, kind: str, role: str, text: str,
                         avatar: dict[str, Any], seed: int) -> str:
    """Insert one POOL segment (body / hook_candidate / cta_candidate).

    Pool segments belong to the BODY, not to a video: `video_id` is null,
    `body_id` is set, `role` records which pool it is, and the seed is the body's
    locked seed (§16.6). The prompt is rendered from the avatar's start-frame
    template, exactly like flow.insert_segment (§11.2). Validates the line first
    (prompts.veo_segment.validate_spoken_text + duration.decide_duration).
    """
    s = get_settings()
    validate_spoken_text(text)
    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, body_id, role)
        values (null, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        returning id
        """,
        (0, 1, kind, text, prompt, int(seed), int(dur), body_id, role),
    ).fetchone()
    return str(row["id"])


def create_body(avatar_id: str, body_lines: list[str], product_id: str | None = None,
                seed: int | None = None) -> str:
    """Create a body: one reusable, identity-locked clip set (SPEC_NEW §16.6).

    Inserts a `bodies` row on a single locked seed (the seed shared by the body
    and ALL of its hook/CTA pool clips, §11.2), then one `role='body'` segment
    per body line — these hold the reusable body clips. Returns the body id.

    The body segments belong to the BODY (`video_id` null); a rendered variation
    later COPIES them (by file_path reference) into its own video.
    """
    avatar = db.fetch_one("select * from avatars where id = %s", (avatar_id,))
    if avatar is None:
        raise ValueError(f"no avatar {avatar_id}")
    if not body_lines:
        raise ValueError("create_body needs at least one body line")

    body_seed = int(seed) if seed is not None else duration.random_seed()
    # Validate every line up front so a bad line fails before we write anything.
    for line in body_lines:
        validate_spoken_text(line)

    with db.connect() as conn:
        brow = conn.execute(
            """
            insert into bodies (avatar_id, product_id, seed, script, status)
            values (%s, %s, %s, %s, 'draft')
            returning id
            """,
            (avatar_id, product_id, body_seed, Json({"body": body_lines})),
        ).fetchone()
        body_id = str(brow["id"])
        for line in body_lines:
            _insert_pool_segment(conn, body_id, "body", "body", line, avatar, body_seed)
    return body_id


def _body_with_avatar(body_id: str) -> tuple[dict[str, Any], dict[str, Any]]:
    body = db.fetch_one("select * from bodies where id = %s", (body_id,))
    if body is None:
        raise ValueError(f"no body {body_id}")
    avatar = db.fetch_one("select * from avatars where id = %s", (body["avatar_id"],))
    if avatar is None:
        raise RuntimeError(f"body {body_id} is missing its avatar {body['avatar_id']}")
    return body, avatar


def generate_hook_pool(body_id: str, hook_lines: list[str]) -> list[str]:
    """Insert a pool of hook candidates around a body (SPEC_NEW §16.6).

    Each line becomes a `role='hook_candidate'`, `kind='hook'` segment on the
    body's locked seed + startframe (so the face matches the body). Returns the
    new candidate segment ids.
    """
    body, avatar = _body_with_avatar(body_id)
    seed = body["seed"]
    ids: list[str] = []
    with db.connect() as conn:
        for line in hook_lines:
            ids.append(
                _insert_pool_segment(conn, body_id, "hook", "hook_candidate", line, avatar, seed)
            )
    return ids


def generate_cta_pool(body_id: str, cta_lines: list[str]) -> list[str]:
    """Insert a pool of CTA candidates around a body (SPEC_NEW §16.6).

    Each line becomes a `role='cta_candidate'`, `kind='cta'` segment on the
    body's locked seed + startframe. Returns the new candidate segment ids.
    """
    body, avatar = _body_with_avatar(body_id)
    seed = body["seed"]
    ids: list[str] = []
    with db.connect() as conn:
        for line in cta_lines:
            ids.append(
                _insert_pool_segment(conn, body_id, "cta", "cta_candidate", line, avatar, seed)
            )
    return ids


# ── Clip generation (body ONCE; pools once each) ─────────────────────────────

def generate_body_clips(body_id: str) -> list[str]:
    """Generate the real clips for a body's segments — body clips ONCE (§16.6).

    Renders every pool segment under the body (role body / hook_candidate /
    cta_candidate) that does not yet have a file_path, via
    veo.generate_and_download (which in DRY_RUN copies the fixture clip). The
    body clips are produced exactly once here and then REUSED by reference across
    all variations (never regenerated). Returns the segment ids generated.
    """
    from . import veo

    segs = db.fetch_all(
        """
        select id, file_path from segments
         where body_id = %s
           and role in ('body', 'hook_candidate', 'cta_candidate')
         order by role, idx
        """,
        (body_id,),
    )
    if not segs:
        raise RuntimeError(f"body {body_id} has no body/pool segments to generate")

    db.execute("update bodies set status = 'generating' where id = %s", (body_id,))
    generated: list[str] = []
    for seg in segs:
        if seg["file_path"]:            # already generated — never regenerate (§16.6)
            continue
        veo.generate_and_download(str(seg["id"]))   # sets segments.file_path
        generated.append(str(seg["id"]))
    db.execute("update bodies set status = 'ready' where id = %s", (body_id,))
    return generated


# ── Variation rendering (reuse the body clip by reference) ───────────────────

def _copy_segment(conn, video_id: str, idx: int, src: dict[str, Any], kind: str) -> str:
    """Copy a pool/body segment into a variation's video, reusing its clip.

    The new segment is a normal (role-null, body_id-null) video segment at `idx`,
    carrying the SAME file_path as the source — the shared body / candidate clip
    is REUSED by reference, never regenerated (§16.6). Prompt/seed/duration are
    carried over so the row is a faithful, re-renderable copy.
    """
    row = conn.execute(
        """
        insert into segments
            (video_id, idx, take, kind, spoken_text, veo_prompt, seed, duration_s,
             file_path, status, role, body_id)
        values (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, null, null)
        returning id
        """,
        (
            video_id, idx, kind, src["spoken_text"], src["veo_prompt"], src["seed"],
            src["duration_s"], src["file_path"],
            "qc_pass" if src["file_path"] else "pending",
        ),
    ).fetchone()
    return str(row["id"])


def build_variation(body_id: str, hook_segment_id: str, cta_segment_id: str) -> dict[str, str]:
    """Render one hook×CTA combination into its own publishable video (§16.6).

    Creates a NEW `videos` row (avatar + product from the body; tags
    `{"source":"variation","body_id":...}` — NO seed key, §12.4) and its
    idx-ordered segments:
        idx 0      = hook, COPIED from the chosen hook_candidate (clip reused)
        idx 1..k   = body, COPIED from the body's body segments (clips REUSED —
                     the shared body clip, never regenerated)
        idx last   = cta,  COPIED from the chosen cta_candidate (clip reused)
    Inserts a `variations` row linking body/hook/cta/video. Returns
    {variation_id, video_id}. The new video is a normal publishable video:
    assemble.assemble + the §7.7 approval/publish gate work unchanged.
    """
    body = db.fetch_one("select * from bodies where id = %s", (body_id,))
    if body is None:
        raise ValueError(f"no body {body_id}")

    hook = db.fetch_one(
        "select * from segments where id = %s and body_id = %s and role = 'hook_candidate'",
        (hook_segment_id, body_id),
    )
    if hook is None:
        raise ValueError(f"no hook candidate {hook_segment_id} under body {body_id}")
    cta = db.fetch_one(
        "select * from segments where id = %s and body_id = %s and role = 'cta_candidate'",
        (cta_segment_id, body_id),
    )
    if cta is None:
        raise ValueError(f"no cta candidate {cta_segment_id} under body {body_id}")

    body_segs = db.fetch_all(
        "select * from segments where body_id = %s and role = 'body' order by idx, id",
        (body_id,),
    )
    if not body_segs:
        raise RuntimeError(f"body {body_id} has no body segments")

    # §12.4: creative levers only — the render seed is NEVER written into tags.
    tags = {"source": "variation", "body_id": body_id}

    with db.connect() as conn:
        vrow = conn.execute(
            "insert into videos (avatar_id, product_id, tags, status) values (%s,%s,%s,'draft') returning id",
            (body["avatar_id"], body.get("product_id"), Json(tags)),
        ).fetchone()
        video_id = str(vrow["id"])

        idx = 0
        _copy_segment(conn, video_id, idx, hook, "hook")          # idx 0 — hook
        for bseg in body_segs:                                    # idx 1..k — body (REUSE)
            idx += 1
            _copy_segment(conn, video_id, idx, bseg, "body")
        idx += 1
        _copy_segment(conn, video_id, idx, cta, "cta")            # idx last — cta

        vr = conn.execute(
            """
            insert into variations (body_id, hook_segment_id, cta_segment_id, video_id, status)
            values (%s, %s, %s, %s, 'ready')
            returning id
            """,
            (body_id, hook_segment_id, cta_segment_id, video_id),
        ).fetchone()
        variation_id = str(vr["id"])

    return {"variation_id": variation_id, "video_id": video_id}


def build_all_variations(body_id: str) -> list[dict[str, str]]:
    """Build every hook×CTA variation under a body — the cartesian product (§16.6).

    For each (hook_candidate, cta_candidate) pair in the body's pools, calls
    build_variation. Returns the list of {variation_id, video_id} dicts (one
    publishable video per combination).
    """
    hooks = db.fetch_all(
        "select id from segments where body_id = %s and role = 'hook_candidate' order by id",
        (body_id,),
    )
    ctas = db.fetch_all(
        "select id from segments where body_id = %s and role = 'cta_candidate' order by id",
        (body_id,),
    )
    if not hooks or not ctas:
        raise RuntimeError(
            f"body {body_id} needs at least one hook candidate and one cta candidate "
            f"(have {len(hooks)} hooks, {len(ctas)} ctas)"
        )

    out: list[dict[str, str]] = []
    for hook, cta in itertools.product(hooks, ctas):
        out.append(build_variation(body_id, str(hook["id"]), str(cta["id"])))
    return out