#!/usr/bin/env python3
"""Generate refined evidence-board and the-catch scenes."""

import base64
import json
import os
import sys
import urllib.request

API_KEY = os.environ.get("GEMINI_API_KEY")
model = "imagen-4.0-generate-001"
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:predict?key={API_KEY}"

scenes = {
    "evidence-board-v2": """Cinematic photograph of a dark moody office at night. Three handmade felt Muppet-style puppet characters investigate a large conspiracy-style corkboard covering most of the wall. The composition is carefully arranged: a warm vintage desk lamp sits on the FAR LEFT of the frame on a desk, casting warm amber light across the scene. The RIGHT SIDE and UPPER RIGHT of the image is darker, with clean dark wall space — intentionally left as negative space for text overlay.

The puppet characters: one wears a grey hoodie and thick-rimmed glasses, pointing at the corkboard. One has bright colorful yarn hair and wears a flannel shirt, holding a large coffee mug. One wears a classic felt trench coat and fedora like a film noir detective, pinning a new index card to the board.

The corkboard has index cards with code snippets, error messages, and printed stack traces, all connected with red string in a conspiracy pattern. A small round fuzzy furry felt BUG creature — like a fluffy caterpillar with big innocent googly eyes, tiny antennae, and little felt legs — is tangled in the red string at the center of the board, looking confused and cute.

All puppets have felt/fleece skin, yarn hair, bead eyes, Muppet-style hinged mouths. The environment is photorealistic — real desk, real coffee cups, real papers, corded telephone. Blue-dark noir atmosphere with the single warm lamp as key light. Shallow depth of field. Wide 16:9 cinematic composition.""",

    "the-catch-v2": """Cinematic photograph of a lived-in, slightly messy modern startup office — real desks with sticker-covered laptops, energy drink cans, tangled headphone cables, sticky notes everywhere. Late afternoon golden window light streaming in.

Three handmade felt Muppet-style puppet characters in a spontaneous moment of triumph. The detective puppet (felt trench coat, fedora, tie) stands on a desk holding up a glass mason jar at eye level — inside is a small round fuzzy furry felt BUG creature with big googly eyes and tiny antennae, pressing its face against the glass looking bewildered. 

One dev puppet in a black oversized hoodie with headphones around their neck is mid-leap from their chair, fist in the air. Another dev puppet with a beanie and band t-shirt leans back in their chair grinning, arms crossed in satisfaction.

All puppets have felt/fleece skin, yarn hair, bead eyes, Muppet hinged mouths. The upper portion of the image has lighter wall/ceiling space suitable for dark text overlay. The scene feels candid and real — like a photo someone snapped at the exact right moment. Warm, natural, not overlit. Shallow depth of field, slight lens flare from the window. Wide 16:9."""
}

for name, prompt in scenes.items():
    print(f"\nGenerating {name}...")
    payload = {
        "instances": [{"prompt": prompt}],
        "parameters": {
            "sampleCount": 2,
            "aspectRatio": "16:9",
            "personGeneration": "allow_adult"
        }
    }
    data = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode()
        print(f"  HTTP {e.code}: {body[:300]}")
        continue

    for i, pred in enumerate(result.get("predictions", [])):
        if "bytesBase64Encoded" in pred:
            img_data = base64.b64decode(pred["bytesBase64Encoded"])
            out = f"/root/.openclaw/workspace/output/bug-hunt-{name}-{i+1}.png"
            os.makedirs(os.path.dirname(out), exist_ok=True)
            with open(out, "wb") as f:
                f.write(img_data)
            print(f"  Saved: {out} ({len(img_data)} bytes)")

print("\nDone!")