pipeline/veo.py

"""Veo 3.1 on Vertex AI — EXACTLY as the MVP donor does it (SPEC_NEW §11.2).

Generation goes through the raw Vertex `:predictLongRunning` / `:fetchPredictOperation`
REST endpoints (a direct port of the donor's worker `veo.ts`), NOT the google-genai
SDK. This guarantees the request matches the proven MVP path byte-for-byte:

  parameters: aspectRatio 9:16 · durationSeconds (dynamic 4/6/8) · sampleCount 1 ·
              seed (locked) · generateAudio true · resolution 720p
  instances : [{ prompt, image: { bytesBase64Encoded, mimeType } }]   ← the startframe

720p on every clip (owner requirement). The seed is passed explicitly in the request
body, so identity continuity can't be silently dropped by an SDK version.

Guarantees (unchanged):
- MAX_GENERATIONS_PER_DAY enforced against generation_log per UTC day (dry-run and
  real counted separately).
- DRY_RUN=true never touches Vertex: the fixture clip is copied; the log records
  duration/seed/image so Phase-1 acceptance is provable without creds.
- Safety blocks raise VeoSafetyBlock and are handled by the RAI policy (§12.5).
"""

from __future__ import annotations

import base64
import json
import re
import shutil
import time
import uuid
from pathlib import Path
from typing import Any
from urllib.parse import quote

import httpx

from . import db, duration
from .config import FIXTURES_DIR, MEDIA_DIR, get_settings

DEFAULT_SEGMENT_SECONDS = 8
POLL_INTERVAL_S = 15
POLL_TIMEOUT_S = 15 * 60
TRANSPORT_RETRIES = 3
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]

SEGMENTS_DIR = MEDIA_DIR / "segments"
DRY_RUN_OP_PREFIX = "dry-run/"
FIXTURE_CLIP = FIXTURES_DIR / "good_clip.mp4"


class DailyCapExceeded(RuntimeError):
    pass


class VeoSafetyBlock(RuntimeError):
    """Safety/RAI filtering. Handled by the RAI policy in generate_and_download."""

    def __init__(self, reasons: list[str]):
        self.reasons = reasons
        super().__init__(f"Veo safety filter blocked the generation: {reasons}")


def _is_rai_message(msg: str | None) -> bool:
    """A Veo operation error that is actually a safety/RAI block (MVP isRaiError)."""
    m = (msg or "").lower()
    return any(k in m for k in (
        "safety", "usage guidelines", "responsible ai", "blocked", "filtered",
        "policy", "violat", "prohibited", "rai",
    ))


# ── Daily cap ────────────────────────────────────────────────────────────────

def generations_today(dry_run: bool) -> int:
    row = db.fetch_one(
        """
        select count(*) as n from generation_log
         where dry_run = %s
           and (started_at at time zone 'utc')::date = (now() at time zone 'utc')::date
        """,
        (dry_run,),
    )
    return int(row["n"]) if row else 0


def enforce_daily_cap() -> None:
    s = get_settings()
    used = generations_today(s.dry_run)
    if used >= s.max_generations_per_day:
        mode = "dry-run" if s.dry_run else "REAL"
        raise DailyCapExceeded(
            f"MAX_GENERATIONS_PER_DAY reached: {used}/{s.max_generations_per_day} {mode} "
            f"generations already submitted today (UTC). The two-take hook spends "
            f"SEED_TAKES extra generations per video (SPEC_NEW §10)."
        )


def _log_generation(segment_id: str, dry_run: bool, model_id: str | None) -> None:
    db.execute(
        "insert into generation_log (segment_id, dry_run, model_id) values (%s, %s, %s)",
        (segment_id, dry_run, model_id),
    )


# ── Vertex REST (port of the donor's veo.ts) ─────────────────────────────────

def _vertex_base() -> str:
    s = get_settings()
    proj = s.require("GOOGLE_CLOUD_PROJECT")
    loc = s.require("VERTEX_LOCATION")
    model = s.require("VEO_MODEL_ID")
    return (
        f"https://{loc}-aiplatform.googleapis.com/v1/projects/{proj}"
        f"/locations/{loc}/publishers/google/models/{model}"
    )


_creds = None


