#!/usr/bin/env python3
"""
三语绘本图像生成脚本 - Nano Banana 2 版本 (v2)
尝试多种 API 格式来调用 Google Gemini 3.1 Flash Image Preview
"""

import os
import requests
import json
import base64
import time
from pathlib import Path
from datetime import datetime

# API 配置
OPENROUTER_API_KEY = os.environ.get('OPENROUTER_API_KEY')
OPENROUTER_BASE = "https://openrouter.ai/api/v1"

# 模型：Nano Banana 2
MODEL = "google/gemini-3.1-flash-image-preview"

# 输出目录
OUTPUT_DIR = Path("/root/.openclaw/workspace/trilingual-picturebook/output")
OUTPUT_DIR.mkdir(exist_ok=True)

# 测试提示词
TEST_PROMPT = "children's book illustration, watercolor style, 2 year old Chinese boy, short black hair, round face, big bright eyes, wearing light blue t-shirt, cute smile, looking at viewer, soft pastel colors, warm lighting, simple white background, for toddlers picture book, square composition, no text, no watermark, high quality, detailed face"


def try_method_1():
    """
    方法 1: 使用 chat/completions + response_format
    """
    print("\n📌 方法 1: chat/completions + response_format")
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator"
    }
    
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "user", "content": f"Generate an image: {TEST_PROMPT}"}
        ],
        "response_format": {"type": "image"}
    }
    
    try:
        response = requests.post(f"{OPENROUTER_BASE}/chat/completions", headers=headers, json=payload, timeout=60)
        print(f"   状态码：{response.status_code}")
        save_response(response, "method1")
        return response
    except Exception as e:
        print(f"   错误：{e}")
        return None


def try_method_2():
    """
    方法 2: 使用 images/generations endpoint
    """
    print("\n📌 方法 2: images/generations endpoint")
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator"
    }
    
    payload = {
        "model": MODEL,
        "prompt": TEST_PROMPT,
        "size": "1024x1024",
        "n": 1
    }
    
    try:
        response = requests.post(f"{OPENROUTER_BASE}/images/generations", headers=headers, json=payload, timeout=60)
        print(f"   状态码：{response.status_code}")
        save_response(response, "method2")
        return response
    except Exception as e:
        print(f"   错误：{e}")
        return None


def try_method_3():
    """
    方法 3: 使用 chat/completions + 多模态响应请求
    """
    print("\n📌 方法 3: chat/completions + 多模态内容请求")
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator"
    }
    
    payload = {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Generate an image based on this description:"},
                    {"type": "text", "text": TEST_PROMPT}
                ]
            }
        ],
        "max_tokens": 4096
    }
    
    try:
        response = requests.post(f"{OPENROUTER_BASE}/chat/completions", headers=headers, json=payload, timeout=60)
        print(f"   状态码：{response.status_code}")
        save_response(response, "method3")
        return response
    except Exception as e:
        print(f"   错误：{e}")
        return None


def try_method_4():
    """
    方法 4: 直接使用 Google AI Studio 格式 (通过 OpenRouter 转发)
    """
    print("\n📌 方法 4: Google AI Studio 格式")
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator"
    }
    
    # Gemini 原生格式
    payload = {
        "model": MODEL,
        "contents": [{
            "parts": [{"text": f"Generate image: {TEST_PROMPT}"}]
        }],
        "generationConfig": {
            "responseModalities": ["IMAGE"]
        }
    }
    
    try:
        response = requests.post(f"{OPENROUTER_BASE}/chat/completions", headers=headers, json=payload, timeout=60)
        print(f"   状态码：{response.status_code}")
        save_response(response, "method4")
        return response
    except Exception as e:
        print(f"   错误：{e}")
        return None


def save_response(response, method_name):
    """保存响应到文件"""
    try:
        output_path = OUTPUT_DIR / f"api-test-{method_name}.json"
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(response.json() if response.content else {"empty": True}, f, indent=2, ensure_ascii=False)
        print(f"   响应已保存：{output_path}")
    except:
        pass


def main():
    print("=" * 60)
    print("🔬 Nano Banana 2 API 测试")
    print("=" * 60)
    
    if not OPENROUTER_API_KEY:
        print("❌ 错误：未设置 OPENROUTER_API_KEY")
        return
    
    print(f"\n🍌 模型：{MODEL}")
    print(f"📁 输出：{OUTPUT_DIR}")
    print("\n尝试多种 API 格式...")
    
    methods = [
        ("方法 1", try_method_1),
        ("方法 2", try_method_2),
        ("方法 3", try_method_3),
        ("方法 4", try_method_4),
    ]
    
    for name, method in methods:
        print(f"\n{'='*60}")
        response = method()
        
        if response and response.status_code == 200:
            try:
                result = response.json()
                # 检查是否有图片数据
                has_image = False
                
                # 检查各种可能的图片字段
                for key in ["images", "image", "image_url", "data"]:
                    if key in result:
                        has_image = True
                        break
                
                if "choices" in result:
                    for choice in result["choices"]:
                        msg = choice.get("message", {})
                        content = msg.get("content", "")
                        if content and (isinstance(content, str) and ("data:image" in content or "http" in content)):
                            has_image = True
                        if isinstance(content, list):
                            for item in content:
                                if isinstance(item, dict) and ("image_url" in item or "image" in item):
                                    has_image = True
                
                if has_image:
                    print(f"   ✅ 找到图片数据!")
                else:
                    print(f"   ⚠️  响应成功但未找到图片数据")
            except:
                pass
    
    print("\n" + "=" * 60)
    print("测试完成！请检查 output/ 目录中的响应文件")
    print("=" * 60)


if __name__ == "__main__":
    main()