#!/usr/bin/env python3
"""
第二本书《食饭啦》图像生成脚本
使用 Gemini 3.1 Flash Image Preview + 角色参考图
主题：饮食认知、吃饭 routine
"""

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"
MODEL = "google/gemini-3.1-flash-image-preview"
IMAGE_SIZE = "512x512"  # 1k 质量

# 输出目录 - 第二本书
OUTPUT_DIR = Path("/root/.openclaw/workspace/trilingual-picturebook/output/book02-shifan")
OUTPUT_DIR.mkdir(exist_ok=True)

# 角色参考图
CHARACTER_SHEET_DIR = Path("/root/.openclaw/workspace/trilingual-picturebook")
YANGYANG_REF = CHARACTER_SHEET_DIR / "_charsheet_yangyang.png"
ZHUAZHUA_REF = CHARACTER_SHEET_DIR / "_charsheet_zhuazhua.png"


def load_image_as_base64(image_path: Path) -> str:
    with open(image_path, "rb") as f:
        image_data = f.read()
    base64_data = base64.b64encode(image_data).decode('utf-8')
    return f"data:image/png;base64,{base64_data}"


def generate_image(prompt_text: str, filename: str):
    """调用 Gemini API 生成图片"""
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Book02 Shifan"
    }
    
    content_parts = []
    
    # 添加角色参考图
    if YANGYANG_REF.exists():
        content_parts.append({
            "type": "image_url",
            "image_url": load_image_as_base64(YANGYANG_REF)
        })
    
    if ZHUAZHUA_REF.exists():
        content_parts.append({
            "type": "image_url",
            "image_url": load_image_as_base64(ZHUAZHUA_REF)
        })
    
    content_parts.append({
        "type": "text",
        "text": prompt_text
    })
    
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": content_parts}],
        "max_tokens": 1000
    }
    
    max_retries = 6
    base_delay = 20
    
    for attempt in range(max_retries):
        try:
            response = requests.post(API_URL, headers=headers, json=payload, timeout=180)
            
            if response.status_code == 429:
                wait_time = min(base_delay * (2 ** attempt), 300)
                print(f"   ⏳ 限流，等待 {wait_time} 秒...")
                time.sleep(wait_time)
                continue
            
            response.raise_for_status()
            result = response.json()
            
            if "choices" in result and len(result["choices"]) > 0:
                message = result["choices"][0].get("message", {})
                if "images" in message and len(message["images"]) > 0:
                    image_url = message["images"][0].get("image_url", "")
                    if isinstance(image_url, dict):
                        image_url = image_url.get("url", "")
                    
                    if image_url.startswith("data:image"):
                        image_data = base64.b64decode(image_url.split(",")[1])
                    else:
                        image_resp = requests.get(image_url, timeout=60)
                        image_resp.raise_for_status()
                        image_data = image_resp.content
                    
                    output_path = OUTPUT_DIR / f"{filename}.png"
                    with open(output_path, "wb") as f:
                        f.write(image_data)
                    
                    file_size_kb = len(image_data) / 1024
                    return {"success": True, "path": str(output_path), "size_kb": file_size_kb}
            
            return {"success": False, "error": "No image"}
            
        except Exception as e:
            print(f"   ❌ 错误：{e}")
            if attempt < max_retries - 1:
                wait_time = min(base_delay * (2 ** attempt), 300)
                time.sleep(wait_time)
            else:
                return {"success": False, "error": str(e)}
    
    return {"success": False, "error": "Max retries"}


