backend/api/gpt_image_frames.py

"""
Per-segment first frames via OpenAI GPT Image models (images.edit + multi-image refs).
Feeds Veo a keyframe that matches 3–4 product reference photos.
"""

from __future__ import annotations

import asyncio
import io
import os
import re
from typing import Any, List

import httpx
from fastapi import APIRouter, HTTPException
from openai import APIError, AsyncOpenAI
from pydantic import BaseModel, Field, field_validator

from utils.image_processor import compress_and_store_image
from utils.public_url import get_public_base_url

router = APIRouter()

_GPT_IMAGE_DEFAULT = "gpt-image-2"
_GPT_IMAGE_RETRY_ATTEMPTS = 3
_PROMPT_REFINER_MODEL_DEFAULT = "gpt-4.1-mini"


def _supports_input_fidelity(model: str) -> bool:
    """gpt-image-2+ models reject `input_fidelity` (invalid_input_fidelity_model)."""
    m = (model or "").strip().lower()
    return m.startswith("gpt-image-1")


def _aspect_to_size(aspect: str) -> str:
    a = (aspect or "9:16").strip()
    if a == "9:16":
        return "1024x1536"
    if a == "16:9":
        return "1536x1024"
    return "1024x1024"


def _sync_actions_text(segment: dict[str, Any]) -> str:
    at = segment.get("action_timeline") or {}
    sync = at.get("synchronized_actions") or {}
    if isinstance(sync, dict) and sync:
        return "; ".join(f"{k}: {v}" for k, v in sorted(sync.items()))
    return ""


def _text_suggests_jewelry(*parts: object) -> bool:
    t = " ".join(str(p or "").lower() for p in parts)
    if not t.strip():
        return False
    if re.search(r"\bring\s+light\b", t):
        return False
    return bool(
        re.search(
            r"\b(?:rings?|necklaces?|bracelets?|earrings?|pendants?|bangles?|anklets?|chokers?|lockets?|"
            r"studs?|hoops?|cartilage|lobe|jewelry|jewellery|925|14k|18k|karat|carat|sterling)\b",
            t,
        )
    )


def _build_frame_prompt(
    segment: dict[str, Any],
    product_name: str,
    *,
    reference_count: int = 1,
) -> str:
    sc = segment.get("scene_continuity") or {}
    at = segment.get("action_timeline") or {}
    ch = segment.get("character_description") or {}
    multi_ref = ""
    if reference_count >= 2:
        multi_ref = (
            f"\nYou are given {reference_count} reference photos of the same product (different angles or crops). "
            "Fuse them into one coherent understanding of the real item: exact shape, materials, proportions, "
            "labels, and distinctive details — not a generic lookalike. Prefer geometry and texture that appear "
            "consistently across the set of references."
        )
    lines = [
        "Generate ONE photorealistic keyframe — the first frame of a premium product video shot.",
        "The product in the references must appear with identical identity: same design, materials, scale, and details. Do not substitute a different SKU or generic item.",
        "No on-image text, logos as graphics, watermarks, or UI. Commercial photography quality.",
        f"Product: {product_name or 'hero product'}.{multi_ref}",
        "",
        "Director notes for this shot:",
    ]
    beat = _sync_actions_text(segment)
    if sc.get("camera_movement"):
        lines.append(f"- Camera / motion intent: {sc.get('camera_movement')}")
    if sc.get("lighting_state"):
        lines.append(f"- Lighting: {sc.get('lighting_state')}")
    if sc.get("environment"):
        lines.append(f"- Environment: {sc.get('environment')}")
    if sc.get("camera_position"):
        lines.append(f"- Framing: {sc.get('camera_position')}")
    if beat:
        lines.append(f"- Beat / blocking: {beat}")
    if at.get("dialogue"):
        lines.append(f"- Dialogue / VO mood (do not render text): {at.get('dialogue')}")
    if _text_suggests_jewelry(
        product_name,
        ch.get("current_state"),
        sc.get("environment"),
        beat,
        at.get("dialogue"),
    ):
        lines.append(
            "- Jewelry on skin (if shown): anatomically correct placement with believable contact—"
            "ring on ring finger between knuckles unless references clearly show another finger; "
            "earrings through lobe or correct cartilage; bracelet at wrist; necklace with natural drape. "
            "No floating jewelry, no metal clipping through skin, scale must match references."
        )
    return "\n".join(lines)


def _looks_like_safety_failure(detail: str) -> bool:
    d = (detail or "").strip().lower()
    if not d:
        return False
    markers = (
        "moderation_blocked",
        "safety system",
        "safety_violations",
        "image_generation_user_error",
        "request was rejected",
        "content policy",
    )
    return any(m in d for m in markers)


async def _refine_prompt_for_retry(
    client: AsyncOpenAI,
    prompt: str,
    failure_detail: str,
    product_name: str,
) -> str:
    """
    Ask an LLM to rewrite a blocked image prompt into a policy-safe variant
    while preserving visual intent and product identity.
    """
    model = (
        os.getenv("OPENAI_PROMPT_REFINER_MODEL") or _PROMPT_REFINER_MODEL_DEFAULT
    ).strip() or _PROMPT_REFINER_MODEL_DEFAULT
    try:
        response = await client.responses.create(
            model=model,
            input=[
                {
                    "role": "system",
                    "content": (
                        "Rewrite prompts for image generation to reduce safety rejections. "
                        "Keep the same commercial product-shot intent and realism. "
                        "Remove or neutralize any sexual, explicit, violent, self-harm, hateful, "
                        "or policy-sensitive phrasing. Return plain prompt text only."
                    ),
                },
                {
                    "role": "user",
                    "content": (
                        f"Product: {product_name or 'hero product'}\n"
                        f"Failure detail: {failure_detail}\n\n"
                        f"Original prompt:\n{prompt}\n\n"
                        "Rewrite this so it is safer while preserving shot composition, "
                        "lighting, camera intent, and product continuity."
                    ),
                },
            ],
            max_output_tokens=500,
        )
        refined = (getattr(response, "output_text", "") or "").strip()
        return refined or prompt
    except Exception:
        return prompt


