#!/usr/bin/env python3
"""Generate stills with Nano Banana (Gemini image models) from a JSON list of prompts.

  python3 gen-images.py prompts.json --out stills/ [--model gemini-3-pro-image] [--aspect 16:9] [--size 2K] [names...]

prompts.json: [["name", "prompt"], ...]. Existing files are skipped, so re-running only fills gaps.
Needs GEMINI_API_KEY (or GOOGLE_API_KEY). Images from the API carry no visible watermark
(they do carry SynthID, which you should leave alone).

Models: gemini-3-pro-image (Nano Banana Pro, best detail; default), gemini-3.1-flash-image
(Nano Banana 2, faster and cheaper), gemini-2.5-flash-image (original Nano Banana).
"""
import argparse, base64, concurrent.futures, glob, json, os, sys
from common import key, http_json

ap = argparse.ArgumentParser()
ap.add_argument("prompts"); ap.add_argument("names", nargs="*")
ap.add_argument("--out", default="stills"); ap.add_argument("--model", default="gemini-3-pro-image")
ap.add_argument("--aspect", default="16:9"); ap.add_argument("--size", default="2K"); ap.add_argument("--parallel", type=int, default=4)
a = ap.parse_args()
K = key("GEMINI_API_KEY", "GOOGLE_API_KEY")
os.makedirs(a.out, exist_ok=True)
shots = [s for s in json.load(open(a.prompts)) if not a.names or s[0] in a.names]


def gen(item):
    name, prompt = item
    if glob.glob(f"{a.out}/{name}.*"):
        return f"{name}: exists"
    body = {"contents": [{"parts": [{"text": prompt}]}],
            "generationConfig": {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": a.aspect, "imageSize": a.size}}}
    try:
        d = http_json(f"https://generativelanguage.googleapis.com/v1beta/models/{a.model}:generateContent", body,
                      {"x-goog-api-key": K}, timeout=240, wait=15)
    except RuntimeError as e:
        return f"{name}: {e}"
    for part in d.get("candidates", [{}])[0].get("content", {}).get("parts", []):
        if "inlineData" in part:
            ext = "png" if "png" in part["inlineData"]["mimeType"] else "jpg"
            data = base64.b64decode(part["inlineData"]["data"])
            open(f"{a.out}/{name}.{ext}", "wb").write(data)
            return f"{name}: ok {len(data) // 1024} KB"
    return f"{name}: no image in the response ({json.dumps(d)[:200]})"


with concurrent.futures.ThreadPoolExecutor(a.parallel) as ex:
    for r in ex.map(gen, shots):
        print(r, flush=True)