#!/usr/bin/env python3
"""Both scenes at 1:1 for Twitter posts - higher quality prompts."""

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 = {
    "board": """A dramatic square photograph that looks like a still from a Wes Anderson stop-motion film. A dark detective office at night. Three HUMAN-SIZED handmade felt Muppet puppet characters stand before a conspiracy corkboard covering the wall behind them.

Center: THE HUNTER — a puppet in a camel felt trench coat and brown fedora, reaching up to pin an index card to the board. Intense, determined expression.

Left: THE SKEPTIC — a puppet in a black oversized hoodie, arms tightly crossed against their chest, weight shifted to one hip. Bright teal yarn hair. Looking sideways at the board with visible doubt. One eyebrow raised.

Right: THE REFEREE — a puppet in a crisp black-and-white vertically striped referee jersey, silver whistle on a lanyard around their neck. Studying the board with one hand raised to their chin. About to make a call.

The corkboard is covered in index cards, printed code, error logs, and red string connecting everything. A cute small fuzzy felt bug creature with googly eyes is caught in the web of red string in the center of the board.

Warm desk lamp on the left casts golden light. The rest of the room is moody dark blue. Coffee cups, scattered papers, a corded telephone on the desk. Every puppet has richly textured felt skin with visible fabric grain, yarn hair, and round bead eyes. The felt textures should be so detailed you can almost feel the fuzz. Cinematic shallow depth of field. Square 1:1 format.""",

    "catch": """A joyful square photograph that looks like a still from a Wes Anderson stop-motion film. A bright modern startup office flooded with warm golden afternoon light from large windows.

Center: THE HUNTER — a HUMAN-SIZED felt Muppet puppet in a camel trench coat and fedora, triumphantly holding up a glass mason jar at eye level. Inside the jar: a small adorable fuzzy felt bug creature with big googly eyes and tiny antennae, pressing its face against the glass looking bewildered. The hunter's expression is proud and satisfied.

Left: THE SKEPTIC — a HUMAN-SIZED felt Muppet puppet in a black oversized hoodie, arms crossed, leaning forward to peer at the jar with a suspicious squint. Bright teal yarn hair. Expression says "I'm not convinced that's the right bug." Dubious body language.

Right: THE REFEREE — a HUMAN-SIZED felt Muppet puppet in a black-and-white vertically striped referee jersey with a silver whistle. One arm raised making an official signal — pointing at the jar, confirming the catch is valid. Decisive expression.

The office has standing desks with sticker-covered laptops, sticky notes, energy drink cans. The puppets are FULL HUMAN SIZE standing on the floor. Every puppet has richly textured felt skin with visible fabric grain, colorful yarn hair, and round bead eyes. Cinematic shallow depth of field. Square 1:1 format."""
}

for name, prompt in scenes.items():
    print(f"\nGenerating {name}...")
    payload = {
        "instances": [{"prompt": prompt}],
        "parameters": {
            "sampleCount": 3,
            "aspectRatio": "1:1",
            "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=180) as resp:
            result = json.loads(resp.read())
    except urllib.error.HTTPError as e:
        print(f"  HTTP {e.code}: {e.read().decode()[: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-tw-{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!")