#!/usr/bin/env python3
"""
Generate image from text prompt using NanoBanana (Gemini 3 Pro).

Text-to-image generation via LaoZhang.ai API or direct Google Gemini API.
Supports multi-turn chat for iterative refinement.

Usage:
    # Single generation
    python nanobanana_text2img.py "prompt text" output.png [aspect_ratio]
    python nanobanana_text2img.py --file prompt.txt output.png [aspect_ratio]

    # Multi-turn chat (refinement)
    python nanobanana_text2img.py --chat history.json "initial prompt" output.png
    python nanobanana_text2img.py --chat history.json --refine "fix frame 7" output2.png

Environment variables:
    NANOBANANA_PROVIDER: "laozhang" (default), "google", or "openrouter"
    LAOZHANG_KEY: API key for LaoZhang.ai
    GOOGLE_API_KEY: API key for direct Google Gemini API
    OPENROUTER_API_KEY: API key for OpenRouter (supports multi-turn natively)

Aspect ratios: 9:16, 16:9, 1:1, 4:3, 3:4
"""

import argparse
import base64
import io
import json
import os
import sys
import time
import requests
from pathlib import Path

try:
    from laozhang_proxy import get_proxied_laozhang_session
except ImportError:
    # Allow running without laozhang_proxy if using Google provider
    def get_proxied_laozhang_session():
        raise ImportError("laozhang_proxy not available, use NANOBANANA_PROVIDER=google")

# API endpoints
LAOZHANG_BASE_URL = "https://api.laozhang.ai/v1beta"
GOOGLE_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"

# Max size per reference image in bytes (before base64). 4MB raw ≈ 5.3MB base64.
MAX_REF_IMAGE_BYTES = 4 * 1024 * 1024


def _get_provider_config() -> tuple:
    """
    Determine which provider to use based on environment.

    Returns:
        Tuple of (provider_name, api_key, base_url, session)
    """
    provider = os.environ.get("MICRODRAMA_NANOBANANA_PROVIDER") or os.environ.get("NANOBANANA_PROVIDER", "laozhang")
    provider = provider.lower()

    if provider == "google":
        api_key = os.environ.get("GOOGLE_API_KEY")
        if not api_key:
            raise ValueError("GOOGLE_API_KEY environment variable not set")
        # Direct Google API - no proxy needed
        session = requests.Session()
        return ("google", api_key, GOOGLE_BASE_URL, session)
    elif provider == "openrouter":
        api_key = os.environ.get("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY environment variable not set")
        session = requests.Session()
        # Configure proxy for OpenRouter
        proxy_url = os.environ.get("OPENROUTER_PROXY", "socks5h://PROXY_HOST:1080")
        if proxy_url and proxy_url.lower() != "none":
            session.proxies = {"http": proxy_url, "https": proxy_url}
        return ("openrouter", api_key, OPENROUTER_BASE_URL, session)
    else:
        # Default: LaoZhang proxy
        api_key = os.environ.get("LAOZHANG_KEY") or os.environ.get("LAOZHANG_API_KEY")
        if not api_key:
            raise ValueError("LAOZHANG_KEY environment variable not set")
        session = get_proxied_laozhang_session()
        return ("laozhang", api_key, LAOZHANG_BASE_URL, session)


