#!/usr/bin/env python3
"""
三语绘本图像生成脚本 - 使用角色参考图
使用 OpenRouter API 调用 Google Gemini 3.1 Flash Image Preview (Nano Banana 2)
支持加载已有角色图作为参考，确保角色一致性
"""

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')
API_URL = "https://openrouter.ai/api/v1/chat/completions"

# 模型：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)

# 图像分辨率设置 (width x height)
# 可选: "512x512", "768x768", "1024x1024", "1280x1280"
# 低分辨率可能减少限流
IMAGE_SIZE = "512x512"

# 角色参考图路径
CHARACTER_SHEET_DIR = Path("/root/.openclaw/workspace/trilingual-picturebook")
CHARACTER_REFERENCES = {
    "yangyang": CHARACTER_SHEET_DIR / "_charsheet_yangyang.png",
    "zhuazhua": CHARACTER_SHEET_DIR / "_charsheet_zhuazhua.png",
}

# 角色提示词（带参考图）
CHARACTER_PROMPTS = {
    "character-boy-front": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "character-boy-side": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "character-girl-front": {
        "text": "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",
        "characters": ["zhuazhua"]
    },
    "character-girl-side": {
        "text": "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",
        "characters": ["zhuazhua"]
    },
}

# 《哥哥妹妹早上好》16 页提示词（带角色引用）
GOODMORNING_PROMPTS = {
    "goodmorning-page-01": {
        "text": "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",
        "characters": []
    },
    "goodmorning-page-02": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "goodmorning-page-03": {
        "text": "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",
        "characters": ["zhuazhua"]
    },
    "goodmorning-page-04": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-05": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "goodmorning-page-06": {
        "text": "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",
        "characters": ["zhuazhua", "yangyang"]
    },
    "goodmorning-page-07": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "goodmorning-page-08": {
        "text": "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",
        "characters": ["zhuazhua"]
    },
    "goodmorning-page-09": {
        "text": "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",
        "characters": ["yangyang"]
    },
    "goodmorning-page-10": {
        "text": "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",
        "characters": ["zhuazhua", "yangyang"]
    },
    "goodmorning-page-11": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-12": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-13": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-14": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-15": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
    "goodmorning-page-16": {
        "text": "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",
        "characters": ["yangyang", "zhuazhua"]
    },
}


def load_character_image(character_name: str) -> str | None:
    """
    加载角色参考图并转换为 base64 data URL
    """
    char_path = CHARACTER_REFERENCES.get(character_name)
    if not char_path or not char_path.exists():
        print(f"   ⚠️  角色图不存在：{char_path}")
        return None
    
    try:
        with open(char_path, "rb") as f:
            image_data = f.read()
        
        base64_data = base64.b64encode(image_data).decode('utf-8')
        file_size_kb = len(image_data) / 1024
        print(f"   📎 已加载参考图：{character_name} ({file_size_kb:.1f} KB)")
        return f"data:image/png;base64,{base64_data}"
    except Exception as e:
        print(f"   ❌ 加载角色图失败：{e}")
        return None


def extract_image_data(result: dict) -> bytes | None:
    """
    从 API 响应中提取图片数据
    """
    try:
        if "choices" in result and len(result["choices"]) > 0:
            message = result["choices"][0].get("message", {})
            
            if "images" in message:
                images = message["images"]
                if images and len(images) > 0:
                    image_obj = images[0]
                    
                    if "image_url" in image_obj:
                        image_url_obj = image_obj["image_url"]
                        url = image_url_obj if isinstance(image_url_obj, str) else image_url_obj.get("url", "")
                        
                        if url.startswith("data:image"):
                            base64_data = url.split(",")[1]
                            return base64.b64decode(base64_data)
                        elif url.startswith("http"):
                            resp = requests.get(url, timeout=30)
                            if resp.status_code == 200:
                                return resp.content
    except Exception as e:
        print(f"   提取图片失败：{e}")
    
    return None


