pipeline/image_gen.py

"""Nano-Banana startframe generation + the `describe_startframe` vision skill
(SPEC_NEW §16.3).

The creative-supply layer builds the startframe bank: take an existing startframe
(an `avatars` row = startframe image + scene-locking prompt) and a product image,
and produce a NEW startframe where the person holds the new product — same person,
pose, framing, lighting and background. The model is **Nano Banana Pro on Vertex
AI**, called the SAME way `pipeline/veo.py` calls Veo: ADC auth via the Veo Google
Cloud project + a raw `generateContent` REST request (NOT the google-genai SDK), so
it reuses one ADC / one project (§16.3).

Each new startframe is then DESCRIBED by Claude vision (`describe_startframe`): it
reads the new image + the base scene prompt + the product knowledge and returns
(a) the `audience` object + a one-line `selection_text` (§16.2, read off the image),
and (b) an ADAPTED `prompt_template` — the base startframe's scene-locking prompt
rewritten to match THIS image (new product in hand, same person/scene). A new
`avatars` row is inserted with all of it.

DRY_RUN / no-creds rule (CRITICAL, §16.3): in dry-run, or when the Vertex creds
(`NANO_BANANA_MODEL` / `GOOGLE_CLOUD_PROJECT`) are unset, we NEVER touch Vertex — we
reuse the base startframe image URL as the new startframe image (no network). With no
`ANTHROPIC_API_KEY`, `describe_startframe` returns a deterministic fallback (carry the
base avatar's audience, derive `selection_text` from names, `prompt_template` = base
prompt unchanged). The whole path is therefore testable with zero credentials.
"""

from __future__ import annotations

import base64
import json
import time
import uuid
from pathlib import Path
from typing import Any

import httpx
from psycopg.types.json import Json

from . import db, storage
from .config import get_settings

# Vertex generateContent for Nano Banana Pro (Gemini image model). The model id comes
# from NANO_BANANA_MODEL and the request/response shape (generateContent with two
# inline_data image parts + a text part → inlineData image bytes on a candidate part)
# is best-effort — VERIFY-AT-BUILD-TIME when the owner enables Nano Banana Pro in the
# Veo Vertex project (§16.9), exactly like VEO_MODEL_ID / META_API_VERSION.
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
TRANSPORT_RETRIES = 3

DEFAULT_INSTRUCTION = (
    "Replace the product the person is holding with the product in the second image; "
    "keep the person, pose, framing, lighting and background identical"
)


# ── Vertex ADC auth + raw REST (mirrors pipeline/veo.py) ─────────────────────

def _vertex_base() -> str:
    """The Vertex publisher-model base URL for Nano Banana Pro (same project/location
    as Veo). VERIFY-AT-BUILD-TIME: NANO_BANANA_MODEL exact id (§16.9)."""
    s = get_settings()
    proj = s.require("GOOGLE_CLOUD_PROJECT")
    loc = s.require("VERTEX_LOCATION")
    model = s.require("NANO_BANANA_MODEL")
    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). Port of
    pipeline/veo.py._access_token so both reuse the one ADC."""
    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 _fetch_image(ref: str) -> tuple[bytes, str]:
    """Return (bytes, mime) for an image ref — a local file OR an http(s) URL (e.g. a
    Supabase Storage public URL). Mirrors veo._image_inline's two branches."""
    if ref.startswith(("http://", "https://")):
        r = httpx.get(ref, timeout=60)
        if r.status_code >= 400:
            raise RuntimeError(f"image fetch {r.status_code}: {r.text[:200]}")
        mime = (r.headers.get("content-type") or "image/jpeg").split(";")[0].strip()
        return r.content, mime
    p = Path(ref)
    mime = {".png": "image/png", ".webp": "image/webp"}.get(p.suffix.lower(), "image/jpeg")
    return p.read_bytes(), mime


def _inline_part(ref: str) -> dict[str, Any]:
    """A Gemini generateContent inline_data image part for the given image ref."""
    data, mime = _fetch_image(ref)
    return {"inline_data": {"mime_type": mime, "data": base64.b64encode(data).decode()}}


