#!/usr/bin/env python3
"""
三语绘本图像生成脚本 - OpenRouter 版本
使用 OpenRouter API 调用支持图像生成的模型
"""

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

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

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

# 角色提示词
CHARACTER_PROMPTS = {
    "character-boy-front": "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",
    "character-boy-side": "children's book illustration, watercolor style, 2 year old Chinese boy, short black hair, side profile view, wearing light blue t-shirt, cute expression, soft pastel colors, warm lighting, simple white background, for toddlers picture book, square composition, no text, no watermark, high quality",
    "character-girl-front": "children's book illustration, watercolor style, 2 year old Chinese girl, black hair with two small pigtails, bangs, round face, big bright eyes, wearing pink dress, 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",
    "character-girl-side": "children's book illustration, watercolor style, 2 year old Chinese girl, black hair with two small pigtails, side profile, wearing pink dress, cute expression, soft pastel colors, warm lighting, simple white background, for toddlers picture book, square composition, no text, no watermark, high quality",
}

# 《哥哥妹妹早上好》16 页提示词
GOODMORNING_PROMPTS = {
    "goodmorning-page-01": "children's book illustration, watercolor style, bright yellow sun rising outside bedroom window, small blue birds flying in sky, fluffy white clouds, morning golden light streaming through window, seen from inside bedroom, window frame visible, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-02": "children's book illustration, watercolor style, 2 year old Chinese boy with short black hair, sitting up in bed, stretching arms, yawning, wearing blue pajamas, messy morning hair, soft morning sunlight from window, bedroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-03": "children's book illustration, watercolor style, 2 year old Chinese girl with black hair and two small pigtails, sitting up in bed, stretching arms, cute yawn, wearing pink pajamas, sleepy expression, soft morning sunlight from window, bedroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-04": "children's book illustration, watercolor style, 2 year old Chinese boy and girl holding hands, boy in blue pajamas, girl in pink pajamas, walking together down hallway, morning light, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-05": "children's book illustration, watercolor style, 2 year old Chinese boy standing at bathroom sink, brushing teeth, toothbrush in hand, toothpaste foam on mouth, looking in bathroom mirror, blue pajamas, bathroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-06": "children's book illustration, watercolor style, 2 year old Chinese girl with pigtails at bathroom sink, brushing teeth, toothbrush in hand, cute expression, looking at boy next to her, pink pajamas, bathroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-07": "children's book illustration, watercolor style, 2 year old Chinese boy washing face with both hands, water splashing, white towel hanging nearby, blue pajamas, bathroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-08": "children's book illustration, watercolor style, 2 year old Chinese girl with pigtails drying face, using soft white towel, clean fresh face, happy expression, pink pajamas, bathroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-09": "children's book illustration, watercolor style, 2 year old Chinese boy putting on light blue t-shirt, arms in sleeves, bedroom scene, morning light from window, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-10": "children's book illustration, watercolor style, 2 year old Chinese girl pointing at boy, excited expression, clapping hands, boy standing proudly in blue t-shirt, bedroom scene, morning light, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-11": "children's book illustration, watercolor style, 2 year old Chinese boy helping girl button her dress, careful expression, girl waiting patiently, boy in blue t-shirt, girl in pink dress, bedroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-12": "children's book illustration, watercolor style, 2 year old Chinese boy and girl standing in front of mirror, both dressed nicely, looking at themselves, smiling, boy in blue outfit, girl in pink dress, bedroom scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-13": "children's book illustration, watercolor style, 2 year old Chinese boy and girl holding hands, walking down wooden stairs, boy leading, girl following, boy in blue outfit, girl in pink dress, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-14": "children's book illustration, watercolor style, cute golden retriever puppy at bottom of stairs, looking up, wagging tail, happy expression, boy and girl visible on stairs above, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-15": "children's book illustration, watercolor style, 2 year old Chinese boy and girl petting golden puppy, gentle touches, puppy enjoying with eyes closed happily, living room scene, morning light, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
    "goodmorning-page-16": "children's book illustration, watercolor style, 2 year old Chinese boy and girl with golden puppy, sitting together by window, looking outside, bright sunshine streaming in, peaceful scene, soft pastel colors, warm lighting, for toddlers picture book, square composition, no text, no watermark, high quality",
}


def get_available_models():
    """
    获取 OpenRouter 可用模型列表
    """
    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
    
    try:
        response = requests.get("https://openrouter.ai/api/v1/models", headers=headers, timeout=10)
        response.raise_for_status()
        return response.json().get("data", [])
    except Exception as e:
        print(f"❌ 获取模型列表失败：{e}")
        return []