def _access_token() -> str:
    """Mint a Vertex access token from the service-account / ADC credentials at
    GOOGLE_APPLICATION_CREDENTIALS (or application-default credentials)."""
    global _creds
    import google.auth
    from google.auth.transport.requests import Request as GoogleAuthRequest

    if _creds is None:
        s = get_settings()
        path = s.google_application_credentials
        if path:
            if not Path(path).is_file():
                raise RuntimeError(
                    f"GOOGLE_APPLICATION_CREDENTIALS points to a missing file: {path!r}."
                )
            _creds, _ = google.auth.load_credentials_from_file(path, scopes=SCOPES)
        else:
            _creds, _ = google.auth.default(scopes=SCOPES)
    if not _creds.valid:
        _creds.refresh(GoogleAuthRequest())
    return _creds.token


def _auth_headers() -> dict[str, str]:
    s = get_settings()
    return {
        "Authorization": f"Bearer {_access_token()}",
        "x-goog-user-project": s.require("GOOGLE_CLOUD_PROJECT"),
        "content-type": "application/json",
    }


def _has_startframe(ref: str | None) -> bool:
    if not ref:
        return False
    return ref.startswith(("http://", "https://")) or Path(ref).is_file()


def _image_inline(ref: str) -> dict[str, str]:
    """Inline base64 the startframe — from a local file OR an http(s) URL (e.g. a
    Supabase Storage public URL). Matches the donor, which fetches the image URL."""
    if ref.startswith(("http://", "https://")):
        r = httpx.get(ref, timeout=60)
        if r.status_code >= 400:
            raise RuntimeError(f"startframe fetch {r.status_code}: {r.text[:200]}")
        mime = (r.headers.get("content-type") or "image/jpeg").split(";")[0].strip()
        return {"bytesBase64Encoded": base64.b64encode(r.content).decode(), "mimeType": mime}
    p = Path(ref)
    mime = {".png": "image/png", ".webp": "image/webp"}.get(p.suffix.lower(), "image/jpeg")
    return {"bytesBase64Encoded": base64.b64encode(p.read_bytes()).decode(), "mimeType": mime}


def _submit(segment: dict[str, Any], avatar: dict[str, Any] | None, duration_s: int, seed: int) -> str:
    """POST :predictLongRunning exactly like the donor's veo.ts. Returns the op name."""
    s = get_settings()
    instance: dict[str, Any] = {"prompt": segment["veo_prompt"]}
    ref = avatar.get("ref_image_path") if avatar else None
    if _has_startframe(ref):
        instance["image"] = _image_inline(ref)  # the startframe, inline base64

    body = {
        "instances": [instance],
        "parameters": {
            "aspectRatio": "9:16",
            "durationSeconds": int(duration_s),
            "sampleCount": 1,
            "seed": int(seed),
            "generateAudio": True,
            "resolution": s.veo_resolution,   # 720p
        },
    }
    url = f"{_vertex_base()}:predictLongRunning"
    last: Exception | None = None
    for attempt in range(TRANSPORT_RETRIES):
        try:
            r = httpx.post(url, headers=_auth_headers(), json=body, timeout=120)
            if r.status_code == 429 or r.status_code >= 500:
                last = RuntimeError(f"Veo start {r.status_code}: {r.text[:300]}")
                time.sleep(2**attempt)
                continue
            if r.status_code >= 400:
                raise RuntimeError(f"Veo start {r.status_code}: {r.text[:500]}")
            return r.json()["name"]
        except httpx.TransportError as exc:
            last = exc
            time.sleep(2**attempt)
    raise RuntimeError(f"Veo submission failed after retries: {last}")


def _download_gcs(gcs_uri: str) -> bytes:
    m = re.match(r"^gs://([^/]+)/(.+)$", gcs_uri)
    if not m:
        raise RuntimeError(f"bad gcsUri: {gcs_uri}")
    url = (
        f"https://storage.googleapis.com/storage/v1/b/{m.group(1)}"
        f"/o/{quote(m.group(2), safe='')}?alt=media"
    )
    s = get_settings()
    r = httpx.get(
        url,
        headers={"Authorization": f"Bearer {_access_token()}",
                 "x-goog-user-project": s.require("GOOGLE_CLOUD_PROJECT")},
        timeout=300,
    )
    if r.status_code >= 400:
        raise RuntimeError(f"GCS download {r.status_code}: {r.text[:300]}")
    return r.content


# ── Per-segment render facts (duration + seed) ───────────────────────────────

