#!/usr/bin/env python3
"""The Catch 1:1 - Imagen with precise character descriptions from reference."""

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}"

prompt = """A square photograph that looks like a still from a high-budget Muppet movie. Three HUMAN-SIZED felt Muppet puppet characters in a modern startup office, golden afternoon light streaming through windows.

EXACT CHARACTER DESCRIPTIONS:

LEFT — THE SKEPTIC: A Muppet puppet with wild messy ORANGE yarn hair sticking up, wearing a BLACK HOODIE with white drawstrings. Arms firmly CROSSED against chest. Frowning expression — one eyebrow lowered, mouth turned down. Big round WHITE eyes with black pupils, round ORANGE felt nose. Peach/tan felt skin. He's leaning slightly toward the jar but his body language screams doubt.

CENTER — THE HUNTER: A Muppet puppet wearing a BROWN FEDORA hat and a CAMEL/BEIGE BELTED TRENCH COAT with collar turned up. Brown yarn hair under the hat. Holding up a clear GLASS MASON JAR at eye level with both hands. Inside the jar is a small adorable FUZZY GREY-GREEN CATERPILLAR BUG puppet with oversized googly eyes looking confused. The hunter has a PROUD GRIN — classic Muppet wide mouth smile. Round ORANGE felt nose. Tan felt skin.

RIGHT — THE REFEREE: A Muppet puppet wearing a BLACK BASEBALL CAP and a BLACK-AND-WHITE VERTICALLY STRIPED REFEREE JERSEY. A SILVER WHISTLE on a black lanyard hangs around his neck. One ARM RAISED HIGH making an official call/signal. Cheerful decisive expression. Round RED/ORANGE felt nose. Tan felt skin.

All three have the classic Muppet look: round heads, wide hinged mouths, felt/fleece skin with visible fabric texture, big spherical white eyes with black pupils. They are chest-up to waist-up in frame, filling the square composition. The jar with the bug is the focal point in the center. Warm office background with monitors showing colorful code, sticky notes, plants. Cinematic shallow depth of field. Rich felt textures you can almost touch."""

payload = {
    "instances": [{"prompt": prompt}],
    "parameters": {
        "sampleCount": 4,
        "aspectRatio": "1:1",
        "personGeneration": "allow_adult"
    }
}

data = json.dumps(payload).encode()
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})

print("Generating catch 1:1 with exact character match...")
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]}")
    sys.exit(1)

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-catch-exact-{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("Done!")