#!/usr/bin/env python3
"""Edit existing images using Gemini API.

Usage:
    python edit_image.py input.jpg "edit instruction" output.jpg [--model MODEL]

Environment:
    GEMINI_API_KEY - Required API key
"""

import argparse
import os
import sys

from PIL import Image
from google import genai
from google.genai import types


def edit_image(
    input_path: str,
    instruction: str,
    output_path: str,
    model: str = "gemini-2.5-flash-image",
    aspect_ratio: str | None = None,
    image_size: str | None = None,
) -> str | None:
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise EnvironmentError("GEMINI_API_KEY environment variable not set")

    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input image not found: {input_path}")

    client = genai.Client(api_key=api_key)
    input_image = Image.open(input_path)

    config_kwargs = {"response_modalities": ["TEXT", "IMAGE"]}

    image_config_kwargs = {}
    if aspect_ratio:
        image_config_kwargs["aspect_ratio"] = aspect_ratio
    if image_size:
        image_config_kwargs["image_size"] = image_size

    if image_config_kwargs:
        config_kwargs["image_config"] = types.ImageConfig(**image_config_kwargs)

    config = types.GenerateContentConfig(**config_kwargs)

    response = client.models.generate_content(
        model=model,
        contents=[instruction, input_image],
        config=config,
    )

    text_response = None
    image_saved = False

    for part in response.parts:
        if part.text is not None:
            text_response = part.text
        elif part.inline_data is not None:
            image = part.as_image()
            image.save(output_path)
            image_saved = True

    if not image_saved:
        raise RuntimeError("No image was generated.")

    return text_response


def main():
    parser = argparse.ArgumentParser(description="Edit images using Gemini API")
    parser.add_argument("input", help="Input image path")
    parser.add_argument("instruction", help="Edit instruction")
    parser.add_argument("output", help="Output file path (.jpg recommended)")
    parser.add_argument("--model", "-m", default="gemini-2.5-flash-image",
                        choices=["gemini-2.5-flash-image", "gemini-3-pro-image-preview"])
    parser.add_argument("--aspect", "-a",
                        choices=["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"])
    parser.add_argument("--size", "-s", choices=["1K", "2K", "4K"])

    args = parser.parse_args()

    try:
        text = edit_image(args.input, args.instruction, args.output, args.model, args.aspect, args.size)
        print(f"Edited image saved to: {args.output}")
        if text:
            print(f"Model response: {text}")
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()