def _resolve_render_facts(seg: dict[str, Any]) -> tuple[int, int]:
    s = get_settings()
    dur = seg.get("duration_s")
    if not dur:
        dur = duration.decide_duration(seg["spoken_text"], s.words_per_second, s.tail_budget_s)
    seed = seg.get("seed")
    if seed is None:
        seed = duration.random_seed()
    dur, seed = int(dur), int(seed)
    db.execute("update segments set duration_s = %s, seed = %s where id = %s", (dur, seed, str(seg["id"])))
    return dur, seed


# ── Contract functions ───────────────────────────────────────────────────────

def generate_segment(segment_id: str) -> str:
    s = get_settings()
    seg = db.fetch_one("select * from segments where id = %s", (segment_id,))
    if seg is None:
        raise RuntimeError(f"no segment {segment_id}")
    avatar = db.fetch_one(
        "select a.* from avatars a join videos v on v.avatar_id = a.id where v.id = %s",
        (seg["video_id"],),
    )

    duration_s, seed = _resolve_render_facts(seg)
    image_attached = _has_startframe(avatar.get("ref_image_path") if avatar else None)

    enforce_daily_cap()
    db.execute("update segments set status = 'generating', attempts = attempts + 1 where id = %s", (segment_id,))

    if s.dry_run:
        op_name = f"{DRY_RUN_OP_PREFIX}{uuid.uuid4()}"
        _log_generation(segment_id, True, f"dry-run-fixture 720p dur={duration_s} seed={seed} img={image_attached}")
        db.execute("update segments set veo_operation = %s where id = %s", (op_name, segment_id))
        return op_name

    _log_generation(segment_id, False, s.veo_model_id)
    op_name = _submit(seg, avatar, duration_s, seed)
    db.execute("update segments set veo_operation = %s where id = %s", (op_name, segment_id))
    return op_name


def poll(operation_name: str, segment_id: str | None = None) -> dict[str, Any]:
    if operation_name.startswith(DRY_RUN_OP_PREFIX):
        seg_id = segment_id or _segment_for_operation(operation_name)
        dest = SEGMENTS_DIR / f"{seg_id}.mp4"
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(FIXTURE_CLIP, dest)
        return {"status": "done", "file": str(dest)}

    r = httpx.post(f"{_vertex_base()}:fetchPredictOperation", headers=_auth_headers(),
                   json={"operationName": operation_name}, timeout=60)
    if r.status_code >= 400:
        raise RuntimeError(f"Veo poll {r.status_code}: {r.text[:300]}")
    data = r.json()
    if not data.get("done"):
        return {"status": "running"}
    if data.get("error"):
        err = data["error"]
        msg = err.get("message") if isinstance(err, dict) else str(err)
        if _is_rai_message(msg):                       # safety block delivered as an op error
            raise VeoSafetyBlock([str(msg)])
        raise RuntimeError(f"Veo operation failed: {err}")

    resp = data.get("response") or {}
    gen_resp = resp.get("generateVideoResponse") or {}
    filtered = resp.get("raiMediaFilteredCount") or gen_resp.get("raiMediaFilteredCount")
    if filtered:
        reasons = resp.get("raiMediaFilteredReasons") or gen_resp.get("raiMediaFilteredReasons") or ["RAI filtered"]
        raise VeoSafetyBlock(list(reasons))

    videos = resp.get("videos") or resp.get("generatedSamples") or gen_resp.get("generatedSamples") or []
    first = videos[0] if videos else None
    if not first:
        raise RuntimeError(f"Veo done with no video; response={json.dumps(resp)[:400]}")

    seg_id = segment_id or _segment_for_operation(operation_name)
    dest = SEGMENTS_DIR / f"{seg_id}.mp4"
    dest.parent.mkdir(parents=True, exist_ok=True)
    b64 = first.get("bytesBase64Encoded")
    if b64:
        dest.write_bytes(base64.b64decode(b64))
        return {"status": "done", "file": str(dest)}
    gcs = first.get("gcsUri") or (first.get("video") or {}).get("uri")
    if gcs:
        dest.write_bytes(_download_gcs(gcs))
        return {"status": "done", "file": str(dest)}
    raise RuntimeError("Veo response had neither bytes nor gcsUri")


def _segment_for_operation(operation_name: str) -> str:
    row = db.fetch_one("select id from segments where veo_operation = %s", (operation_name,))
    if row is None:
        raise RuntimeError(f"no segment recorded for operation {operation_name}")
    return str(row["id"])


def _poll_until_done(op_name: str, segment_id: str) -> str:
    deadline = time.monotonic() + POLL_TIMEOUT_S
    while True:
        result = poll(op_name, segment_id=segment_id)
        if result["status"] == "done":
            return result["file"]
        if time.monotonic() > deadline:
            raise TimeoutError(f"Veo operation {op_name} still running after {POLL_TIMEOUT_S}s")
        time.sleep(POLL_INTERVAL_S)


