#!/usr/bin/env python3
"""
Generate object reference sheets using NanoBanana.

Creates consistent reference images for props/objects used across scenes.

Usage:
    python generate_object_sheet.py -d "Old leather folder with label 'АЛЯСКА 1867'" -o obj_folder.jpg
    python generate_object_sheet.py -d "Treaty document with red wax seal" -o obj_document.jpg --views 4
"""

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

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

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


PROMPT_OBJECT_SHEET = """Create a professional object reference sheet of {description}.

Use a clean, neutral white/grey background. Present as technical product photography showing the object from multiple angles.

Arrange the composition showing:
- Main view (largest, center): The object in its most recognizable angle
- Top view: Looking down at the object
- Side view: Profile of the object
- Detail view: Close-up of key identifying features (text, seal, texture)

Maintain perfect consistency across all views - same object, same colors, same details, same text/labels.
Lighting should be even and professional, no dramatic shadows.
Output crisp, detailed reference sheet.

IMPORTANT: Any text on the object must be exactly as described - do not change or omit text."""


def call_nanobanana(
    api_key: str,
    prompt: str,
    reference_image: bytes = None,
    aspect_ratio: str = "16:9",
) -> tuple:
    """Call NanoBanana to generate an image."""
    parts = []

    if reference_image:
        ref_b64 = base64.b64encode(reference_image).decode('utf-8')
        parts.append({"inline_data": {"mime_type": "image/jpeg", "data": ref_b64}})

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

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

    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} - {response.text[:300]}")
        return None, None

    result = response.json()

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

    return None, None


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

    args = parser.parse_args()

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

    prompt = PROMPT_OBJECT_SHEET.format(description=args.description)

    print(f"Generating object sheet...")
    print(f"Description: {args.description[:80]}...")

    ref_image = None
    if args.reference:
        with open(args.reference, "rb") as f:
            ref_image = f.read()
        print(f"Using reference: {args.reference}")

    img_data, img_mime = call_nanobanana(
        api_key, prompt, ref_image, args.aspect
    )

    if img_data:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as f:
            f.write(img_data)
        print(f"Saved: {output_path} ({len(img_data)} bytes)")
    else:
        print("Failed to generate")
        sys.exit(1)


if __name__ == "__main__":
    main()
