#!/usr/bin/env python3
"""Fruitdrama still-generator met Nano Banana (gemini-2.5-flash-image).
Genereert een hero-referentie en trekt de characters consistent door alle scenes
door telkens de hero-image als referentie mee te sturen."""
import base64, json, os, sys, time, urllib.request

KEY = open(os.path.join(os.path.dirname(__file__), ".env")).read().split("=", 1)[1].strip()
MODEL = "gemini-2.5-flash-image"
URL = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={KEY}"
OUT = os.path.join(os.path.dirname(__file__), "assets/stills/nb")
os.makedirs(OUT, exist_ok=True)

LOCK = ("Consistent cartoon characters across all images. "
        "BANANA-MAN: a big yellow banana with a friendly cartoon face (eyes, eyebrows, mouth) "
        "and muscular human-like arms, wearing a black sleeveless tank top. "
        "TOMATO-WOMAN: a round red tomato with a cute cartoon face, green leaf hair on top, "
        "slim arms, wearing an elegant red dress. They are FRUIT, not humans. "
        "Style: high-quality 3D Pixar animation render, romantic candlelit restaurant, "
        "night city skyline through the window, warm cinematic bokeh lighting, vertical 9:16 portrait.")

SCENES = [
    ("01_setup", "Wide shot: the banana-man relaxes drinking from a beer bottle. "
                 "The tomato-woman sits close beside him looking annoyed and jealous."),
    ("02_throw", "The tomato-woman angrily snatches the beer bottle and throws it away, "
                 "beer splashing through the air. The banana-man looks shocked."),
    ("03_shock", "Close-up of the banana-man with a big shocked open-mouth face, "
                 "his hand empty, the beer is gone."),
    ("04_grab",  "The banana-man smirks confidently and pulls the tomato-woman into a warm "
                 "happy hug. Both are smiling, glowing love-hearts float around them."),
]

def call(prompt, ref_bytes=None):
    parts = [{"text": prompt}]
    if ref_bytes:
        parts.append({"inline_data": {"mime_type": "image/jpeg",
                      "data": base64.b64encode(ref_bytes).decode()}})
    body = json.dumps({"contents": [{"parts": parts}],
                       "generationConfig": {"responseModalities": ["IMAGE"],
                                            "imageConfig": {"aspectRatio": "9:16"}}}).encode()
    req = urllib.request.Request(URL, data=body, headers={"Content-Type": "application/json"})
    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                d = json.load(r)
            for p in d["candidates"][0]["content"]["parts"]:
                if "inlineData" in p:
                    return base64.b64decode(p["inlineData"]["data"])
                if "inline_data" in p:
                    return base64.b64decode(p["inline_data"]["data"])
            raise RuntimeError("geen image in response: " + json.dumps(d)[:300])
        except urllib.error.HTTPError as e:
            print("  HTTP", e.code, e.read().decode()[:200]); time.sleep(5)
    raise RuntimeError("3x gefaald")

# 1) hero referentie = scene 1
print("hero (01_setup)...")
hero = call(LOCK + " " + SCENES[0][1])
open(f"{OUT}/01_setup.jpg", "wb").write(hero)
print("  ok", len(hero), "bytes")

# 2) overige scenes met hero als referentie voor consistentie
for name, action in SCENES[1:]:
    print(f"{name} (met referentie)...")
    img = call(LOCK + " SAME two characters as in the reference image. " + action, ref_bytes=hero)
    open(f"{OUT}/{name}.jpg", "wb").write(img)
    print("  ok", len(img), "bytes")
print("KLAAR ->", OUT)