# ── RAI policy (SPEC_NEW §12.5) ──────────────────────────────────────────────

def _rephrase_segment_line(segment_id: str) -> str | None:
    """Sonnet rephrases the spoken line (softer/less imperative) and swaps it into
    the segment's prompt — the scene/startframe stay identical. Returns the new line."""
    s = get_settings()
    seg = db.fetch_one("select * from segments where id = %s", (segment_id,))
    if seg is None:
        return None
    from . import agents
    from .prompts.veo_segment import validate_spoken_text

    new_text = agents.rephrase_spoken_line(seg["spoken_text"])
    validate_spoken_text(new_text)
    old_q = f'"{seg["spoken_text"].strip()}"'
    new_prompt = seg["veo_prompt"].replace(old_q, f'"{new_text.strip()}"', 1)
    new_dur = duration.decide_duration(new_text, s.words_per_second, s.tail_budget_s)
    db.execute(
        "update segments set spoken_text = %s, veo_prompt = %s, duration_s = %s where id = %s",
        (new_text, new_prompt, new_dur, segment_id),
    )
    return new_text


def _set_gen_note(segment_id: str, note: str | None) -> None:
    """Surface a transient generation status to the UI (e.g. RAI seed-swap progress).
    Stored on videos.gen_note — not a creative attribute, never in videos.tags."""
    db.execute(
        "update videos set gen_note = %s where id = (select video_id from segments where id = %s)",
        (note, str(segment_id)),
    )


def _handle_safety_block(segment_id: str, first_exc: VeoSafetyBlock) -> str:
    """On a Veo safety block (§12.5): Veo's RAI filter is non-deterministic and usually
    trips on the rendered VIDEO, not the words — so we retry on a **fresh seed** (the words
    and startframe stay identical), looping up to VEO_RAI_MAX_RETRIES times. A new seed
    almost always escapes the false positive. The retry runs **automatically**; its progress
    is surfaced live via videos.gen_note. Only when the loop is exhausted does it surface as
    failed — and only then does the UI show a 'Försök igen' button."""
    s = get_settings()
    if not s.veo_rai_autoretry:
        raise first_exc
    last_exc = first_exc
    for attempt in range(1, max(1, s.veo_rai_max_retries) + 1):
        new_seed = duration.random_seed()
        db.execute("update segments set seed = %s where id = %s", (new_seed, segment_id))
        _set_gen_note(segment_id,
                      f"Veo blockade klippet — byter seed (försök {attempt}/{s.veo_rai_max_retries})…")
        print(f"[RAI] segment {segment_id}: safety block — auto-retry {attempt}/{s.veo_rai_max_retries} "
              f"on a fresh seed {new_seed} (same words + startframe).")
        op_name = generate_segment(segment_id)
        try:
            return _poll_until_done(op_name, segment_id)
        except VeoSafetyBlock as e2:
            last_exc = e2
            continue
    print(f"[RAI] segment {segment_id}: still blocked after {s.veo_rai_max_retries} fresh seeds — surfacing "
          f"(the UI now shows 'Försök igen' for a fresh {s.veo_rai_max_retries}-retry loop).")
    raise last_exc


def _maybe_upload_clip(segment_id: str, local_path: str) -> str:
    """Upload a generated clip to Supabase Storage (so nothing lives only locally);
    return the public URL, or the local path on failure / dry-run."""
    from . import storage

    if get_settings().dry_run or not storage.enabled():
        return local_path
    try:
        return storage.upload_bytes("gens", f"segments/{segment_id}.mp4",
                                    Path(local_path).read_bytes(), "video/mp4")
    except Exception as exc:  # noqa: BLE001 — keep local on failure
        print(f"[storage] clip upload failed ({exc}); keeping local")
        return local_path


def generate_and_download(segment_id: str) -> str:
    op_name = generate_segment(segment_id)
    try:
        path = _poll_until_done(op_name, segment_id)
    except VeoSafetyBlock as exc:
        path = _handle_safety_block(segment_id, exc)
    file_path = _maybe_upload_clip(segment_id, path)
    db.execute("update segments set file_path = %s where id = %s", (file_path, segment_id))
    _set_gen_note(segment_id, None)   # clear any RAI seed-swap note — this segment is done
    return file_path