import json, os, sys, base64, time, urllib.request, urllib.error
D = os.path.expanduser("~/mnt/KIOXIA/r2r2r-data/videos/_ショート専用/AIあるある_背景")
KEY = open(os.path.expanduser("~/mnt/KIOXIA/gemini_api_key.txt")).read().strip()
MODEL = "gemini-3.1-flash-image"
COMMON = (" Photorealistic photograph, vertical 9:16 composition. Absolutely no text, no letters, no numbers, no logos, no signs with writing, no people."
          " Cinematic, slightly moody lighting, rich colors. Keep the lower third of the frame calm and uncluttered (characters will be placed there).")
JOBS = {
 "02": "A cozy home recording desk at night: studio headphones, a microphone on a boom arm, a laptop glowing, warm lamp light.",
 "03": "A vintage printing press in a dim workshop, freshly printed blank sheets stacked high, ink rollers, dramatic side light.",
 "05": "A wall of thirty identical retro television sets stacked in a grid, every screen glowing with the same plain blue light, in a dark warehouse.",
 "06": "An extreme close-up of a large magnifying glass over a sheet of paper covered in microscopic unreadable marks, on a wooden desk.",
 "07": "New York City skyline at dusk seen from the water, with the Statue of Liberty silhouette, glowing orange sky.",
 "09": "A surreal crossroads in a desert with dozens of blank wooden arrow signposts pointing in every possible direction, blue sky.",
 "11": "An office paper shredder overflowing with long curly strips of shredded paper piled on the floor, harsh fluorescent light.",
 "12": "A professional voice recording booth: a studio condenser microphone with a pop filter, acoustic foam walls, headphones hanging.",
 "13": "An artist studio with a wooden easel holding a blank canvas, paint brushes and palettes, soft window light.",
 "14": "An empty comedy club stage with a single microphone stand under a bright spotlight, red brick wall, dark audience area.",
}
want = sys.argv[1:] or list(JOBS)
for k in want:
    out = os.path.join(D, f"bg_{k}.png")
    if os.path.exists(out):
        continue
    body = {"contents": [{"role": "user", "parts": [{"text": JOBS[k] + COMMON}]}],
            "generationConfig": {"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": "9:16", "imageSize": "1K"}}}
    req = urllib.request.Request(f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent",
                                 data=json.dumps(body).encode(), headers={"Content-Type": "application/json", "x-goog-api-key": KEY})
    for a in range(3):
        try:
            res = json.load(urllib.request.urlopen(req, timeout=170)); break
        except Exception as e:
            print(k, "retry", e, flush=True); time.sleep(15)
    else:
        print(k, "失敗", flush=True); continue
    img = [p["inlineData"] for c in res.get("candidates", []) for p in c.get("content", {}).get("parts", []) if "inlineData" in p]
    if img:
        open(out, "wb").write(base64.b64decode(img[0]["data"])); print(k, "ok", flush=True)
    else:
        print(k, "画像なし", flush=True)
print("おわり", flush=True)