#!/usr/bin/env python3
"""
FF image generator with model fallback.
Tries gemini-3-pro-image-preview first, falls back to gemini-2.0-flash-preview-image-generation.
Both support --input-image for logo embedding.
"""
import argparse
import sys
import os
from pathlib import Path
from google import genai
from google.genai import types
from PIL import Image as PILImage
from io import BytesIO

MODELS = [
    "gemini-3-pro-image-preview",
    "gemini-2.5-flash-image",
    "gemini-3.1-flash-image-preview",
]

def build_contents(prompt, logo_path):
    parts = []
    if logo_path and Path(logo_path).exists():
        img = PILImage.open(logo_path)
        buf = BytesIO()
        img.save(buf, format='PNG')
        parts.append(types.Part.from_bytes(data=buf.getvalue(), mime_type='image/png'))
    parts.append(types.Part.from_text(text=prompt))
    return [types.Content(role="user", parts=parts)]

def try_generate(client, model, contents, aspect_ratio, output_path):
    image_cfg_kwargs = {}
    if aspect_ratio:
        image_cfg_kwargs["aspect_ratio"] = aspect_ratio

    response = client.models.generate_content(
        model=model,
        contents=contents,
        config=types.GenerateContentConfig(
            response_modalities=["TEXT", "IMAGE"],
            image_config=types.ImageConfig(**image_cfg_kwargs) if image_cfg_kwargs else None
        )
    )

    for part in response.parts:
        if part.inline_data is not None:
            image_data = part.inline_data.data
            if isinstance(image_data, str):
                import base64
                image_data = base64.b64decode(image_data)
            image = PILImage.open(BytesIO(image_data))
            if image.mode == 'RGBA':
                rgb = PILImage.new('RGB', image.size, (255, 255, 255))
                rgb.paste(image, mask=image.split()[3])
                rgb.save(str(output_path), 'PNG')
            else:
                image.convert('RGB').save(str(output_path), 'PNG')
            return True
    return False

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--prompt', '-p', required=True)
    parser.add_argument('--filename', '-f', required=True)
    parser.add_argument('--input-image', '-i', default=None)
    parser.add_argument('--aspect-ratio', '-a', default='16:9')
    parser.add_argument('--api-key', '-k', default=None)
    args = parser.parse_args()

    api_key = args.api_key or os.environ.get('GEMINI_API_KEY')
    if not api_key:
        print("Error: no API key", file=sys.stderr)
        sys.exit(1)

    client = genai.Client(api_key=api_key)
    contents = build_contents(args.prompt, args.input_image)
    output_path = Path(args.filename)

    last_err = None
    for model in MODELS:
        try:
            print(f"Trying model: {model}")
            ok = try_generate(client, model, contents, args.aspect_ratio, output_path)
            if ok:
                print(f"Success with {model}")
                sys.exit(0)
            else:
                print(f"No image in response from {model}", file=sys.stderr)
        except Exception as e:
            last_err = e
            print(f"Model {model} failed: {e}", file=sys.stderr)
            continue

    print(f"All models failed. Last error: {last_err}", file=sys.stderr)
    sys.exit(1)

if __name__ == '__main__':
    main()