def _generate_image_bytes(base_ref: str, product_ref: str, prompt: str) -> bytes:
    """Call Nano Banana Pro on Vertex (generateContent) with BOTH images + a prompt;
    return the generated PNG bytes. Raw REST + ADC, mirroring veo._submit. The image
    parts are ordered base-startframe first, product second (the prompt's "second
    image"). VERIFY-AT-BUILD-TIME: request/response shape (§16.9)."""
    body = {
        "contents": [
            {
                "role": "user",
                "parts": [
                    _inline_part(base_ref),     # first image — the person/scene
                    _inline_part(product_ref),  # second image — the product to swap in
                    {"text": prompt},
                ],
            }
        ],
        "generationConfig": {"responseModalities": ["IMAGE"]},
    }
    url = f"{_vertex_base()}:generateContent"
    last: Exception | None = None
    for attempt in range(TRANSPORT_RETRIES):
        try:
            r = httpx.post(url, headers=_auth_headers(), json=body, timeout=180)
            if r.status_code == 429 or r.status_code >= 500:
                last = RuntimeError(f"Nano Banana {r.status_code}: {r.text[:300]}")
                time.sleep(2**attempt)
                continue
            if r.status_code >= 400:
                raise RuntimeError(f"Nano Banana {r.status_code}: {r.text[:500]}")
            return _extract_image_bytes(r.json())
        except httpx.TransportError as exc:
            last = exc
            time.sleep(2**attempt)
    raise RuntimeError(f"Nano Banana request failed after retries: {last}")


def _extract_image_bytes(data: dict[str, Any]) -> bytes:
    """Pull the first inline image out of a generateContent response. Tolerant of
    camelCase/snake_case (inlineData/inline_data) — VERIFY-AT-BUILD-TIME (§16.9)."""
    for cand in data.get("candidates") or []:
        for part in ((cand.get("content") or {}).get("parts") or []):
            blob = part.get("inlineData") or part.get("inline_data")
            if blob and (blob.get("data")):
                return base64.b64decode(blob["data"])
    raise RuntimeError(f"Nano Banana response had no image; response={json.dumps(data)[:400]}")


# ── Claude vision skill: describe_startframe (§16.2 + §16.3) ──────────────────

DESCRIBE_SYSTEM = """You are a casting director for short Swedish UGC video ads. You \
look at ONE start-frame image (a person holding a product, talking to camera) and \
produce structured selection metadata so the right start-frame can later be matched \
to an ad audience.

Return ONLY a single JSON object, no prose, with EXACTLY these keys:
{
  "audience": {
    "gender": "kvinna" | "man" | "annat",
    "age": "<free text/number, e.g. 'ca 55' — NO fixed buckets>",
    "appearance": "<free text in Swedish, hair/makeup/clothing>",
    "setting": "<e.g. 'ljust kök', 'vardagsrum', 'utomhus'>",
    "vibe": "<e.g. 'varm', 'energisk', 'saklig'>",
    "language": "sv",
    "product_in_hand": true | false
  },
  "selection_text": "<ONE Swedish line describing this person+scene for matching>",
  "prompt_template": "<the BASE scene-locking prompt, rewritten to match THIS image: \
the NEW product now in the person's hand, but the SAME person, pose, framing, \
lighting and background. Keep the {script_chunk} placeholder if the base prompt had \
one.>"
}

Estimate the age as free text — there are NO age buckets. Write Swedish values."""


def _describe_fallback(base_audience: dict | None, base_prompt: str | None,
                       base_name: str | None, product_name: str | None) -> dict:
    """Deterministic, no-key fallback (§16.3): carry the base avatar's audience (or
    {}), derive a one-line selection_text from the base/product names, and leave the
    base scene prompt UNCHANGED as the adapted prompt_template."""
    audience = dict(base_audience or {})
    bits = [b for b in (base_name, product_name) if b]
    selection_text = " · ".join(bits) if bits else (base_name or product_name or "Startframe")
    return {
        "audience": audience,
        "selection_text": selection_text,
        "prompt_template": base_prompt or "",
    }


