#!/usr/bin/env python3
"""
Generate single images using NanoBanana (Gemini 3 Pro).

For Ken Burns effect scenes where we need a single photo-style image,
not a full character turnaround sheet.

Usage:
    python generate_single_image.py -d "Description of image" -o output.jpg -a "9:16"
"""

import os
import sys
import argparse
import base64
import requests
from pathlib import Path

# Load .env from parent directory
from dotenv import load_dotenv
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)

BASE_URL = "https://api.laozhang.ai/v1beta"


def call_nanobanana(
    api_key: str,
    prompt: str,
    image_data: bytes = None,
    mime_type: str = "image/jpeg",
    aspect_ratio: str = "9:16",
) -> tuple:
    """
    Call NanoBanana Pro (Gemini 3 Pro) via LaoZhang.ai.
    """
    parts = []

    if image_data:
        image_base64 = base64.b64encode(image_data).decode('utf-8')
        parts.append({"inline_data": {"mime_type": mime_type, "data": image_base64}})

    parts.append({"text": prompt})

    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
                "imageSize": "1K"
            }
        }
    }

    print(f"Calling NanoBanana Pro via LaoZhang.ai...")
    print(f"Aspect ratio: {aspect_ratio}")

    response = requests.post(
        f"{BASE_URL}/models/gemini-3-pro-image-preview:generateContent",
        headers={
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        },
        json=payload,
        timeout=120
    )

    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(f"Response: {response.text}")
        return None, None

    result = response.json()

    if "candidates" not in result:
        print(f"No candidates in response: {result}")
        return None, None

    for candidate in result.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            if "inlineData" in part:
                data = part["inlineData"]
                image_bytes = base64.b64decode(data["data"])
                return image_bytes, data.get("mimeType", "image/png")

    print("No image data in response")
    return None, None


def main():
    parser = argparse.ArgumentParser(description="Generate single image")
    parser.add_argument("--description", "-d", required=True, help="Image description")
    parser.add_argument("--output", "-o", required=True, help="Output file path")
    parser.add_argument("--aspect", "-a", default="9:16", help="Aspect ratio (default: 9:16)")
    parser.add_argument("--reference", "-r", help="Optional reference image")

    args = parser.parse_args()

    api_key = os.getenv("LAOZHANG_KEY")
    if not api_key:
        print("Error: Set LAOZHANG_KEY environment variable")
        sys.exit(1)

    image_data = None
    mime_type = None
    if args.reference:
        print(f"Using reference image: {args.reference}")
        with open(args.reference, "rb") as f:
            image_data = f.read()
        mime_type = "image/jpeg"
        if args.reference.lower().endswith(".png"):
            mime_type = "image/png"

    print(f"Generating: {args.description[:100]}...")

    result_bytes, result_mime = call_nanobanana(
        api_key=api_key,
        prompt=args.description,
        image_data=image_data,
        mime_type=mime_type,
        aspect_ratio=args.aspect,
    )

    if result_bytes:
        ext = "png"
        if result_mime == "image/jpeg":
            ext = "jpg"
        elif result_mime == "image/webp":
            ext = "webp"

        output_path = Path(args.output)
        if not output_path.suffix:
            output_path = output_path.with_suffix(f".{ext}")

        with open(output_path, "wb") as f:
            f.write(result_bytes)

        print(f"Image saved to: {output_path}")
        print(f"Size: {len(result_bytes)} bytes")
    else:
        print("Failed to generate image")
        sys.exit(1)


if __name__ == "__main__":
    main()