# 《食饭啦》16 页提示词
SHIFAN_PROMPTS = {
    "page-01": "children's book illustration, watercolor style, kitchen scene, round wooden dining table with colorful tablecloth, steaming bowls of food on table, warm family atmosphere, soft morning light from window, boy in blue outfit and girl in pink dress with pigtails sitting at table waiting, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-02": "children's book illustration, watercolor style, 2 year old Chinese boy (yangyang: short black hair, blue outfit) washing hands at bathroom sink, water bubbles, soap suds, happy expression, bathroom scene with towel, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-03": "children's book illustration, watercolor style, 2 year old Chinese girl (zhuazhua: black hair with two pigtails, pink dress) washing hands at sink next to brother, water flowing, cute smile, bathroom scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-04": "children's book illustration, watercolor style, both children drying hands with colorful towels, boy in blue and girl in pink standing side by side, happy expressions, bathroom or kitchen background, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-05": "children's book illustration, watercolor style, grandmother (Chinese grandma with warm smile) serving food from kitchen, carrying bowls of rice, boy and girl watching excitedly at dining table, warm kitchen scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-06": "children's book illustration, watercolor style, close-up of dining table with various Chinese foods, white rice in bowls, green vegetables, steamed fish, soup, colorful and appetizing, chopsticks and spoons, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-07": "children's book illustration, watercolor style, boy (yangyang) holding chopsticks trying to pick up food, concentrating expression, food on plate, girl watching and learning beside him, dining table scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-08": "children's book illustration, watercolor style, girl (zhuazhua) using spoon to eat rice, some rice on her cheek, cute messy eating, boy laughing gently beside her, dining table scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-09": "children's book illustration, watercolor style, close-up of vegetables on plate, broccoli, carrots, green beans, colorful and fresh, children's hands reaching for veggies with chopsticks, educational food scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-10": "children's book illustration, watercolor style, boy eating vegetables with happy expression, thumb up gesture, girl watching and trying vegetables too, encouraging scene, dining table, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-11": "children's book illustration, watercolor style, bowl of soup with steam rising, spoon in soup, children blowing on soup to cool it down, cute expressions, close-up food scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-12": "children's book illustration, watercolor style, both children drinking soup carefully from spoons, satisfied expressions, warm cozy feeling, dining table scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-13": "children's book illustration, watercolor style, almost empty plates and bowls, children patting their full tummies, satisfied happy expressions, finished meal scene, dining table with crumbs, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-14": "children's book illustration, watercolor style, boy helping to carry empty bowl to kitchen, responsible expression, girl following with her bowl, learning to help, kitchen background, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-15": "children's book illustration, watercolor style, grandmother praising the children, thumbs up gesture, children looking proud and happy, warm family moment in kitchen or dining area, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
    
    "page-16": "children's book illustration, watercolor style, both children saying thank you with hands together (gesture), bowing slightly to grandmother, grateful expressions, warm ending scene, soft pastel colors, for toddlers picture book, square composition, no text, no watermark",
}


def main():
    print("=" * 60)
    print("📖 第二本书《食饭啦》生成开始")
    print("=" * 60)
    print(f"📁 输出目录：{OUTPUT_DIR}")
    print(f"📐 图片尺寸：{IMAGE_SIZE}")
    print(f"🎨 模型：{MODEL}")
    print(f"📎 角色参考：已加载")
    print("-" * 60)
    
    results = []
    total = len(SHIFAN_PROMPTS)
    
    for i, (page_name, prompt) in enumerate(SHIFAN_PROMPTS.items(), 1):
        print(f"\n[{i}/{total}] 📸 {page_name}")
        result = generate_image(prompt, page_name)
        results.append((page_name, result))
        
        if result.get("success"):
            print(f"   ✅ {result['size_kb']:.1f} KB")
        else:
            print(f"   ❌ {result.get('error', 'Failed')}")
        
        # 每页间隔 5 秒，避免限流
        if i < total:
            time.sleep(5)
    
    # 保存日志
    log_path = OUTPUT_DIR / f"generation-log-{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
    with open(log_path, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    success_count = sum(1 for _, r in results if r.get("success"))
    
    print("\n" + "=" * 60)
    print(f"✅ 完成：{success_count}/{total} 页")
    print(f"📋 日志：{log_path}")
    print("=" * 60)


if __name__ == "__main__":
    main()