"""Replicate Flux 2 Pro client — scene generation only (no text overlay).

Verbatim transcription of the proven flow from wom-flux-generate.py:
- POST to black-forest-labs/flux-2-pro
- Poll until succeeded
- Download output URL
- Rate-limit spacing between calls

Params (width=1200, height=640, guidance=3.5, safety=5, output_format=png) are
the exact values that produced the working WOM scenes.
"""
import os
import time
import requests
from pathlib import Path

API      = "https://api.replicate.com/v1/models/black-forest-labs/flux-2-pro/predictions"
POLL_URL = "https://api.replicate.com/v1/predictions/{pid}"

POLL_INTERVAL_SEC   = 5
RATE_LIMIT_SLEEP_SEC = 12   # Proven safe spacing from the recovered flow
MAX_POLL_SEC        = 300   # Hard timeout per prediction


class FluxError(Exception):
    pass


def _token() -> str:
    tok = os.environ.get("REPLICATE_API_TOKEN", "").strip()
    if not tok:
        raise FluxError("REPLICATE_API_TOKEN not set in environment")
    return tok


def generate_scene(prompt: str, out_path: Path) -> Path:
    """Generate a single Flux 2 Pro scene and save to out_path.

    Parameters match the proven WOM generator exactly.
    Retries once on submission failure.
    """
    headers = {
        "Authorization": f"Bearer {_token()}",
        "Content-Type": "application/json",
    }
    payload = {
        "input": {
            "prompt": prompt,
            "width": 1200,
            "height": 640,
            "guidance": 3.5,
            "output_format": "png",
            "safety_tolerance": 5,
        }
    }

    # Submit (with one retry on transient failure)
    data = None
    for attempt in (1, 2):
        resp = requests.post(API, headers=headers, json=payload, timeout=60)
        data = resp.json() if resp.content else {}
        if "id" in data:
            break
        if attempt == 1:
            print(f"  [flux] submit failed ({data}), retrying in 15s…")
            time.sleep(15)
    if not data or "id" not in data:
        raise FluxError(f"Flux submission failed after retry: {data}")

    pid = data["id"]
    print(f"  [flux] submitted: {pid}")

    # Poll
    waited = 0
    while waited < MAX_POLL_SEC:
        time.sleep(POLL_INTERVAL_SEC)
        waited += POLL_INTERVAL_SEC
        r = requests.get(POLL_URL.format(pid=pid), headers=headers, timeout=30)
        d = r.json()
        status = d.get("status")
        if status == "succeeded":
            url = d["output"]
            if isinstance(url, list):
                url = url[0]
            img = requests.get(url, timeout=60).content
            out_path.parent.mkdir(parents=True, exist_ok=True)
            out_path.write_bytes(img)
            print(f"  [flux] ✓ {out_path.name} ({len(img)//1024}KB)")
            return out_path
        if status == "failed":
            raise FluxError(f"Flux prediction failed: {d.get('error')}")
        print(f"  [flux] … {status}")

    raise FluxError(f"Flux prediction timed out after {MAX_POLL_SEC}s")


def pause_for_rate_limit():
    """Sleep between back-to-back Flux calls to stay under rate limits.

    12s is the proven interval from the WOM generator.
    """
    time.sleep(RATE_LIMIT_SLEEP_SEC)