#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
FLUX风格塑形API的Python客户端示例
"""

import requests
import argparse
import os
from datetime import datetime

def generate_image(server_url, prompt, structure_image_path, style_image_path, depth_strength=15.0, style_strength=0.5):
    """
    调用API生成风格化图像
    
    参数:
        server_url (str): API服务器地址
        prompt (str): 文本提示
        structure_image_path (str): 结构图像的路径
        style_image_path (str): 风格图像的路径
        depth_strength (float): 深度强度
        style_strength (float): 风格强度
    
    返回:
        str: 生成图像的保存路径
    """
    # API端点
    url = f"{server_url}/generate"
    
    # 参数
    data = {
        "prompt": prompt,
        "depth_strength": depth_strength,
        "style_strength": style_strength
    }
    
    # 图像文件
    files = {
        "structure_image": open(structure_image_path, "rb"),
        "style_image": open(style_image_path, "rb")
    }
    
    print(f"正在发送请求到 {url}...")
    print(f"参数: 提示词='{prompt}', 深度强度={depth_strength}, 风格强度={style_strength}")
    print(f"结构图像: {structure_image_path}")
    print(f"风格图像: {style_image_path}")
    
    try:
        # 发送请求
        response = requests.post(url, data=data, files=files)
        
        # 检查响应
        if response.status_code == 200:
            # 生成输出文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_filename = f"生成图像_{timestamp}.png"
            
            # 保存生成的图像
            with open(output_filename, "wb") as f:
                f.write(response.content)
            
            print(f"成功! 图像已保存到 {output_filename}")
            return output_filename
        else:
            print(f"生成失败: {response.json()}")
            return None
    except Exception as e:
        print(f"请求出错: {e}")
        return None
    finally:
        # 关闭文件
        for f in files.values():
            f.close()

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="FLUX风格塑形API客户端")
    parser.add_argument("--server", default="http://localhost:8000", help="API服务器地址")
    parser.add_argument("--prompt", default="", help="文本提示")
    parser.add_argument("--structure", required=True, help="结构图像路径")
    parser.add_argument("--style", required=True, help="风格图像路径")
    parser.add_argument("--depth-strength", type=float, default=15.0, help="深度强度 (默认: 15.0)")
    parser.add_argument("--style-strength", type=float, default=0.5, help="风格强度 (默认: 0.5)")
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.structure):
        print(f"错误: 结构图像文件不存在 {args.structure}")
        return
    
    if not os.path.exists(args.style):
        print(f"错误: 风格图像文件不存在 {args.style}")
        return
    
    # 调用API
    generate_image(
        server_url=args.server,
        prompt=args.prompt,
        structure_image_path=args.structure,
        style_image_path=args.style,
        depth_strength=args.depth_strength,
        style_strength=args.style_strength
    )

if __name__ == "__main__":
    main()