def _image_block(ref: str) -> dict[str, Any] | None:
    """A Claude vision image block for an image ref (URL or local file). Mirrors
    agents._product_image_block; returns None if the image can't be read."""
    try:
        data, mime = _fetch_image(ref)
    except Exception:  # noqa: BLE001 — image unavailable → caller falls back
        return None
    return {"type": "image", "source": {"type": "base64", "media_type": mime,
                                        "data": base64.standard_b64encode(data).decode()}}


def _parse_describe_json(raw: str, base_audience: dict | None, base_prompt: str | None,
                         base_name: str | None, product_name: str | None) -> dict:
    """Defensive parse of the vision JSON (first/last brace). Missing pieces fall back
    to the deterministic values so a partial answer is still usable."""
    s = raw.strip()
    a, b = s.find("{"), s.rfind("}")
    if a == -1 or b == -1:
        raise ValueError("no JSON object in describe_startframe output")
    obj = json.loads(s[a:b + 1])
    fb = _describe_fallback(base_audience, base_prompt, base_name, product_name)
    audience = obj.get("audience")
    selection_text = str(obj.get("selection_text") or "").strip()
    prompt_template = str(obj.get("prompt_template") or "").strip()
    return {
        "audience": audience if isinstance(audience, dict) else fb["audience"],
        "selection_text": selection_text or fb["selection_text"],
        "prompt_template": prompt_template or fb["prompt_template"],
    }


def describe_startframe(image_ref: str, base_prompt: str | None, product: dict | None) -> dict:
    """Claude VISION skill (§16.3): read the new startframe image + the base scene
    prompt + the product knowledge and return:
      {audience: {...}, selection_text: "<1 line>", prompt_template: "<adapted>"}.

    `audience` is the §16.2 object (gender, age FREE text, appearance, setting, vibe,
    language, product_in_hand). `prompt_template` is the base prompt rewritten to match
    THIS image (new product in hand, same person/scene). Defensive JSON parse + retry
    once. DETERMINISTIC FALLBACK with no ANTHROPIC_API_KEY (carry base audience,
    selection_text from names, prompt_template = base prompt unchanged)."""
    product = product or {}
    base_audience = product.get("_base_audience") if isinstance(product, dict) else None
    base_name = (product.get("_base_name") if isinstance(product, dict) else None)
    product_name = product.get("name") if isinstance(product, dict) else None

    s = get_settings()
    block = _image_block(image_ref) if image_ref else None
    if not s.anthropic_api_key or block is None:
        return _describe_fallback(base_audience, base_prompt, base_name, product_name)

    import anthropic

    knowledge = product.get("knowledge") or {}
    user_text = (
        f"Product: {product_name or ''}\n"
        f"Product knowledge: {json.dumps(knowledge, ensure_ascii=False)}\n\n"
        f"Base scene-locking prompt (to ADAPT for this image):\n{base_prompt or ''}\n\n"
        f"Look at the attached start-frame image and return the JSON object."
    )
    client = anthropic.Anthropic(api_key=s.anthropic_api_key)
    last_err: str | None = None
    for attempt in range(2):
        text_part = user_text if attempt == 0 else (
            f"{user_text}\n\nYour previous answer could not be parsed as JSON "
            f"({last_err}). Reply with ONLY valid JSON in the required format."
        )
        resp = client.messages.create(
            model=s.anthropic_model, max_tokens=1024, system=DESCRIBE_SYSTEM,
            messages=[{"role": "user", "content": [block, {"type": "text", "text": text_part}]}],
        )
        text = "".join(b.text for b in resp.content if getattr(b, "type", None) == "text")
        try:
            return _parse_describe_json(text, base_audience, base_prompt, base_name, product_name)
        except Exception as exc:  # noqa: BLE001 — retry once, then deterministic fallback
            last_err = str(exc)
    return _describe_fallback(base_audience, base_prompt, base_name, product_name)


# ── Startframe generation (§16.3) ────────────────────────────────────────────