def _downsize_image(image_path: str, max_bytes: int = MAX_REF_IMAGE_BYTES) -> bytes:
    """Load image, downsize to fit within max_bytes. Returns JPEG bytes."""
    with open(image_path, 'rb') as f:
        raw = f.read()

    if len(raw) <= max_bytes:
        return raw

    from PIL import Image
    img = Image.open(io.BytesIO(raw))
    if img.mode == 'RGBA':
        img = img.convert('RGB')

    w, h = img.size
    # Scale down large images to ~2MP
    if w * h > 4_000_000:
        scale = (2_000_000 / (w * h)) ** 0.5
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)

    for quality in [85, 70, 55, 40]:
        buf = io.BytesIO()
        img.save(buf, format='JPEG', quality=quality)
        if buf.tell() <= max_bytes:
            print(f"  Downsized {os.path.basename(image_path)}: {len(raw)//1024}KB -> {buf.tell()//1024}KB (q={quality})")
            return buf.getvalue()

    # Last resort
    w, h = img.size
    img = img.resize((w // 2, h // 2), Image.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format='JPEG', quality=50)
    print(f"  Downsized {os.path.basename(image_path)}: {len(raw)//1024}KB -> {buf.tell()//1024}KB (aggressive)")
    return buf.getvalue()


def _load_chat_history(history_path: str) -> list:
    """Load conversation history from JSON file."""
    path = Path(history_path)
    if path.exists():
        with open(path) as f:
            return json.load(f)
    return []


def _save_chat_history(history_path: str, contents: list):
    """Save conversation history to JSON file."""
    path = Path(history_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w') as f:
        json.dump(contents, f, indent=2)
    print(f"  Saved chat history: {path} ({len(contents)} turns)")


def _image_to_base64(image_path: str) -> str:
    """Read image file and return base64 encoded string."""
    with open(image_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8')


def _google_to_or_messages(contents: list) -> list:
    """Convert Google API format contents to OpenRouter/OpenAI format messages."""
    messages = []
    for turn in contents:
        role = "assistant" if turn["role"] == "model" else turn["role"]
        parts = turn.get("parts", [])

        has_images = any("inline_data" in p for p in parts)

        if has_images:
            content_parts = []
            for p in parts:
                if "text" in p:
                    content_parts.append({"type": "text", "text": p["text"]})
                if "inline_data" in p:
                    mime = p["inline_data"].get("mime_type", "image/png")
                    data = p["inline_data"]["data"]
                    content_parts.append({
                        "type": "image_url",
                        "image_url": {"url": f"data:{mime};base64,{data}"}
                    })
            messages.append({"role": role, "content": content_parts})
        else:
            texts = [p["text"] for p in parts if "text" in p]
            messages.append({"role": role, "content": " ".join(texts) if texts else ""})

    return messages


def _or_response_to_google_parts(message: dict) -> list:
    """Convert OpenRouter response message to Google format model parts."""
    parts = []

    content = message.get("content", "")
    if isinstance(content, str) and content:
        parts.append({"text": content})
    elif isinstance(content, list):
        for item in content:
            if item.get("type") == "text":
                parts.append({"text": item["text"]})

    for img in message.get("images", []):
        url = img.get("image_url", {}).get("url", "")
        if url.startswith("data:"):
            header, b64data = url.split(",", 1)
            mime_type = header.split(":")[1].split(";")[0] if ":" in header else "image/png"
            parts.append({
                "inline_data": {
                    "mime_type": mime_type,
                    "data": b64data
                }
            })

    return parts


def generate_image_chat(
    prompt: str,
    output_path: str,
    history_path: str,
    aspect_ratio: str = "9:16",
    image_size: str = "4K",
    reference_image_paths: list = None,
    is_refinement: bool = False
) -> str:
    """
    Generate image using multi-turn chat for iterative refinement.

    On first call (is_refinement=False):
    - Creates new conversation with prompt + reference images
    - Saves model response to history

    On refinement (is_refinement=True):
    - Loads existing history
    - Appends refinement prompt
    - Sends full conversation to API
    - Updates history with new response

    Args:
        prompt: Text prompt (initial or refinement)
        output_path: Where to save the output image
        history_path: Path to conversation history JSON file
        aspect_ratio: Output aspect ratio
        image_size: Output size (1K, 2K, 4K)
        reference_image_paths: Optional reference images (can be used on first call or refinement)
        is_refinement: True if this is a follow-up refinement

    Returns:
        Path to saved image
    """
    provider, api_key, base_url, session = _get_provider_config()

    print(f"Generating image via {provider} (chat mode)...")
    print(f"  History: {history_path}")
    print(f"  Refinement: {is_refinement}")
    print(f"  Aspect ratio: {aspect_ratio}")
    print(f"  Size: {image_size}")

    # Load or initialize conversation history
    contents = _load_chat_history(history_path) if is_refinement else []

    if is_refinement and not contents:
        raise ValueError(f"No chat history found at {history_path} for refinement")

    # Build user message parts
    user_parts = [{"text": prompt}]

    # Add reference images (can be used in any turn)
    if reference_image_paths:
        for idx, ref_path in enumerate(reference_image_paths, 1):
            with open(ref_path, 'rb') as f:
                raw_bytes = f.read()
            print(f"  Reference {idx}: {ref_path} ({len(raw_bytes) // 1024}KB)")
            user_parts.append({
                "inline_data": {
                    "mime_type": "image/jpeg",
                    "data": base64.b64encode(raw_bytes).decode('utf-8')
                }
            })

    # Add user turn to contents
    contents.append({
        "role": "user",
        "parts": user_parts
    })

    print(f"  Prompt: {prompt[:100]}{'...' if len(prompt) > 100 else ''}")
    print(f"  Total turns: {len(contents)}")

    # Calculate payload size
    payload_size = sum(
        len(p.get("inline_data", {}).get("data", ""))
        for c in contents for p in c.get("parts", [])
    ) * 3 // 4  # base64 to bytes approximation
    print(f"  Estimated payload: {payload_size // 1024 // 1024}MB")

    # Make API request with retry logic for transient errors
    max_retries = 5
    last_error = None
    response = None

    if provider == "openrouter":
        # Convert Google format contents to OpenRouter/OpenAI format
        or_messages = _google_to_or_messages(contents)
        request_json = {
            "model": "google/gemini-3-pro-image-preview",
            "messages": or_messages,
            "modalities": ["image", "text"],
            "image_config": {
                "aspect_ratio": aspect_ratio,
                "image_size": image_size
            }
        }
        request_headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        request_url = f"{base_url}/chat/completions"
    else:
        # Google/LaoZhang native format
        request_json = {
            "contents": contents,
            "generationConfig": {
                "responseModalities": ["TEXT", "IMAGE"],
                "imageConfig": {
                    "aspectRatio": aspect_ratio,
                    "imageSize": image_size
                }
            }
        }
        request_headers = {
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        }
        request_url = f"{base_url}/models/gemini-3-pro-image-preview:generateContent"

    for attempt in range(max_retries):
        try:
            response = session.post(
                request_url,
                headers=request_headers,
                json=request_json,
                timeout=300
            )
            if response.status_code in (429, 502, 503, 504):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    wait = min(30, 5 * (2 ** attempt))  # 5, 10, 20, 30, 30
                    print(f"  Retry {attempt + 1}/{max_retries} after {last_error} (wait {wait}s)")
                    time.sleep(wait)
                    continue
            break
        except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
            last_error = str(e)
            if attempt < max_retries - 1:
                wait = min(30, 5 * (2 ** attempt))
                print(f"  Retry {attempt + 1}/{max_retries} after error: {last_error} (wait {wait}s)")
                time.sleep(wait)
            else:
                raise RuntimeError(f"API request failed after {max_retries} retries: {last_error}")

    # LaoZhang proxy fallback: if multi-turn fails (400 or HTML response),
    # retry with model turns stripped to text-only, moving last model image
    # to next user turn. Not needed for OpenRouter (handles multi-turn natively).
    is_html = response is not None and response.content.startswith(b'<!')
    if provider != "openrouter" and is_refinement and response is not None and (response.status_code == 400 or is_html):
        print(f"  Multi-turn failed ({response.status_code}), retrying with proxy workaround...")
        import copy
        fallback_contents = copy.deepcopy(_load_chat_history(history_path))
        last_model_img = None
        for turn in fallback_contents:
            if turn.get("role") == "model":
                new_parts = []
                for part in turn.get("parts", []):
                    if "inline_data" in part:
                        last_model_img = part["inline_data"]
                    elif "text" in part:
                        new_parts.append({"text": part["text"]})
                if not new_parts:
                    new_parts = [{"text": "Here is the generated image."}]
                turn["parts"] = new_parts
        fb_user_parts = [{"text": prompt}]
        if last_model_img:
            fb_user_parts.append({"inline_data": last_model_img})
            print(f"  Re-attached model image as user ref ({len(last_model_img.get('data',''))//1024}KB)")
        if reference_image_paths:
            for ref_path in reference_image_paths:
                with open(ref_path, 'rb') as f:
                    fb_user_parts.append({
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": base64.b64encode(f.read()).decode('utf-8')
                        }
                    })
        fallback_contents.append({"role": "user", "parts": fb_user_parts})
        request_json["contents"] = fallback_contents
        contents = fallback_contents
        response = session.post(request_url, headers=request_headers, json=request_json, timeout=300)
        print(f"  Fallback status: {response.status_code}")

    if response is None or response.status_code != 200:
        status = response.status_code if response else "no response"
        text = response.text[:500] if response else str(last_error)
        raise RuntimeError(f"API error: {status} - {text}")

    data = response.json()

    # Parse response based on provider format
    if provider == "openrouter":
        # OpenRouter/OpenAI response format
        choices = data.get("choices", [])
        if not choices:
            error_log_path = output_path.rsplit('.', 1)[0] + '_error_response.json'
            with open(error_log_path, 'w') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            print(f"  ERROR: No choices in response. Saved to: {error_log_path}")
            if "error" in data:
                print(f"  API error: {data['error']}")
            raise RuntimeError(f"No choices in API response. See {error_log_path}")

        message = choices[0].get("message", {})
        finish_reason = choices[0].get("finish_reason", "unknown")

        # Log usage
        usage = data.get("usage", {})
        if usage:
            cost = usage.get('completion_tokens', 0) * 120 / 1_000_000
            print(f"  Tokens: {usage.get('prompt_tokens', '?')} in / {usage.get('completion_tokens', '?')} out (est ${cost:.4f})")

        # Log text content
        msg_content = message.get("content", "")
        if isinstance(msg_content, str) and msg_content:
            print(f"  Model text: {msg_content[:300]}{'...' if len(msg_content) > 300 else ''}")

        # Convert OpenRouter response to Google format parts for storage
        model_parts = _or_response_to_google_parts(message)

        # Save image from model_parts
        saved_image = False
        for part in model_parts:
            if "inline_data" in part and not saved_image:
                image_data = base64.b64decode(part["inline_data"]["data"])
                mime_type = part["inline_data"].get("mime_type", "image/png")
                ext = ".png"
                if "jpeg" in mime_type or "jpg" in mime_type:
                    ext = ".jpg"
                elif "webp" in mime_type:
                    ext = ".webp"
                if not output_path.lower().endswith(ext):
                    output_path = output_path.rsplit('.', 1)[0] + ext
                with open(output_path, 'wb') as f:
                    f.write(image_data)
                print(f"  Saved: {output_path} ({len(image_data) // 1024}KB)")
                saved_image = True

        print(f"  Response: {len(model_parts)} parts, finish_reason={finish_reason}")

    else:
        # Google/LaoZhang response format
        candidates = data.get("candidates", [])
        if not candidates:
            error_log_path = output_path.rsplit('.', 1)[0] + '_error_response.json'
            with open(error_log_path, 'w') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            print(f"  ERROR: No candidates in response. Saved to: {error_log_path}")
            if "promptFeedback" in data:
                print(f"  Prompt feedback: {data['promptFeedback']}")
            if "error" in data:
                print(f"  API error: {data['error']}")
            raise RuntimeError(f"No candidates in API response. See {error_log_path}")

        candidate = candidates[0]
        finish_reason = candidate.get("finishReason", "unknown")
        content = candidate.get("content") or {}
        response_parts = content.get("parts") or []

        if not response_parts:
            error_log_path = output_path.rsplit('.', 1)[0] + '_error_response.json'
            debug_data = {k: v for k, v in data.items()}
            with open(error_log_path, 'w') as f:
                json.dump(debug_data, f, indent=2, ensure_ascii=False, default=str)
            print(f"  ERROR: No parts in response. finishReason={finish_reason}")
            print(f"  Saved response to: {error_log_path}")
            if "safetyRatings" in candidate:
                print(f"  Safety ratings: {candidate['safetyRatings']}")
            raise RuntimeError(f"No parts in API response (finishReason={finish_reason}). See {error_log_path}")

        print(f"  Response: {len(response_parts)} parts, finishReason={finish_reason}")

        # Log text response
        for part in response_parts:
            if "text" in part:
                text = part["text"]
                print(f"  Model text: {text[:300]}{'...' if len(text) > 300 else ''}")

        # Find and save image, preserving ALL model response parts (including thoughtSignature)
        saved_image = False
        model_parts = []

        for part in response_parts:
            model_part = {}

            if "inlineData" in part:
                image_data = base64.b64decode(part["inlineData"]["data"])
                mime_type = part["inlineData"].get("mimeType", "image/png")

                if not saved_image:
                    ext = ".png"
                    if "jpeg" in mime_type or "jpg" in mime_type:
                        ext = ".jpg"
                    elif "webp" in mime_type:
                        ext = ".webp"
                    if not output_path.lower().endswith(ext):
                        output_path = output_path.rsplit('.', 1)[0] + ext
                    with open(output_path, 'wb') as f:
                        f.write(image_data)
                    print(f"  Saved: {output_path} ({len(image_data) // 1024}KB)")
                    saved_image = True

                model_part["inline_data"] = {
                    "mime_type": mime_type,
                    "data": part["inlineData"]["data"]
                }

            if "text" in part:
                model_part["text"] = part["text"]

            if "thoughtSignature" in part:
                model_part["thought_signature"] = part["thoughtSignature"]
                print(f"  Preserved thought_signature ({len(part['thoughtSignature'])} chars)")

            if "thought" in part:
                model_part["thought"] = part["thought"]

            if model_part:
                model_parts.append(model_part)

    if not saved_image:
        raise RuntimeError("No image found in API response")

    # Add model response to history (with all parts including signatures)
    contents.append({
        "role": "model",
        "parts": model_parts
    })

    # Save updated history
    _save_chat_history(history_path, contents)

    return output_path


def _save_request_debug(output_path, request_json, provider,
                        primary_ref_path=None, additional_ref_paths=None):
    """Save a debug copy of the NanoBanana request with image data replaced by metadata."""
    try:
        debug_path = Path(output_path).with_suffix('.request.json')
        debug = {"provider": provider, "output_path": output_path}

        if provider == "openrouter":
            debug["model"] = request_json.get("model")
            debug["image_config"] = request_json.get("image_config")
            # Extract text parts, replace image data with metadata
            messages = request_json.get("messages", [])
            debug_messages = []
            for msg in messages:
                content = msg.get("content", "")
                if isinstance(content, list):
                    debug_parts = []
                    img_idx = 0
                    for part in content:
                        if part.get("type") == "text":
                            debug_parts.append(part)
                        elif part.get("type") == "image_url":
                            url = part["image_url"]["url"]
                            size_kb = len(url.split(",", 1)[-1]) * 3 // 4 // 1024 if "," in url else 0
                            ref_path = primary_ref_path if img_idx == 0 else (
                                additional_ref_paths[img_idx - 1] if additional_ref_paths and img_idx - 1 < len(additional_ref_paths) else None
                            )
                            debug_parts.append({
                                "type": "image_url",
                                "image_index": img_idx + 1,
                                "source_path": ref_path,
                                "approx_size_kb": size_kb
                            })
                            img_idx += 1
                    debug_messages.append({"role": msg["role"], "content": debug_parts})
                else:
                    debug_messages.append(msg)
            debug["messages"] = debug_messages
        else:
            debug["generation_config"] = request_json.get("generationConfig")
            # Extract parts, replace inline_data with metadata
            contents = request_json.get("contents", [{}])
            debug_parts = []
            img_idx = 0
            for part in contents[0].get("parts", []):
                if "text" in part:
                    debug_parts.append(part)
                elif "inline_data" in part:
                    data_len = len(part["inline_data"]["data"])
                    size_kb = data_len * 3 // 4 // 1024
                    ref_path = primary_ref_path if img_idx == 0 else (
                        additional_ref_paths[img_idx - 1] if additional_ref_paths and img_idx - 1 < len(additional_ref_paths) else None
                    )
                    debug_parts.append({
                        "inline_data_placeholder": True,
                        "image_index": img_idx + 1,
                        "mime_type": part["inline_data"]["mime_type"],
                        "source_path": ref_path,
                        "approx_size_kb": size_kb
                    })
                    img_idx += 1
            debug["contents"] = [{"parts": debug_parts}]

        with open(debug_path, 'w') as f:
            json.dump(debug, f, indent=2, ensure_ascii=False)
        print(f"  Debug JSON: {debug_path}")
    except Exception as e:
        print(f"  Warning: failed to save debug JSON: {e}")


def generate_image(
    prompt: str,
    output_path: str,
    aspect_ratio: str = "9:16",
    image_size: str = "4K",
    reference_image_path: str = None,
    reference_strength: float = 0.5,
    additional_reference_paths: list = None
) -> str:
    """
    Generate image from text prompt using NanoBanana.

    Args:
        prompt: Text description of the image to generate
        output_path: Where to save the output image
        aspect_ratio: Output aspect ratio (9:16, 16:9, 1:1, 4:3, 3:4)
        image_size: Output size (1K, 2K, 4K)
        reference_image_path: Optional path to primary reference image for image-to-image generation
        reference_strength: Strength of reference image influence (0.0-1.0), default 0.5 (DEPRECATED - not used by Gemini)
        additional_reference_paths: Optional list of additional reference image paths for context

    Returns:
        Path to saved image
    """
    provider, api_key, base_url, session = _get_provider_config()

    print(f"Generating image via {provider}...")
    print(f"  Aspect ratio: {aspect_ratio}")
    print(f"  Size: {image_size}")
    if reference_image_path:
        print(f"  Primary reference: {reference_image_path}")
    if additional_reference_paths:
        print(f"  Additional references: {len(additional_reference_paths)} images")
        for idx, path in enumerate(additional_reference_paths, 1):
            print(f"    {idx}. {path}")
    print(f"  Prompt: {prompt[:100]}{'...' if len(prompt) > 100 else ''}")

    # Build parts array (text first, then images — matches official LaoZhang format)
    parts = [{"text": prompt}]

    # Add primary reference image if provided
    if reference_image_path:
        with open(reference_image_path, 'rb') as f:
            raw_bytes = f.read()
        print(f"  Primary ref size: {len(raw_bytes) // 1024}KB")
        ref_image_data = base64.b64encode(raw_bytes).decode('utf-8')
        # Use snake_case field names (inline_data, mime_type) to match official API format
        parts.append({
            "inline_data": {
                "mime_type": "image/jpeg",
                "data": ref_image_data
            }
        })

    # Add additional reference images at full resolution
    if additional_reference_paths:
        for idx, ref_path in enumerate(additional_reference_paths, 1):
            with open(ref_path, 'rb') as f:
                raw_bytes = f.read()
            print(f"  Additional ref {idx} size: {len(raw_bytes) // 1024}KB")
            ref_image_data = base64.b64encode(raw_bytes).decode('utf-8')
            parts.append({
                "inline_data": {
                    "mime_type": "image/jpeg",
                    "data": ref_image_data
                }
            })

    if provider == "openrouter":
        # Build OpenRouter/OpenAI format content parts
        or_content = [{"type": "text", "text": prompt}]
        for p in parts[1:]:  # skip text part, add images
            if "inline_data" in p:
                mime = p["inline_data"]["mime_type"]
                data_b64 = p["inline_data"]["data"]
                or_content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:{mime};base64,{data_b64}"}
                })

        request_json = {
            "model": "google/gemini-3-pro-image-preview",
            "messages": [{"role": "user", "content": or_content if len(or_content) > 1 else prompt}],
            "modalities": ["image", "text"],
            "image_config": {
                "aspect_ratio": aspect_ratio,
                "image_size": image_size
            }
        }
        request_headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        request_url = f"{base_url}/chat/completions"
    else:
        request_json = {
            "contents": [{"parts": parts}],
            "generationConfig": {
                "responseModalities": ["TEXT", "IMAGE"],
                "imageConfig": {
                    "aspectRatio": aspect_ratio,
                    "imageSize": image_size
                }
            }
        }
        request_headers = {
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        }
        request_url = f"{base_url}/models/gemini-3-pro-image-preview:generateContent"

    # Save debug JSON (image data replaced with metadata for readability)
    _save_request_debug(output_path, request_json, provider,
                        reference_image_path, additional_reference_paths)

    max_retries = 5
    last_error = None
    response = None
    for attempt in range(max_retries):
        try:
            response = session.post(
                request_url,
                headers=request_headers,
                json=request_json,
                timeout=180
            )
            if response.status_code in (429, 502, 503, 504):
                last_error = f"HTTP {response.status_code}"
                if attempt < max_retries - 1:
                    wait = min(30, 5 * (2 ** attempt))
                    print(f"  Retry {attempt + 1}/{max_retries} after {last_error} (wait {wait}s)")
                    time.sleep(wait)
                    continue
            break
        except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
            last_error = str(e)
            if attempt < max_retries - 1:
                wait = min(30, 5 * (2 ** attempt))
                print(f"  Retry {attempt + 1}/{max_retries} after error: {last_error} (wait {wait}s)")
                time.sleep(wait)
            else:
                raise RuntimeError(f"API request failed after {max_retries} retries: {last_error}")

    if response is None or response.status_code != 200:
        status = response.status_code if response else "no response"
        text = response.text[:500] if response else str(last_error)
        raise RuntimeError(f"API error: {status} - {text}")

    data = response.json()

    if provider == "openrouter":
        # OpenRouter response format
        choices = data.get("choices", [])
        if not choices:
            raise RuntimeError(f"No choices in API response: {str(data)[:300]}")

        message = choices[0].get("message", {})
        usage = data.get("usage", {})
        if usage:
            cost = usage.get('completion_tokens', 0) * 120 / 1_000_000
            print(f"  Tokens: {usage.get('prompt_tokens', '?')} in / {usage.get('completion_tokens', '?')} out (est ${cost:.4f})")

        msg_content = message.get("content", "")
        if isinstance(msg_content, str) and msg_content:
            print(f"  Model text: {msg_content[:300]}{'...' if len(msg_content) > 300 else ''}")

        for img in message.get("images", []):
            url = img.get("image_url", {}).get("url", "")
            if url.startswith("data:"):
                header, b64data = url.split(",", 1)
                mime_type = header.split(":")[1].split(";")[0] if ":" in header else "image/png"
                image_data = base64.b64decode(b64data)
                ext = ".png"
                if "jpeg" in mime_type or "jpg" in mime_type:
                    ext = ".jpg"
                elif "webp" in mime_type:
                    ext = ".webp"
                if not output_path.lower().endswith(ext):
                    output_path = output_path.rsplit('.', 1)[0] + ext
                with open(output_path, 'wb') as f:
                    f.write(image_data)
                print(f"  Saved: {output_path} ({len(image_data) // 1024}KB)")
                return output_path

        raise RuntimeError("No image found in API response")
    else:
        # Google/LaoZhang response format
        resp_parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])

        for part in resp_parts:
            if "text" in part:
                text = part["text"]
                print(f"  Model text: {text[:300]}{'...' if len(text) > 300 else ''}")

        for part in resp_parts:
            if "inlineData" in part:
                image_data = base64.b64decode(part["inlineData"]["data"])
                mime_type = part["inlineData"].get("mimeType", "image/png")
                ext = ".png"
                if "jpeg" in mime_type or "jpg" in mime_type:
                    ext = ".jpg"
                elif "webp" in mime_type:
                    ext = ".webp"
                if not output_path.lower().endswith(ext):
                    output_path = output_path.rsplit('.', 1)[0] + ext
                with open(output_path, 'wb') as f:
                    f.write(image_data)
                print(f"  Saved: {output_path} ({len(image_data) // 1024}KB)")
                return output_path

        raise RuntimeError("No image found in API response")


