#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
双图合成处理 - 基于 multi_image_merge 方法
使用 Gemini 3 Pro Image Preview 模型进行双图创意合成

特点：
- 支持本地图片和网络图片 URL 两种输入方式
- 支持自定义提示词
- 创意合成、场景融合

使用方法：
1. 安装依赖：pip install openai requests
2. 在下方配置区填入您的 API Key
3. 修改文件底部的参数变量
4. 运行：python 双图合成处理.py

获取 API Key：https://api.laozhang.ai/token
"""

import base64
import re
import os
from openai import OpenAI

# ========== 配置区（请填入您的 API Key）==========
API_KEY = "sk-KZ9TGgZAWnzY1T3c90B277F72e184c26A53f22440e08E86e"  # 替换为您的 API Key
BASE_URL = "https://api.laozhang.ai/v1"

# 模型选择
MODEL = "gemini-3-pro-image-preview"  # 最新版，支持更高质量
# MODEL = "gemini-2.5-flash-image"    # 稳定版，价格更低（$0.025 vs $0.05）
# ================================================


def extract_and_save_image(content: str, filename: str) -> bool:
    """从响应内容中提取 base64 图片并保存"""
    # 匹配 markdown 格式的 base64 图片
    match = re.search(r'!\[.*?\]\((data:image/\w+;base64,([^)]+))\)', content)
    if match:
        base64_data = match.group(2)
        # 确保 base64 填充正确
        padding = 4 - len(base64_data) % 4
        if padding != 4:
            base64_data += '=' * padding
        image_data = base64.b64decode(base64_data)
        with open(filename, 'wb') as f:
            f.write(image_data)
        return True
    return False


def is_url(path: str) -> bool:
    """判断路径是否为 URL"""
    return path.startswith('http://') or path.startswith('https://')


def multi_image_merge(image1_path: str, image2_path: str, prompt: str, output_filename: str = None) -> bool:
    """
    多图合成
    
    参数:
        image1_path: 第一张图片路径（本地文件路径或 URL）
        image2_path: 第二张图片路径（本地文件路径或 URL）
        prompt: 合成提示词
        output_filename: 输出文件名（可选）
    
    返回: 是否成功
    """
    print("\n" + "="*60)
    print("🎨 双图合成处理")
    print("="*60)
    
    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
    
    # 构建 content 数组
    content = [{"type": "text", "text": prompt}]
    
    # 处理第一张图片
    if is_url(image1_path):
        # 网络图片 URL
        content.append({"type": "image_url", "image_url": {"url": image1_path}})
        print(f"🖼️  图片1 (URL): {image1_path[:50]}...")
    else:
        # 本地图片 Base64
        if not os.path.exists(image1_path):
            print(f"❌ 错误：图片1不存在: {image1_path}")
            return False
        with open(image1_path, "rb") as f:
            image1_b64 = base64.b64encode(f.read()).decode("utf-8")
        content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{image1_b64}"
            }
        })
        print(f"🖼️  图片1 (本地): {image1_path}")
    
    # 处理第二张图片
    if is_url(image2_path):
        # 网络图片 URL
        content.append({"type": "image_url", "image_url": {"url": image2_path}})
        print(f"🖼️  图片2 (URL): {image2_path[:50]}...")
    else:
        # 本地图片 Base64
        if not os.path.exists(image2_path):
            print(f"❌ 错误：图片2不存在: {image2_path}")
            return False
        with open(image2_path, "rb") as f:
            image2_b64 = base64.b64encode(f.read()).decode("utf-8")
        content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{image2_b64}"
            }
        })
        print(f"🖼️  图片2 (本地): {image2_path}")
    
    print(f"📝 提示词: {prompt}")
    print("📡 发送请求...")
    
    try:
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {
                    "role": "user",
                    "content": content
                }
            ]
        )
        
        result_content = response.choices[0].message.content
        
        # 确定输出文件名
        if output_filename:
            output_file = output_filename
        else:
            output_file = "result_multi_merge.png"
        
        if extract_and_save_image(result_content, output_file):
            print(f"✅ 成功！图片已保存: {output_file}")
            return True
        else:
            print(f"⚠️ 未找到图片数据")
            return False
            
    except Exception as e:
        print(f"❌ 错误: {str(e)}")
        return False


def main():
    """主函数"""
    print("="*60)
    print("🎨 双图合成处理")
    print(f"模型: {MODEL}")
    print("="*60)
    
    # 检查 API Key
    if API_KEY == "sk-YOUR_API_KEY" or not API_KEY:
        print("\n❌ 请先配置您的 API Key！")
        print("   获取地址: https://api.laozhang.ai/token")
        print("   然后修改脚本顶部的 API_KEY 变量")
        return
    
    # 检查图片是否存在（仅本地文件）
    if not is_url(IMAGE1_PATH) and not os.path.exists(IMAGE1_PATH):
        print(f"\n❌ 错误：图片1文件不存在: {IMAGE1_PATH}")
        print("   请检查文件路径是否正确")
        return
    
    if not is_url(IMAGE2_PATH) and not os.path.exists(IMAGE2_PATH):
        print(f"\n❌ 错误：图片2文件不存在: {IMAGE2_PATH}")
        print("   请检查文件路径是否正确")
        return
    
    print(f"\n📋 配置信息：")
    print(f"   图片1: {IMAGE1_PATH}")
    print(f"   图片2: {IMAGE2_PATH}")
    print(f"   提示词: {PROMPT}")
    print()
    
    # 执行合成
    result = multi_image_merge(
        image1_path=IMAGE1_PATH,
        image2_path=IMAGE2_PATH,
        prompt=PROMPT,
        output_filename=OUTPUT_FILENAME
    )
    
    if result:
        print("\n" + "="*60)
        print("🎉 合成完成！")
        print("="*60)
    else:
        print("\n" + "="*60)
        print("❌ 合成失败，请检查错误信息")
        print("="*60)


# ========== 参数配置区（请修改以下参数）==========

# 1. 第一张图片路径（支持本地文件路径或网络 URL）
IMAGE1_PATH = "5.png"  # 本地文件示例: "1.png" 或 "C:/images/photo1.jpg"
# IMAGE1_PATH = "https://images.unsplash.com/photo-1514888286974-6c03e2ca1dba?w=800&q=80"  # URL 示例

# 2. 第二张图片路径（支持本地文件路径或网络 URL）
IMAGE2_PATH = "6.png"  # 本地文件示例: "2.png" 或 "C:/images/photo2.jpg"
# IMAGE2_PATH = "https://images.unsplash.com/photo-1560806887-1e4cd0b6cbd6?w=800&q=80"  # URL 示例

# 3. 合成提示词（描述如何合成这两张图片）
PROMPT = "把第2张图的衣股花型图案替换成图1的花型图案"  # 修改为您想要的合成效果描述

# 4. 输出文件名（可选，留空则使用默认名称）
OUTPUT_FILENAME = None  # 例如: "合成结果.png" 或 None（使用默认名称）

# ================================================


if __name__ == "__main__":
    main()