#!/usr/bin/env python3
"""
三语绘本图像生成脚本 - WellAPI + Gemini 3.1 Flash Image
使用图像参考确保角色一致性
"""

import os
import requests
import json
import base64
import re
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/chat/completions"

# 模型：Gemini 3.1 Flash Image Preview
MODEL = "gemini-3.1-flash-image-preview"

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

# 角色参考图路径
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",
}

# 16 页提示词（带角色引用）
PROMPTS = {
    "gemini-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", "chars": []},
    "gemini-page-02": {"text": "children's book illustration, watercolor style, 2 year old Chinese boy 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", "chars": ["yangyang"]},
    "gemini-page-03": {"text": "children's book illustration, watercolor style, 2 year old Chinese girl 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", "chars": ["zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang"]},
    "gemini-page-06": {"text": "children's book illustration, watercolor style, 2 year old Chinese girl 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", "chars": ["zhuazhua", "yangyang"]},
    "gemini-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", "chars": ["yangyang"]},
    "gemini-page-08": {"text": "children's book illustration, watercolor style, 2 year old Chinese girl drying face, using soft white towel, clean fresh face, happy expression, pink pajamas, bathroom scene, soft pastel colors, warm lighting", "chars": ["zhuazhua"]},
    "gemini-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", "chars": ["yangyang"]},
    "gemini-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", "chars": ["zhuazhua", "yangyang"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
    "gemini-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", "chars": ["yangyang", "zhuazhua"]},
}


def load_character_image(char_name: str) -> str | None:
    """加载角色图并转 base64"""
    path = CHARACTER_REFERENCES.get(char_name)
    if not path or not path.exists():
        return None
    try:
        with open(path, "rb") as f:
            data = base64.b64encode(f.read()).decode('utf-8')
        return f"data:image/png;base64,{data}"
    except Exception as e:
        print(f"   ❌ 加载失败: {e}")
        return None


def extract_image_from_result(result: dict) -> bytes | None:
    """从 WellAPI 响应中提取图片"""
    try:
        msg = result["choices"][0]["message"]
        
        # 方式 1: content 中的 markdown 图片 ![image](data:image/png;base64,...)
        if "content" in msg and msg["content"]:
            content = msg["content"]
            if "![image](data:image/" in content:
                match = re.search(r'!\[image\]\((data:image/[^;]+;base64,([^)]+))\)', content)
                if match:
                    return base64.b64decode(match.group(2))
        
        # 方式 2: images 数组
        if "images" in msg and msg["images"]:
            url = msg["images"][0].get("image_url", {}).get("url", "")
            if url.startswith("data:image"):
                return base64.b64decode(url.split(",")[1])
            elif url.startswith("http"):
                return requests.get(url, timeout=30).content
    except Exception as e:
        print(f"   提取失败: {e}")
    return None


def generate_with_refs(filename: str, prompt_data: dict) -> dict:
    """生成图片（带角色参考）"""
    headers = {
        "Authorization": f"Bearer {WELL_API_KEY}",
        "Content-Type": "application/json"
    }
    
    content = []
    chars = prompt_data.get("chars", [])
    
    # 添加角色参考图
    for char in chars:
        img_url = load_character_image(char)
        if img_url:
            content.append({"type": "image_url", "image_url": {"url": img_url}})
            print(f"   📎 参考图: {char}")
    
    # 添加文字提示
    content.append({"type": "text", "text": f"Generate an image: {prompt_data['text']}"})
    
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": content}],
        "response_format": {"type": "image"}
    }
    
    output_path = OUTPUT_DIR / f"{filename}.png"
    
    try:
        print(f"📸 生成: {filename}...")
        resp = requests.post(API_URL, headers=headers, json=payload, timeout=180)
        
        if resp.status_code == 200:
            result = resp.json()
            img_data = extract_image_from_result(result)
            
            if img_data:
                with open(output_path, "wb") as f:
                    f.write(img_data)
                
                size_kb = len(img_data) / 1024
                print(f"✅ 已保存: {filename} ({size_kb:.1f} KB)")
                return {"success": True, "path": str(output_path), "size_kb": size_kb}
            else:
                return {"success": False, "error": "无法提取图片"}
        else:
            return {"success": False, "error": f"HTTP {resp.status_code}: {resp.text[:200]}"}
    
    except Exception as e:
        return {"success": False, "error": str(e)}


def main():
    print("=" * 60)
    print("📚 绘本生成 - WellAPI + Gemini 3.1 Flash Image")
    print("=" * 60)
    
    if not WELL_API_KEY:
        print("❌ WELL_API_KEY 未设置")
        return
    
    # 检查参考图
    print("\n📎 检查角色参考图...")
    for name, path in CHARACTER_REFERENCES.items():
        if path.exists():
            print(f"   ✅ {name}: {path.name}")
        else:
            print(f"   ❌ {name}: 未找到")
    
    # 检查已有文件，计算需要生成的
    to_generate = []
    for filename, prompt_data in PROMPTS.items():
        output_path = OUTPUT_DIR / f"{filename}.png"
        if output_path.exists():
            size_mb = output_path.stat().st_size / 1024 / 1024
            print(f"   ⏭️  {filename}.png 已存在 ({size_mb:.1f} MB)")
        else:
            to_generate.append((filename, prompt_data))
    
    if not to_generate:
        print("\n✅ 所有图片已生成完毕！")
        return
    
    print(f"\n📝 总共 {len(PROMPTS)} 张，已生成 {len(PROMPTS) - len(to_generate)} 张")
    print(f"🎯 本次将生成剩余 {len(to_generate)} 张图片")
    print("=" * 60)
    
    results = []
    total_size = 0
    
    for filename, prompt_data in to_generate:
        result = generate_with_refs(filename, prompt_data)
        results.append((filename, result))
        
        if result["success"]:
            total_size += result.get("size_kb", 0)
        
        time.sleep(5)  # 限流控制
    
    # 总结
    success = sum(1 for _, r in results if r["success"])
    print(f"\n📊 本次完成! 成功: {success}/{len(results)}, 新增: {total_size:.1f} KB")
    
    # 统计总数
    existing = len([f for f in OUTPUT_DIR.glob("gemini-page-*.png")])
    print(f"📁 目前共生成: {existing}/{len(PROMPTS)} 张")
    
    # 保存日志
    log = OUTPUT_DIR / f"gemini-ref-log-{datetime.now():%Y%m%d-%H%M%S}.json"
    with open(log, "w") as f:
        json.dump({"results": results, "total_kb": total_size}, f, indent=2)


if __name__ == "__main__":
    main()