def generate_image_with_references(prompt_text: str, filename: str, character_refs: list[str] = None, max_retries: int = 5, base_delay: int = 10) -> dict:
    """
    调用 OpenRouter API 生成图片（带角色参考图）
    支持自动重试（指数退避）
    """
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Trilingual Picturebook Generator with References"
    }
    
    # 构建消息内容（文字 + 参考图）
    content_parts = []
    
    # 添加角色参考图
    if character_refs:
        for char_name in character_refs:
            image_url = load_character_image(char_name)
            if image_url:
                content_parts.append({
                    "type": "image_url",
                    "image_url": image_url
                })
    
    # 添加文字提示（包含分辨率要求）
    size_hint = f" Square format {IMAGE_SIZE}." if IMAGE_SIZE else ""
    content_parts.append({
        "type": "text",
        "text": f"Generate an image: {prompt_text}{size_hint}"
    })
    
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "user", "content": content_parts}
        ],
        "response_format": {"type": "image"}
    }
    
    output_path = OUTPUT_DIR / f"{filename}.png"
    
    # 重试逻辑（指数退避）
    last_error = None
    for attempt in range(1, max_retries + 1):
        try:
            if attempt > 1:
                delay = base_delay * (2 ** (attempt - 1))  # 指数退避：10s, 20s, 40s, 80s, 160s
                print(f"   ⏳ 等待 {delay} 秒后重试 ({attempt}/{max_retries})...")
                time.sleep(delay)
            
            print(f"📸 正在生成：{filename}... (尝试 {attempt}/{max_retries})")
            if character_refs and attempt == 1:
                print(f"   🎨 使用角色参考：{', '.join(character_refs)}")
            
            response = requests.post(API_URL, headers=headers, json=payload, timeout=180)
            
            if response.status_code == 200:
                result = response.json()
                image_data = extract_image_data(result)
                
                if image_data:
                    with open(output_path, "wb") as f:
                        f.write(image_data)
                    
                    file_size = len(image_data) / 1024
                    print(f"✅ 已保存：{output_path} ({file_size:.1f} KB)")
                    return {"success": True, "path": str(output_path), "size_kb": file_size}
                else:
                    print(f"⚠️  未找到图片数据")
                    return {"success": False, "path": None, "error": "No image data"}
            elif response.status_code == 429:
                # 限流，继续重试
                last_error = f"Rate limited (429): {response.text[:150]}"
                print(f"   ⚠️  API 限流，将重试...")
                continue
            elif response.status_code >= 500:
                # 服务器错误，可重试
                last_error = f"Server error ({response.status_code})"
                print(f"   ⚠️  服务器错误，将重试...")
                continue
            else:
                # 其他错误，不重试
                error_msg = f"API error {response.status_code}: {response.text[:200]}"
                print(f"❌ {error_msg}")
                return {"success": False, "path": None, "error": error_msg}
                
        except requests.exceptions.Timeout:
            last_error = "Timeout"
            print(f"   ⏱️  请求超时")
            continue
        except Exception as e:
            last_error = str(e)
            print(f"   ❌ 错误：{e}")
            continue
    
    # 所有重试失败
    print(f"❌ {filename} 生成失败（重试 {max_retries} 次后放弃）: {last_error}")
    return {"success": False, "path": None, "error": f"Max retries exceeded: {last_error}"}


def test_model():
    """测试模型连接"""
    print("🔍 测试 Nano Banana 2 模型...")
    
    headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
    
    try:
        response = requests.get("https://openrouter.ai/api/v1/models", headers=headers, timeout=10)
        if response.status_code == 200:
            models = response.json().get("data", [])
            nano_banana = [m for m in models if "gemini-3.1-flash-image" in m.get("id", "")]
            
            if nano_banana:
                model = nano_banana[0]
                print(f"✅ 模型：{model.get('id')}")
                print(f"   名称：{model.get('name', 'N/A')}")
                return True
        print("❌ 未找到模型")
        return False
    except Exception as e:
        print(f"❌ 测试失败：{e}")
        return False


def check_character_references():
    """检查角色参考图是否存在"""
    print("\n📎 检查角色参考图...")
    found = []
    missing = []
    
    for name, path in CHARACTER_REFERENCES.items():
        if path.exists():
            size_kb = path.stat().st_size / 1024
            print(f"   ✅ {name}: {path.name} ({size_kb:.1f} KB)")
            found.append(name)
        else:
            print(f"   ❌ {name}: 未找到 {path}")
            missing.append(name)
    
    return found, missing