def main():
    parser = argparse.ArgumentParser(
        description="Generate image from text using NanoBanana (Gemini 3 Pro)"
    )
    parser.add_argument(
        "prompt",
        nargs="?",
        help="Text prompt or --file flag (or 'test' for API test)"
    )
    parser.add_argument(
        "output",
        nargs="?",
        help="Output image path"
    )
    parser.add_argument(
        "--file", "-f",
        action="store_true",
        help="Treat prompt argument as path to file containing prompt"
    )
    parser.add_argument(
        "--aspect-ratio", "-a",
        default="9:16",
        choices=["9:16", "16:9", "1:1", "4:3", "3:4"],
        help="Output aspect ratio (default: 9:16)"
    )
    parser.add_argument(
        "--size", "-s",
        default="4K",
        choices=["1K", "2K", "4K"],
        help="Output size (default: 4K)"
    )
    parser.add_argument(
        "--provider", "-p",
        choices=["laozhang", "google", "openrouter"],
        help="API provider (default: from NANOBANANA_PROVIDER env var or 'laozhang')"
    )
    parser.add_argument(
        "--test",
        action="store_true",
        help="Run API connectivity test"
    )
    parser.add_argument(
        "--chat",
        metavar="HISTORY_FILE",
        help="Enable chat mode with conversation history file (JSON)"
    )
    parser.add_argument(
        "--refine",
        action="store_true",
        help="This is a refinement (requires --chat with existing history)"
    )
    parser.add_argument(
        "--ref",
        action="append",
        metavar="IMAGE",
        help="Reference image path (can be specified multiple times)"
    )

    args = parser.parse_args()

    # Override provider if specified
    if args.provider:
        os.environ["NANOBANANA_PROVIDER"] = args.provider

    # Test mode
    if args.test or args.prompt == "test":
        print("Testing NanoBanana API...")
        try:
            provider, api_key, base_url, session = _get_provider_config()
            print(f"  Provider: {provider}")
            print(f"  Base URL: {base_url}")

            response = session.post(
                f"{base_url}/models/gemini-3-pro-image-preview:generateContent",
                headers={
                    "x-goog-api-key": api_key,
                    "Content-Type": "application/json"
                },
                json={
                    "contents": [{"parts": [{"text": "A simple red circle on white background"}]}],
                    "generationConfig": {
                        "responseModalities": ["TEXT", "IMAGE"],
                        "imageConfig": {"aspectRatio": "1:1", "imageSize": "1K"}
                    }
                },
                timeout=60
            )

            print(f"  Status: {response.status_code}")
            if response.status_code == 200:
                data = response.json()
                parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
                for part in parts:
                    if "inlineData" in part:
                        size = len(part["inlineData"]["data"]) * 3 // 4
                        mime = part["inlineData"].get("mimeType", "unknown")
                        print(f"  Success! Image: {size // 1024}KB, {mime}")
                        break
                else:
                    print(f"  Response (no image): {str(data)[:300]}")
            else:
                print(f"  Error: {response.text[:500]}")
        except Exception as e:
            print(f"  Error: {e}")
        return

    # Normal generation mode
    if not args.prompt or not args.output:
        parser.error("prompt and output are required for generation (or use --test)")

    # Get prompt text
    if args.file:
        with open(args.prompt) as f:
            prompt = f.read().strip()
    else:
        prompt = args.prompt

    # Chat mode (multi-turn refinement)
    if args.chat:
        if args.refine and not Path(args.chat).exists():
            parser.error(f"--refine requires existing history file: {args.chat}")

        generate_image_chat(
            prompt=prompt,
            output_path=args.output,
            history_path=args.chat,
            aspect_ratio=args.aspect_ratio,
            image_size=args.size,
            reference_image_paths=args.ref,
            is_refinement=args.refine
        )
    else:
        # Single-shot generation (original behavior)
        generate_image(
            prompt=prompt,
            output_path=args.output,
            aspect_ratio=args.aspect_ratio,
            image_size=args.size,
            reference_image_path=args.ref[0] if args.ref else None,
            additional_reference_paths=args.ref[1:] if args.ref and len(args.ref) > 1 else None
        )


if __name__ == "__main__":
    main()
