#!/usr/bin/env python3
"""Generate animated cell GIF using Nano Banana 2 (gemini-3.1-flash-image-preview)."""

import os, sys, base64, time, io, textwrap
from pathlib import Path
from PIL import Image, ImageFilter

API_KEY = "AIzaSyBM-NVkvNgTD5Jt_l-xhvFaSGCeaMKzB98"
OUTPUT_DIR = Path("/Users/bowang/.openclaw/workspace")

# ── Shared style block ──────────────────────────────────────────────────────
STYLE = textwrap.dedent("""
    Ultra-detailed 3D scientific illustration of a eukaryotic animal cell,
    hemispherical cross-section cut-away view revealing interior organelles.
    Dark teal-black background. Cinematic rim-lighting, dramatic multi-directional fill.
    Color palette: nucleus deep purple-blue, mitochondria rust-orange (visible cristae),
    rough ER magenta-pink with orange ribosomes, smooth ER cyan-blue tubular,
    Golgi violet-purple stacked cisternae lower-right quadrant, cytoskeleton bright
    green filaments radiating from nucleus, vesicles lavender and burnt-orange spheres,
    cell membrane teal-blue with visible protein channels.
    Subsurface scattering on membranes, molecular-level textures, depth-of-field blur
    on background neighboring cells. Photorealistic, no text, no labels, 16:9 aspect.
""").strip().replace("\n", " ")

# ── 6 animation frames ──────────────────────────────────────────────────────
FRAMES = [
    # 0 — rest
    f"{STYLE} Animation frame 1 of 6: RESTING STATE. "
    "Nucleus centered, normal size, soft glow. "
    "Mitochondria evenly distributed at original positions. "
    "Vesicles near Golgi, cytoskeleton calm, gentle ambient light.",

    # 1 — nucleus begins to pulse / mitochondria start drifting
    f"{STYLE} Animation frame 2 of 6: EARLY ACTIVATION. "
    "Nucleus slightly enlarged (5%), brighter inner glow, nuclear pores more visible. "
    "Two mitochondria have drifted slightly clockwise. "
    "A small vesicle budding from Golgi apparatus.",

    # 2 — peak pulse
    f"{STYLE} Animation frame 3 of 6: PEAK ACTIVITY. "
    "Nucleus at maximum pulse — noticeably larger with intense purple-blue radiance, "
    "nucleolus glowing brightly. "
    "Mitochondria drifted 20 degrees clockwise from starting positions. "
    "Several vesicles released from Golgi, moving outward. "
    "Cytoskeleton filaments shimmering with energy.",

    # 3 — active secretion
    f"{STYLE} Animation frame 4 of 6: ACTIVE SECRETION. "
    "Nucleus returning to normal size but still glowing. "
    "Mitochondria spread further clockwise — one near upper membrane. "
    "Vesicles visibly fusing with cell membrane at upper right, secreting contents. "
    "Cytoskeleton under dynamic tension.",

    # 4 — returning
    f"{STYLE} Animation frame 5 of 6: RETURNING TO REST. "
    "Nucleus back to normal size, glow fading. "
    "Mitochondria drifting back counterclockwise toward original positions. "
    "New small vesicles forming near ER. Cell membrane calm.",

    # 5 — rest again (close to frame 0 for seamless loop)
    f"{STYLE} Animation frame 6 of 6: REST (LOOP CLOSE). "
    "Nearly identical to frame 1 — resting state with very subtle membrane shimmer. "
    "Nucleus normal, mitochondria at original positions, Golgi calm. "
    "Soft warm glow over entire cell interior.",
]

def generate_frame(prompt: str, frame_idx: int) -> Image.Image | None:
    from google import genai
    from google.genai import types

    client = genai.Client(api_key=API_KEY)
    print(f"  Generating frame {frame_idx+1}/6...", flush=True)

    for attempt in range(3):
        try:
            response = client.models.generate_content(
                model="models/gemini-3.1-flash-image-preview",
                contents=prompt,
                config=types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"]
                ),
            )
            for part in response.candidates[0].content.parts:
                if part.inline_data and part.inline_data.data:
                    data = part.inline_data.data
                    if isinstance(data, str):
                        data = base64.b64decode(data)
                    img = Image.open(io.BytesIO(data)).convert("RGB")
                    print(f"    Frame {frame_idx+1}: {img.size} ✓", flush=True)
                    return img
            print(f"    Frame {frame_idx+1}: no image data in response", flush=True)
            return None
        except Exception as e:
            print(f"    Frame {frame_idx+1} attempt {attempt+1} failed: {e}", flush=True)
            if attempt < 2:
                time.sleep(5)
    return None


def main():
    print("🎬 Generating cell animation frames with Nano Banana 2...\n")
    frames = []

    for i, prompt in enumerate(FRAMES):
        img = generate_frame(prompt, i)
        if img:
            frames.append(img)
            path = OUTPUT_DIR / f"cell_frame_{i:02d}.png"
            img.save(path)
            print(f"    Saved: {path.name}")
        else:
            print(f"    SKIPPED frame {i+1}")
        time.sleep(2)  # be kind to the API

    if len(frames) < 2:
        print("Not enough frames generated. Aborting.")
        sys.exit(1)

    print(f"\n✅ Got {len(frames)} frames. Building animated GIF...")

    # Resize all frames to a consistent size (use first frame's size)
    target_w, target_h = frames[0].size
    # Keep aspect ratio, max width 1024 for GIF
    if target_w > 1024:
        scale = 1024 / target_w
        target_w = 1024
        target_h = int(target_h * scale)

    resized = [f.resize((target_w, target_h), Image.LANCZOS) for f in frames]

    # Forward + reverse for smooth loop (frames 0-5 then 4-1)
    loop_frames = resized + resized[-2:0:-1]

    gif_path = OUTPUT_DIR / "cell-animation-nano-banana-2.gif"
    loop_frames[0].save(
        gif_path,
        save_all=True,
        append_images=loop_frames[1:],
        duration=600,       # ms per frame
        loop=0,             # infinite loop
        optimize=False,
    )

    size_mb = gif_path.stat().st_size / 1_000_000
    print(f"🎉 GIF saved: {gif_path.name} ({size_mb:.1f} MB, {len(loop_frames)} frames)")
    print(f"   Full path: {gif_path}")


if __name__ == "__main__":
    main()