#!/usr/bin/env python3
"""
Generate NxN storyboard grid using NanoBanana (Gemini 3 Pro) via LaoZhang.ai.

Usage:
    python generate_grid.py --prompt prompt.txt --output-dir /path/to/output [--reference ref.jpg] [--layout 4x4]
    python generate_grid.py -p prompt.txt -o /path/to/output -r ref.jpg -l 4x4
"""

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

# Add scripts dir to path for laozhang_proxy import
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

from dotenv import load_dotenv
import requests

# Load environment
load_dotenv(Path(__file__).parent.parent.parent / ".env")

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


def generate_grid(
    api_key: str,
    prompt: str,
    output_dir: Path,
    reference_image: Path = None,
    layout: str = "3x3",
    aspect_ratio: str = "9:16",
    image_size: str = "4K"
) -> Path:
    """
    Generate NxN grid using NanoBanana with optional reference image.

    Args:
        api_key: LaoZhang API key
        prompt: Grid generation prompt text
        output_dir: Directory to save output
        reference_image: Optional reference image for character consistency
        layout: Grid layout (e.g., "3x3", "4x4")
        aspect_ratio: Output aspect ratio
        image_size: Output size (1K, 2K, 4K)

    Returns:
        Path to generated grid image
    """
    print(f"Generating {layout} grid with NanoBanana...")
    print(f"  Prompt length: {len(prompt)} chars")
    print(f"  Aspect ratio: {aspect_ratio}")
    print(f"  Size: {image_size}")

    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)

    # Build parts list
    parts = []

    # Add reference image if provided
    if reference_image and reference_image.exists():
        print(f"  Reference: {reference_image}")
        with open(reference_image, 'rb') as f:
            ref_data = f.read()

        # Determine mime type
        if str(reference_image).lower().endswith('.png'):
            mime = 'image/png'
        elif str(reference_image).lower().endswith('.webp'):
            mime = 'image/webp'
        else:
            mime = 'image/jpeg'

        ref_base64 = base64.b64encode(ref_data).decode('utf-8')
        parts.append({"inline_data": {"mime_type": mime, "data": ref_base64}})

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

    # Call NanoBanana
    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
                "imageSize": image_size
            }
        }
    }

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

    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(response.text[:500])
        raise RuntimeError(f"NanoBanana API error: {response.status_code}")

    data = response.json()

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

            # Determine extension
            ext = ".png" if "png" in mime_type else ".jpg"

            # Save grid image
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            grid_path = output_dir / f"grid_{layout}_{timestamp}{ext}"

            with open(grid_path, 'wb') as f:
                f.write(image_bytes)

            print(f"  Saved: {grid_path} ({len(image_bytes) // 1024}KB)")
            return grid_path

    raise RuntimeError("No image in response")


def main():
    parser = argparse.ArgumentParser(description="Generate NxN storyboard grid with NanoBanana")
    parser.add_argument("--prompt", "-p", required=True, help="Prompt text or path to .txt file")
    parser.add_argument("--output-dir", "-o", required=True, help="Output directory")
    parser.add_argument("--reference", "-r", help="Reference image for character consistency")
    parser.add_argument("--layout", "-l", default="3x3", help="Grid layout (3x3, 4x4, etc.)")
    parser.add_argument("--aspect-ratio", "-a", default="9:16", help="Output aspect ratio")
    parser.add_argument("--size", "-s", default="4K", choices=["1K", "2K", "4K"], help="Output size")

    args = parser.parse_args()

    api_key = os.environ.get("LAOZHANG_KEY")
    if not api_key:
        print("Error: LAOZHANG_KEY not set in environment")
        sys.exit(1)

    # Get prompt
    if args.prompt.endswith('.txt') and Path(args.prompt).exists():
        with open(args.prompt) as f:
            prompt = f.read().strip()
    else:
        prompt = args.prompt

    # Paths
    output_dir = Path(args.output_dir)
    reference = Path(args.reference) if args.reference else None

    try:
        grid_path = generate_grid(
            api_key=api_key,
            prompt=prompt,
            output_dir=output_dir,
            reference_image=reference,
            layout=args.layout,
            aspect_ratio=args.aspect_ratio,
            image_size=args.size
        )
        print(f"\nSuccess! Grid: {grid_path}")
    except Exception as e:
        print(f"\nError: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
