# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "google-genai",
#     "pillow",
#     "truststore>=0.10",
# ]
# ///
"""Wrapper that injects truststore (Windows OS cert store) before calling Gemini.
Use to bypass corporate-cert/anti-virus HTTPS interception in environments where the
default certifi bundle does not include the intercepting CA."""

import truststore
truststore.inject_into_ssl()

import argparse
import os
import sys
from google import genai
from google.genai import types
from PIL import Image
import io

MODEL_IDS = {
    "flash": "gemini-2.5-flash-image",
    "pro": "gemini-3-pro-image-preview",
    "2": "gemini-3.1-flash-image-preview",
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--model", default="2", choices=MODEL_IDS.keys())
    parser.add_argument("--size", default="2K")
    parser.add_argument("--aspect", default="16:9")
    args = parser.parse_args()

    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        print("ERROR: GEMINI_API_KEY not set", file=sys.stderr)
        sys.exit(1)

    client = genai.Client(api_key=api_key)
    model_id = MODEL_IDS[args.model]

    cfg = types.GenerateContentConfig(
        response_modalities=["IMAGE"],
        image_config=types.ImageConfig(
            aspect_ratio=args.aspect,
            image_size=args.size,
        ),
    )

    print(f"Calling {model_id} aspect={args.aspect} size={args.size}...", flush=True)
    response = client.models.generate_content(
        model=model_id,
        contents=args.prompt,
        config=cfg,
    )

    for cand in response.candidates:
        for part in cand.content.parts:
            if part.inline_data and part.inline_data.data:
                img = Image.open(io.BytesIO(part.inline_data.data))
                ext = os.path.splitext(args.output)[1].lower()
                if ext == ".webp":
                    img.save(args.output, "WEBP", quality=90, method=6)
                elif ext in (".jpg", ".jpeg"):
                    img.convert("RGB").save(args.output, "JPEG", quality=90)
                else:
                    img.save(args.output, "PNG")
                print(f"OK saved {args.output} ({img.size[0]}x{img.size[1]})", flush=True)
                return
    print("ERROR no image data in response", file=sys.stderr)
    sys.exit(2)


if __name__ == "__main__":
    main()