def _new_startframe_image(base_ref: str, product_ref: str | None, product_id: str,
                          new_id: str, instruction: str | None) -> str:
    """Produce the new startframe image and return its URL/path.

    DRY_RUN or no Vertex creds (NANO_BANANA_MODEL / GOOGLE_CLOUD_PROJECT): NEVER call
    Vertex — reuse the base startframe image URL as the new image (no network). Live:
    call Nano Banana Pro with both images + the prompt → new bytes → Supabase
    `startframes/{product_id}/{new_id}.png` → public URL (§16.3, §16.8)."""
    s = get_settings()
    use_vertex = (not s.dry_run) and bool(s.nano_banana_model) and bool(s.google_cloud_project)
    if not use_vertex or not product_ref:
        return base_ref  # reuse the base image — deterministic, no network
    prompt = (instruction or DEFAULT_INSTRUCTION).strip()
    image_bytes = _generate_image_bytes(base_ref, product_ref, prompt)
    return storage.upload_bytes(
        s.supabase_bucket, f"startframes/{product_id}/{new_id}.png", image_bytes, "image/png"
    )


def generate_startframe(base_startframe_id: str, product_id: str,
                        instruction: str | None = None) -> str:
    """Generate ONE new startframe from a base startframe + a product (SPEC_NEW §16.3).

    1. Load the base startframe (avatars.ref_image_path + prompt_template + audience)
       and the product (products.image_path + name + knowledge).
    2. Produce the new startframe image (Nano Banana Pro live; reuse base image in
       dry-run / no-creds — see _new_startframe_image).
    3. Describe it with Claude vision → audience + selection_text + adapted prompt.
    4. Insert a new avatars row (product_id, ref_image_path, audience, selection_text,
       adapted prompt_template, base_startframe_id, generator='nano_banana',
       persona={}). Return the new id (stringified).
    """
    base = db.fetch_one("select * from avatars where id = %s", (str(base_startframe_id),))
    if base is None:
        raise RuntimeError(f"no base startframe {base_startframe_id}")
    product = db.fetch_one("select * from products where id = %s", (str(product_id),))
    if product is None:
        raise RuntimeError(f"no product {product_id}")

    base_ref = base.get("ref_image_path")
    if not base_ref:
        raise RuntimeError(f"base startframe {base_startframe_id} has no ref_image_path")
    product_ref = product.get("image_path")

    new_id = str(uuid.uuid4())
    new_image = _new_startframe_image(base_ref, product_ref, str(product_id), new_id, instruction)

    # Describe the new image: carry base audience/name through to the deterministic
    # fallback so it stays demoable with no ANTHROPIC_API_KEY (§16.3).
    desc = describe_startframe(
        new_image,
        base.get("prompt_template"),
        {
            "name": product.get("name"),
            "knowledge": product.get("knowledge") or {},
            "_base_audience": base.get("audience") or {},
            "_base_name": base.get("name"),
        },
    )

    name = _derive_name(base.get("name"), product.get("name"))
    row = db.fetch_one(
        """
        insert into avatars
            (id, name, ref_image_path, persona, prompt_template, product_id, audience,
             selection_text, base_startframe_id, generator)
        values (%s, %s, %s, %s, %s, %s, %s, %s, %s, 'nano_banana')
        returning id
        """,
        (
            new_id, name, new_image, Json({}), desc["prompt_template"], str(product_id),
            Json(desc["audience"] or {}), desc["selection_text"], str(base_startframe_id),
        ),
    )
    assert row is not None
    return str(row["id"])


def _derive_name(base_name: str | None, product_name: str | None) -> str:
    """A readable name for the new startframe: base name + ' · ' + product name."""
    bits = [b for b in (base_name, product_name) if b]
    return " · ".join(bits) if bits else (base_name or product_name or "Startframe")


def generate_startframe_batch(base_startframe_ids: list[str], product_id: str,
                              n: int | None = None) -> list[str]:
    """Generate many startframes at once ("massa nya avatarer", §16.3). Calls
    generate_startframe `n` times (default 1) across the base ids round-robin; returns
    the new ids in creation order."""
    bases = list(base_startframe_ids or [])
    if not bases:
        raise RuntimeError("generate_startframe_batch needs at least one base startframe id")
    count = n if n is not None else 1
    out: list[str] = []
    for i in range(max(0, count)):
        base_id = bases[i % len(bases)]
        out.append(generate_startframe(base_id, product_id))
    return out