#!/usr/bin/env python3
"""
三语绘本图像生成脚本 - WellAPI 版本
使用 Flux Pro / DALL-E 3 生成高质量绘本插图
"""

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

# API 配置 - WellAPI
WELL_API_KEY = os.environ.get('WELL_API_KEY')
API_URL = "https://wellapi.ai/v1/images/generations"

# 图像生成模型
# 可选: "flux-pro", "flux.1-dev", "dall-e-3", "gemini-2.5-flash-image-preview"
MODEL = "flux-pro"
IMAGE_SIZE = "1024x1024"  # 可选: "512x512", "768x768", "1024x1024"

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

# 《哥哥妹妹早上好》16 页提示词
GOODMORNING_PROMPTS = {
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
    "well-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",
}


def generate_image(prompt: str, filename: str) -> dict:
    """
    调用 WellAPI 生成图片
    """
    headers = {
        "Authorization": f"Bearer {WELL_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": MODEL,
        "prompt": prompt,
        "size": IMAGE_SIZE,
        "n": 1
    }
    
    output_path = OUTPUT_DIR / f"{filename}.png"
    
    try:
        print(f"📸 正在生成：{filename}...")
        
        response = requests.post(API_URL, headers=headers, json=payload, timeout=120)
        
        if response.status_code == 200:
            result = response.json()
            
            if "data" in result and len(result["data"]) > 0:
                image_url = result["data"][0].get("url")
                
                if image_url:
                    # 下载图片
                    img_response = requests.get(image_url, timeout=30)
                    if img_response.status_code == 200:
                        with open(output_path, "wb") as f:
                            f.write(img_response.content)
                        
                        file_size = len(img_response.content) / 1024
                        print(f"✅ 已保存：{output_path} ({file_size:.1f} KB)")
                        return {"success": True, "path": str(output_path), "size_kb": file_size}
                    else:
                        return {"success": False, "path": None, "error": f"Download failed: {img_response.status_code}"}
                else:
                    return {"success": False, "path": None, "error": "No image URL in response"}
            else:
                return {"success": False, "path": None, "error": "No data in response"}
        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:
        return {"success": False, "path": None, "error": "Timeout"}
    except Exception as e:
        return {"success": False, "path": None, "error": str(e)}


def main():
    """主函数"""
    print("=" * 60)
    print("📚 三语绘本图像生成器 (WellAPI)")
    print(f"🎨 模型：{MODEL}")
    print(f"📐 尺寸：{IMAGE_SIZE}")
    print("=" * 60)
    
    if not WELL_API_KEY:
        print("❌ 错误：未设置 WELL_API_KEY")
        return
    
    print(f"\n📝 将生成 {len(GOODMORNING_PROMPTS)} 张图片")
    print(f"📁 输出目录：{OUTPUT_DIR}")
    print("\n" + "=" * 60)
    
    # 生成图片
    results = []
    total_size = 0
    
    for filename, prompt in GOODMORNING_PROMPTS.items():
        result = generate_image(prompt, filename)
        results.append((filename, result))
        
        if result["success"]:
            total_size += result.get("size_kb", 0)
        
        # API 限流控制 - WellAPI 也需要一定延迟
        time.sleep(3)
    
    # 总结
    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"wellapi-log-{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
    log_data = {
        "model": MODEL,
        "size": IMAGE_SIZE,
        "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()