def find_image_models(models):
    """
    查找支持图像生成的模型
    """
    image_models = []
    
    for model in models:
        model_id = model.get("id", "").lower()
        name = model.get("name", "").lower()
        
        # 查找图像生成相关模型
        if any(keyword in model_id or keyword in name for keyword in 
               ["image", "diffusion", "flux", "stable", "playground", "nano", "banana", "dall", "midjourney"]):
            image_models.append(model)
    
    return image_models


def generate_image(prompt: str, filename: str, model: str = "google/imagen-3") -> bool:
    """
    调用 OpenRouter API 生成图片
    """
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator"
    }
    
    # 尝试不同的 API 格式
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"Generate an image: {prompt}"
                    }
                ]
            }
        ],
        "max_tokens": 1
    }
    
    try:
        print(f"📸 正在生成：{filename} (模型：{model})...")
        response = requests.post(API_URL, headers=headers, json=payload, timeout=120)
        
        print(f"   状态码：{response.status_code}")
        
        if response.status_code == 200:
            result = response.json()
            print(f"   响应：{json.dumps(result, ensure_ascii=False)[:500]}...")
            
            # 检查是否有图片数据
            if "choices" in result and len(result["choices"]) > 0:
                content = result["choices"][0].get("message", {}).get("content", "")
                
                # 检查是否包含图片 URL 或 base64
                if "https://" in content:
                    # 提取 URL
                    import re
                    urls = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', content)
                    if urls:
                        image_url = urls[0]
                        print(f"   图片 URL: {image_url}")
                        
                        # 下载图片
                        img_response = requests.get(image_url, timeout=30)
                        if img_response.status_code == 200:
                            output_path = OUTPUT_DIR / f"{filename}.png"
                            with open(output_path, "wb") as f:
                                f.write(img_response.content)
                            print(f"✅ 已保存：{output_path}")
                            return True
                elif "data:image" in content:
                    # base64 格式
                    import re
                    match = re.search(r'data:image/png;base64,([A-Za-z0-9+/=]+)', content)
                    if match:
                        img_data = base64.b64decode(match.group(1))
                        output_path = OUTPUT_DIR / f"{filename}.png"
                        with open(output_path, "wb") as f:
                            f.write(img_data)
                        print(f"✅ 已保存：{output_path}")
                        return True
            
            print(f"⚠️  响应中未找到图片数据")
            return False
        else:
            print(f"❌ API 错误：{response.status_code}")
            print(f"   响应：{response.text[:500]}")
            return False
            
    except requests.exceptions.Timeout:
        print(f"❌ 请求超时")
        return False
    except Exception as e:
        print(f"❌ 生成失败：{e}")
        return False


def main():
    """
    主函数
    """
    print("=" * 60)
    print("📚 三语绘本图像生成器 - OpenRouter 版")
    print("=" * 60)
    
    if not OPENROUTER_API_KEY:
        print("❌ 错误：未设置 OPENROUTER_API_KEY 环境变量")
        return
    
    # 获取可用模型
    print("\n🔍 获取可用模型...")
    models = get_available_models()
    
    if models:
        print(f"✅ 找到 {len(models)} 个模型")
        
        # 查找图像生成模型
        image_models = find_image_models(models)
        if image_models:
            print(f"\n🎨 找到 {len(image_models)} 个图像生成相关模型:")
            for i, model in enumerate(image_models[:10], 1):
                print(f"   {i}. {model.get('id')} - {model.get('name', 'N/A')}")
        else:
            print("\n⚠️  未找到明确的图像生成模型")
            print("   OpenRouter 主要提供文本 LLM 模型")
            print("   建议使用专门的图像生成服务:")
            print("   - Leonardo.ai (推荐)")
            print("   - Stability AI API")
            print("   - DALL-E 3 API")
            print("   - Playground AI")
    else:
        print("❌ 无法获取模型列表")
    
    # 显示提示
    print("\n" + "=" * 60)
    print("⚠️  注意：OpenRouter 主要提供文本 LLM 模型，不是专门的图像生成平台")
    print("   如果需要使用 nano banana 2 生成图片，请确认:")
    print("   1. 该模型确实在 OpenRouter 上可用")
    print("   2. 该模型支持图像生成输出")
    print("   3. API endpoint 和参数格式正确")
    print("\n   建议使用以下替代方案生成绘本图片:")
    print("   - Leonardo.ai (有免费额度，质量优秀)")
    print("   - 使用 prompts/003-optimized-prompts.md 中的提示词")
    print("=" * 60)


if __name__ == "__main__":
    main()