def main():
    """主函数"""
    print("=" * 60)
    print("📚 三语绘本图像生成器 (带角色参考)")
    print("🍌 Google Nano Banana 2 (Gemini 3.1 Flash Image)")
    print("=" * 60)
    
    if not OPENROUTER_API_KEY:
        print("❌ 错误：未设置 OPENROUTER_API_KEY")
        return
    
    # 检查角色参考图
    found_chars, missing_chars = check_character_references()
    
    if not found_chars:
        print("\n⚠️  警告：未找到任何角色参考图!")
        print("   请确保以下文件存在:")
        for name, path in CHARACTER_REFERENCES.items():
            print(f"   - {path}")
        print("\n将继续但不会使用角色参考...")
    
    if not test_model():
        print("\n⚠️  模型测试失败，但仍尝试生成...")
    
    # 选择生成模式
    print("\n请选择生成模式:")
    print("1. 仅生成角色设定图 (4 张)")
    print("2. 生成《哥哥妹妹早上好》全书 (16 张)")
    print("3. 生成全部 (20 张)")
    print("4. 自定义页码")
    print("5. 测试生成 (1 张，带参考)")
    
    choice = input("\n请输入选项 (1-5): ").strip()
    
    prompts_to_use = {}
    
    if choice == "1":
        prompts_to_use = CHARACTER_PROMPTS
    elif choice == "2":
        prompts_to_use = GOODMORNING_PROMPTS
    elif choice == "3":
        prompts_to_use = {**CHARACTER_PROMPTS, **GOODMORNING_PROMPTS}
    elif choice == "4":
        print("\n可用页码:")
        print("角色图：character-boy-front, character-boy-side, character-girl-front, character-girl-side")
        print("绘本页：goodmorning-page-01 到 goodmorning-page-16")
        pages = input("请输入页码 (逗号分隔): ").strip()
        for page in pages.split(","):
            page = page.strip()
            if page in CHARACTER_PROMPTS:
                prompts_to_use[page] = CHARACTER_PROMPTS[page]
            elif page in GOODMORNING_PROMPTS:
                prompts_to_use[page] = GOODMORNING_PROMPTS[page]
            else:
                print(f"⚠️  未找到页码：{page}")
    elif choice == "5":
        # 测试：使用角色参考生成一张图
        prompts_to_use = {
            "test-with-ref": {
                "text": CHARACTER_PROMPTS["character-boy-front"]["text"],
                "characters": ["yangyang"]
            }
        }
    else:
        print("❌ 无效选项")
        return
    
    if not prompts_to_use:
        print("❌ 没有要生成的内容")
        return
    
    print(f"\n📝 将生成 {len(prompts_to_use)} 张图片")
    print(f"📁 输出目录：{OUTPUT_DIR}")
    print(f"🍌 模型：{MODEL}")
    print(f"📐 分辨率：{IMAGE_SIZE}")
    print(f"🎨 角色参考：{', '.join(found_chars) if found_chars else '无'}")
    print("\n" + "=" * 60)
    
    # 生成图片
    results = []
    total_size = 0
    
    for filename, prompt_data in prompts_to_use.items():
        # 处理新旧格式兼容
        if isinstance(prompt_data, dict):
            prompt_text = prompt_data.get("text", str(prompt_data))
            character_refs = prompt_data.get("characters", [])
        else:
            prompt_text = prompt_data
            character_refs = []
        
        result = generate_image_with_references(
            prompt_text, 
            filename, 
            character_refs if found_chars else [],
            max_retries=5,
            base_delay=15  # 15s 基础延迟，指数退避：15s, 30s, 60s, 120s, 240s
        )
        results.append((filename, result))
        
        if result["success"]:
            total_size += result.get("size_kb", 0)
        
        # 成功生成后等待 10 秒，避免连续请求触发限流
        if result["success"] and len(prompts_to_use) > 1:
            print(f"   ⏳ 等待 10 秒后继续下一张...")
            time.sleep(10)
    
    # 总结
    print("\n" + "=" * 60)
    print("📊 生成完成!")
    
    success_count = sum(1 for _, r in results if r["success"])
    fail_count = len(results) - success_count
    
    print(f"✅ 成功：{success_count} 张")
    print(f"❌ 失败：{fail_count} 张")
    print(f"💾 总大小：{total_size:.1f} KB")
    print(f"📁 输出目录：{OUTPUT_DIR}")
    
    # 保存生成日志
    log_file = OUTPUT_DIR / f"generation-log-{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
    log_data = {
        "model": MODEL,
        "timestamp": datetime.now().isoformat(),
        "total": len(results),
        "success": success_count,
        "failed": fail_count,
        "total_size_kb": total_size,
        "results": results
    }
    
    with open(log_file, "w", encoding="utf-8") as f:
        json.dump(log_data, f, indent=2, ensure_ascii=False)
    
    print(f"📝 日志已保存：{log_file}")
    print("=" * 60)


if __name__ == "__main__":
    main()