#!/usr/bin/env python3
"""
Generate character reference sheets using NanoBanana (Gemini 3 Pro).

Usage:
    # From description only:
    python generate_character_sheet.py -d "a 45-year-old male historian, glasses, tweed jacket" -o char_historian.jpg

    # From reference image:
    python generate_character_sheet.py -r ref_image.jpg -o char_from_ref.jpg

    # With name and sepia style:
    python generate_character_sheet.py -r tsar_ref.jpg -o char_tsar.jpg -n "Tsar Alexander II" --sepia
"""

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"

# Tested prompt for character sheet with reference image
PROMPT_WITH_REF = """Create a professional character reference sheet based strictly on the uploaded
reference image{name_suffix}. Use a clean, neutral plain background and present the sheet
as a technical model turnaround while matching the exact visual style of the
reference (same realism level, rendering approach, texture, color treatment,
and overall aesthetic).{sepia_note}

Arrange the composition into two horizontal rows.
Top row: four full-body standing views placed side-by-side in this order:
front view, left profile view (facing left), right profile view (facing right), back view.
Bottom row: three highly detailed close-up portraits aligned beneath the full-body row
in this order: front portrait, left profile portrait (facing left), right profile portrait (facing right).

Maintain perfect identity consistency across every panel. Keep the subject in a
relaxed A-pose and with consistent scale and alignment between views, accurate
anatomy, and clear silhouette; ensure even spacing and clean panel separation,
with uniform framing and consistent head height across the full-body lineup and
consistent facial scale across the portraits. Lighting should be consistent across
all panels (same direction, intensity, and softness), with natural, controlled
shadows that preserve detail without dramatic mood shifts. Output a crisp,
print-ready reference sheet look, sharp details."""

# Tested prompt for character sheet from description only
PROMPT_DESCRIPTION_TEMPLATE = """Create a professional character reference sheet of {description}.
Use a clean, neutral plain background and present the sheet as a technical model
turnaround in a photographic style.

Arrange the composition into two horizontal rows.
Top row: four full-body standing views placed side-by-side in this order:
front view, left profile view (facing left), right profile view (facing right), back view.
Bottom row: three highly detailed close-up portraits aligned beneath the full-body row
in this order: front portrait, left profile portrait (facing left), right profile portrait (facing right).

Maintain perfect identity consistency across every panel. Keep the subject in a
relaxed A-pose and with consistent scale and alignment between views, accurate
anatomy, and clear silhouette; ensure even spacing and clean panel separation,
with uniform framing and consistent head height across the full-body lineup and
consistent facial scale across the portraits. Lighting should be consistent across
all panels (same direction, intensity, and softness), with natural, controlled
shadows that preserve detail without dramatic mood shifts. Output a crisp,
print-ready reference sheet look, sharp details."""


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

    Args:
        api_key: LaoZhang API key
        prompt: Generation instructions
        image_data: Optional input image bytes
        mime_type: Input image MIME type
        aspect_ratio: Output aspect ratio

    Returns:
        Tuple of (image_bytes, mime_type)
    """
    # Build parts list
    parts = []

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

    # Add prompt
    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()

    # Extract image from response
    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 character reference sheet")
    parser.add_argument("--reference", "-r", help="Reference image file")
    parser.add_argument("--description", "-d", help="Character description (if no reference)")
    parser.add_argument("--output", "-o", required=True, help="Output file path")
    parser.add_argument("--aspect", "-a", default="16:9", help="Aspect ratio (default: 16:9)")
    parser.add_argument("--name", "-n", help="Character name (e.g., 'Tsar Alexander II')")
    parser.add_argument("--sepia", "-s", action="store_true", help="Apply sepia/vintage tone")

    args = parser.parse_args()

    if not args.reference and not args.description:
        print("Error: Must provide either --reference or --description")
        sys.exit(1)

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

    # Build prompt
    name_suffix = f" of {args.name}" if args.name else ""
    sepia_note = "\n\nApply SEPIA VINTAGE TONE throughout for historical period authenticity." if args.sepia else ""

    if args.reference:
        print(f"Using reference image: {args.reference}")
        if args.name:
            print(f"Character name: {args.name}")
        if args.sepia:
            print("Style: Sepia/vintage")
        with open(args.reference, "rb") as f:
            image_data = f.read()
        prompt = PROMPT_WITH_REF.format(name_suffix=name_suffix, sepia_note=sepia_note)
        # Detect mime type
        mime_type = "image/jpeg"
        if args.reference.lower().endswith(".png"):
            mime_type = "image/png"
        elif args.reference.lower().endswith(".webp"):
            mime_type = "image/webp"
    else:
        print(f"Using description: {args.description[:100]}...")
        if args.sepia:
            print("Style: Sepia/vintage")
        image_data = None
        mime_type = None
        desc_with_sepia = args.description + (". Apply sepia vintage tone for historical authenticity" if args.sepia else "")
        prompt = PROMPT_DESCRIPTION_TEMPLATE.format(description=desc_with_sepia)

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

    if result_bytes:
        # Determine extension
        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"Character sheet saved to: {output_path}")
        print(f"Size: {len(result_bytes)} bytes")
    else:
        print("Failed to generate character sheet")
        sys.exit(1)


if __name__ == "__main__":
    main()
