#!/usr/bin/env python3
"""
重新生成第 13 页 - 强化角色一致性
使用更详细的角色描述 + 参考图
"""

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"

# 输出目录
OUTPUT_DIR = Path("/root/.openclaw/workspace/trilingual-picturebook/output")
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"

# 图像分辨率 - 1k 质量
IMAGE_SIZE = "512x512"


def load_image_as_base64(image_path: Path) -> str:
    """加载图片为 base64 data URL"""
    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, system_prompt: str = None):
    """生成单张图片"""
    print(f"\n📸 生成：{filename}")
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://openclaw.ai",
        "X-Title": "Picturebook Regenerator v2"
    }
    
    # 构建消息内容
    content_parts = []
    
    # 添加参考图
    if YANGYANG_REF.exists():
        print(f"   📎 参考图：yangyang ({YANGYANG_REF.stat().st_size // 1024} KB)")
        content_parts.append({
            "type": "image_url",
            "image_url": load_image_as_base64(YANGYANG_REF)
        })
    
    if ZHUAZHUA_REF.exists():
        print(f"   📎 参考图：zhuazhua ({ZHUAZHUA_REF.stat().st_size // 1024} KB)")
        content_parts.append({
            "type": "image_url",
            "image_url": load_image_as_base64(ZHUAZHUA_REF)
        })
    
    # 添加文字提示词
    content_parts.append({
        "type": "text",
        "text": prompt_text
    })
    
    messages = [{"role": "user", "content": content_parts}]
    
    # 可选的系统提示词
    if system_prompt:
        messages.insert(0, {"role": "system", "content": system_prompt})
    
    payload = {
        "model": MODEL,
        "messages": messages,
        "max_tokens": 1000
    }
    
    # 发送请求（带重试，更长间隔）
    max_retries = 8
    base_delay = 15
    
    for attempt in range(max_retries):
        try:
            print(f"   🔄 尝试 {attempt + 1}/{max_retries}...")
            response = requests.post(API_URL, headers=headers, json=payload, timeout=180)
            
            if response.status_code == 429:
                wait_time = base_delay * (2 ** attempt)
                if wait_time > 300:
                    wait_time = 300  # 最多等 5 分钟
                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", "")
                    
                    print(f"   📥 下载图片...")
                    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
                    print(f"   ✅ 已保存：{filename} ({file_size_kb:.1f} KB)")
                    return {"success": True, "path": str(output_path), "size_kb": file_size_kb}
            
            print(f"   ❌ 未返回图片数据")
            return {"success": False, "error": "No image data"}
            
        except requests.exceptions.Timeout as e:
            print(f"   ⏱️  请求超时：{e}")
            if attempt < max_retries - 1:
                wait_time = base_delay * (2 ** attempt)
                time.sleep(wait_time)
            else:
                return {"success": False, "error": "Timeout"}
        except Exception as e:
            print(f"   ❌ 错误：{e}")
            if attempt < max_retries - 1:
                wait_time = base_delay * (2 ** attempt)
                time.sleep(wait_time)
            else:
                return {"success": False, "error": str(e)}
    
    return {"success": False, "error": "Max retries exceeded"}


def main():
    print("=" * 60)
    print("📖 重新生成第 13 页（强化角色一致性）")
    print("=" * 60)
    print(f"📁 输出目录：{OUTPUT_DIR}")
    print(f"📐 图片尺寸：{IMAGE_SIZE}")
    
    # 系统提示词 - 强调角色一致性
    system_prompt = """You are a professional children's book illustrator specializing in consistent character design.
When provided with character reference images, you MUST maintain exact character appearance including:
- Facial features (eye shape, nose, mouth, face shape)
- Hairstyle and hair color
- Clothing style and colors
- Age appearance and proportions
- Overall art style (watercolor, soft edges)

Reference images are provided for character consistency. Study them carefully before generating."""
    
    # 第 13 页提示词 - 极度详细的角色描述
    prompt_13 = """
【CRITICAL: Character Consistency】
Please study the provided reference images carefully. The characters MUST match exactly:

**Boy Character (from reference image "yangyang"):**
- Age: 2 year old Chinese toddler boy
- Face: Round chubby face, big bright eyes, small nose, cute smile
- Hair: Short black hair, neat cut, slightly rounded bangs
- Outfit: Light blue t-shirt and shorts (casual home wear)
- Expression: Gentle, caring, protective big brother look

**Girl Character (from reference image "zhuazhua"):**
- Age: 2 year old Chinese toddler girl (twin sister)
- Face: Round chubby face, big bright eyes, similar features to brother
- Hair: Black hair styled in TWO SMALL PIGTAILS on top of head, cute bangs
- Outfit: Pink dress (casual home wear)
- Expression: Happy, trusting, following her brother

【Scene Description】
Setting: Warm wooden staircase inside a cozy home, morning sunlight streaming through nearby window

Action: 
- Boy (yangyang) is walking down the stairs FIRST, holding girl's hand
- Girl (zhuazhua) is FOLLOWING brother, one step behind
- Boy is looking back at girl with caring expression, guiding her safely
- Girl is looking at the steps carefully, trusting her brother
- They are holding HANDS - this is the key interaction

Atmosphere:
- Warm golden morning light
- Soft shadows
- Cozy home feeling
- Sibling love and care

【Art Style】
- Watercolor children's book illustration
- Soft pastel colors (blues, pinks, warm wood tones)
- Gentle brush strokes, no harsh lines
- Square composition
- NO text, NO watermark
- High quality, suitable for 2 year old toddlers

【Important】
The characters must look EXACTLY like the reference images provided. This is page 13 of a consistent picture book series.
"""
    
    # 生成第 13 页
    result_13 = generate_image(prompt_13, "gemini-page-13-v2", system_prompt)
    
    # 保存结果
    log_data = [("gemini-page-13-v2", result_13)]
    log_path = OUTPUT_DIR / f"regenerate-log-{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
    with open(log_path, "w", encoding="utf-8") as f:
        json.dump(log_data, f, ensure_ascii=False, indent=2)
    
    print(f"\n📋 结果日志：{log_path}")
    print("\n" + "=" * 60)
    
    if result_13["success"]:
        print("✅ 第 13 页生成成功！")
    else:
        print(f"❌ 第 13 页生成失败：{result_13.get('error', 'Unknown')}")
    
    print("=" * 60)


if __name__ == "__main__":
    main()