class SegmentFirstFrameRequest(BaseModel):
    segment: dict[str, Any]
    reference_image_urls: List[str] = Field(..., min_length=1, max_length=4)
    aspect_ratio: str = "9:16"
    product_name: str = ""

    @field_validator("reference_image_urls")
    @classmethod
    def must_be_http(cls, urls: List[str]) -> List[str]:
        out: List[str] = []
        for u in urls:
            u = (u or "").strip()
            if u.startswith(("http://", "https://")):
                out.append(u)
        if not out:
            raise ValueError("At least one valid http(s) image URL is required")
        return out[:4]


async def _download_image(client: httpx.AsyncClient, url: str) -> tuple[bytes, str]:
    r = await client.get(
        url,
        follow_redirects=True,
        timeout=30.0,
        headers={"User-Agent": "ProductShowcase/1.0"},
    )
    r.raise_for_status()
    ct = (r.headers.get("content-type") or "image/jpeg").split(";")[0].strip()
    if not ct.startswith("image/"):
        raise ValueError(f"Not an image: {url} ({ct})")
    return r.content, ct


@router.post("/showcase/segment-first-frame")
async def generate_segment_first_frame(body: SegmentFirstFrameRequest):
    """
    Uses OpenAI Images `edits` with 1–4 reference images to synthesize a segment keyframe,
    then hosts it for Veo (same as upload-image pipeline).
    """
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise HTTPException(
            status_code=503,
            detail="OPENAI_API_KEY is required for GPT Image first frames.",
        )

    model = (os.getenv("GPT_IMAGE_MODEL") or _GPT_IMAGE_DEFAULT).strip() or _GPT_IMAGE_DEFAULT
    public_url = get_public_base_url()

    n_refs = len(body.reference_image_urls)
    prompt = _build_frame_prompt(body.segment, body.product_name, reference_count=n_refs)
    size = _aspect_to_size(body.aspect_ratio)

    file_tuples: list[tuple[str, io.BytesIO, str]] = []
    failed_refs: list[str] = []
    urls_in_order = body.reference_image_urls[:4]
    async with httpx.AsyncClient() as dl:
        for i, url in enumerate(urls_in_order):
            try:
                raw, ctype = await _download_image(dl, url)
                ext = "png" if "png" in ctype else "jpeg"
                file_tuples.append((f"ref_{i}.{ext}", io.BytesIO(raw), ctype))
            except httpx.HTTPError:
                failed_refs.append(url)
            except ValueError:
                failed_refs.append(url)

    if not file_tuples:
        if failed_refs:
            raise HTTPException(
                status_code=502,
                detail=f"Could not download any reference image. Failed URLs: {failed_refs}",
            )
        raise HTTPException(status_code=400, detail="No valid reference image URLs were provided.")

    # Regenerate prompt to match the actual number of successfully downloaded references.
    if len(file_tuples) != n_refs:
        prompt = _build_frame_prompt(
            body.segment,
            body.product_name,
            reference_count=len(file_tuples),
        )

    client = AsyncOpenAI(api_key=api_key)
    result = None
    prompt_for_attempt = prompt
    last_detail = ""
    for attempt in range(1, _GPT_IMAGE_RETRY_ATTEMPTS + 1):
        try:
            if _supports_input_fidelity(model):
                result = await client.images.edit(
                    model=model,
                    image=file_tuples,
                    prompt=prompt_for_attempt,
                    size=size,  # type: ignore[arg-type]
                    quality="high",
                    input_fidelity="high",
                    output_format="png",
                )
            else:
                result = await client.images.edit(
                    model=model,
                    image=file_tuples,
                    prompt=prompt_for_attempt,
                    size=size,  # type: ignore[arg-type]
                    quality="high",
                    output_format="png",
                )
            break
        except Exception as e:
            detail = e.message if isinstance(e, APIError) else str(e)
            last_detail = detail
            should_retry = attempt < _GPT_IMAGE_RETRY_ATTEMPTS and _looks_like_safety_failure(
                detail
            )
            if not should_retry:
                raise HTTPException(
                    status_code=502,
                    detail=f"OpenAI image edit failed ({model}): {detail}",
                )
            prompt_for_attempt = await _refine_prompt_for_retry(
                client,
                prompt_for_attempt,
                detail,
                body.product_name,
            )

    if result is None:
        raise HTTPException(
            status_code=502,
            detail=f"OpenAI image edit failed ({model}) after {_GPT_IMAGE_RETRY_ATTEMPTS} attempts: {last_detail}",
        )

    if not result.data or not result.data[0].b64_json:
        raise HTTPException(
            status_code=502,
            detail="OpenAI returned no image (b64_json missing). Check model access and billing.",
        )

    b64 = result.data[0].b64_json
    data_url = f"data:image/png;base64,{b64}"

    try:
        hosted = await compress_and_store_image(
            data_url,
            public_url,
            max_width=1920,
            max_height=1080,
            quality=92,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Could not host generated frame: {e}")

    return {
        "url": hosted,
        "model": model,
        "size": size,
    }