#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "google-genai>=1.0.0",
#     "pillow>=10.0.0",
# ]
# ///
"""
Generate images using imagen-4.0-ultra-generate-001 (Nano Banana Pro 2).
"""
import argparse
import os
import sys
from pathlib import Path
from PIL import Image
import io

RESOLUTION_MAP = {"1K": 1024, "2K": 2048, "4K": 4096}
ASPECT_RATIOS = ["1:1", "3:4", "4:3", "9:16", "16:9"]


def get_api_key(provided_key=None):
    return provided_key or os.environ.get("GEMINI_API_KEY")


def main():
    parser = argparse.ArgumentParser(description="Generate images via imagen-4.0-ultra (Nano Banana Pro 2)")
    parser.add_argument("--prompt", "-p", required=True)
    parser.add_argument("--filename", "-f", required=True)
    parser.add_argument("--aspect-ratio", "-a", default="16:9", choices=ASPECT_RATIOS)
    parser.add_argument("--resolution", "-r", default="2K", choices=["1K", "2K", "4K"])
    parser.add_argument("--api-key", "-k", default=None)
    parser.add_argument("--no-media-line", action="store_true")
    args = parser.parse_args()

    api_key = get_api_key(args.api_key)
    if not api_key:
        print("ERROR: No GEMINI_API_KEY found.", file=sys.stderr)
        sys.exit(1)

    import google.genai as genai
    from google.genai import types

    client = genai.Client(api_key=api_key)

    print(f"Generating with imagen-4.0-ultra | {args.aspect_ratio} | {args.resolution}...")

    result = client.models.generate_images(
        model="imagen-4.0-ultra-generate-001",
        prompt=args.prompt,
        config=types.GenerateImagesConfig(
            number_of_images=1,
            aspect_ratio=args.aspect_ratio,
        )
    )

    if not result.generated_images:
        print("ERROR: No images returned.", file=sys.stderr)
        sys.exit(1)

    img_bytes = result.generated_images[0].image.image_bytes
    img = Image.open(io.BytesIO(img_bytes))

    # Resize to target resolution
    target = RESOLUTION_MAP[args.resolution]
    w, h = img.size
    if max(w, h) < target:
        scale = target / max(w, h)
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)

    out_path = Path(args.filename)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    img.save(out_path, "PNG", optimize=False)
    print(f"Image saved: {out_path}")

    if not args.no_media_line:
        print(f"MEDIA:{out_path}")


if __name__ == "__main__":
    main()