import requests
import re

# API 配置
API_KEY = "YOUR_API_KEY"
API_URL = "https://api.apiyi.com/v1/chat/completions"

def generate_image(prompt, ratio="2:3"):
    """
    使用 Sora Image 生成图片
    
    Args:
        prompt: 图片描述文本
        ratio: 图片比例，支持 "2:3", "3:2", "1:1"
    """
    # 在提示词末尾添加比例标记
    if ratio and ratio in ["2:3", "3:2", "1:1"]:
        prompt = f"{prompt}【{ratio}】"
    
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": "sora_image",
        "messages": [
            {
                "role": "user",
                "content": prompt
            }
        ]
    }
    
    response = requests.post(API_URL, headers=headers, json=payload)
    result = response.json()
    
    # 提取图片 URL
    content = result['choices'][0]['message']['content']
    image_urls = re.findall(r'!\[.*?\]\((https?://[^)]+)\)', content)
    
    return image_urls

# 使用示例
urls = generate_image("一只可爱的猫咪在花园里玩耍", "2:3")
print(f"生成的图片: {urls[0]}")

def batch_generate_images(prompts, ratio="2:3"):
    """批量生成图片"""
    results = []
    
    for prompt in prompts:
        try:
            urls = generate_image(prompt, ratio)
            results.append({
                "prompt": prompt,
                "url": urls[0] if urls else None,
                "success": bool(urls)
            })
            print(f"✅ 成功生成: {prompt}")
        except Exception as e:
            results.append({
                "prompt": prompt,
                "url": None,
                "success": False,
                "error": str(e)
            })
            print(f"❌ 生成失败: {prompt} - {e}")
    
    return results

# 批量生成示例
prompts = [
    "夕阳下的海滩",
    "未来科技城市",
    "梦幻森林中的精灵"
]

results = batch_generate_images(prompts, "3:2")

# 竖版人像
portrait = generate_image("优雅的女士肖像，专业摄影风格【2:3】")

# 横版风景
landscape = generate_image("日出时分的山谷，光线柔和【3:2】")

# 正方形图标
icon = generate_image("简约现代的应用图标设计【1:1】")

# ❌ 不推荐：过于简单
prompt = "猫"

# ✅ 推荐：详细描述
prompt = """
一只毛茸茸的橘色猫咪，
坐在阳光明媚的窗台上，
背景是模糊的城市景观，
摄影风格，高清细节
【2:3】
"""

import time

def generate_with_retry(prompt, max_retries=3):
    """带重试机制的图片生成"""
    for attempt in range(max_retries):
        try:
            urls = generate_image(prompt)
            if urls:
                return urls[0]
        except Exception as e:
            if attempt < max_retries - 1:
                wait_time = 2 ** attempt  # 指数退避
                print(f"重试 {attempt + 1}/{max_retries}，等待 {wait_time}s...")
                time.sleep(wait_time)
            else:
                raise e
    
    return None

import requests
from datetime import datetime

def save_generated_image(url, prompt):
    """保存生成的图片到本地"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"sora_{timestamp}.png"
    
    response = requests.get(url)
    with open(filename, 'wb') as f:
        f.write(response.content)
    
    # 保存提示词信息
    with open(f"sora_{timestamp}_prompt.txt", 'w', encoding='utf-8') as f:
        f.write(f"Prompt: {prompt}\n")
        f.write(f"URL: {url}\n")
        f.write(f"Time: {datetime.now()}\n")
    
    return filename

# 艺术风格模板
art_styles = {
    "油画": "油画风格，厚重的笔触，丰富的色彩层次",
    "水彩": "水彩画风格，透明感，色彩流动",
    "素描": "铅笔素描风格，黑白，细腻的线条",
    "动漫": "日本动漫风格，大眼睛，鲜艳的色彩",
    "写实": "超写实摄影，高清细节，专业摄影"
}

def generate_with_style(subject, style_name):
    """使用预设风格生成图片"""
    style = art_styles.get(style_name, "")
    prompt = f"{subject}，{style}【2:3】"
    return generate_image(prompt)

# 使用示例
url = generate_with_style("美丽的玫瑰花", "水彩")

# 场景生成模板
def generate_scene(subject, time="日落", weather="晴朗", mood="宁静"):
    """根据参数生成场景图片"""
    prompt = f"""
    {subject}
    时间：{time}
    天气：{weather}
    氛围：{mood}
    摄影风格，专业构图
    【3:2】
    """
    return generate_image(prompt.strip())

# 生成不同场景
sunset_beach = generate_scene("海滩", "日落", "微风", "浪漫")
morning_forest = generate_scene("森林小径", "清晨", "薄雾", "神秘")

def generate_variations(base_prompt, variations, ratio="2:3"):
    """生成同一主题的多个变体"""
    results = []
    
    for variation in variations:
        full_prompt = f"{base_prompt}，{variation}【{ratio}】"
        try:
            url = generate_image(full_prompt)[0]
            results.append({
                "variation": variation,
                "url": url,
                "prompt": full_prompt
            })
        except Exception as e:
            print(f"变体生成失败: {variation}")
    
    return results

# 生成猫咪的不同变体
base = "一只可爱的猫咪"
variations = [
    "橘色虎斑",
    "黑白奶牛色",
    "纯白色长毛",
    "灰色短毛"
]

cat_variations = generate_variations(base, variations)

def calculate_cost(num_images):
    """计算生成图片的成本"""
    cost_per_image = 0.01  # $0.01 per image
    total_cost = num_images * cost_per_image
    
    # 与 GPT-Image-1 对比
    # GPT-Image-1 按 Token 计费，假设平均每张图约 1000 tokens
    avg_tokens_per_image = 1000
    gpt_image_1_cost = num_images * (avg_tokens_per_image / 1_000_000) * 40  # 输出 token费用
    
    print(f"Sora Image 成本计算:")
    print(f"- 图片数量: {num_images}")
    print(f"- 单价: ${cost_per_image} (固定价格)")
    print(f"- 总成本: ${total_cost:.2f}")
    print(f"\n对比 GPT-Image-1:")
    print(f"- GPT-Image-1 计费: 按Token ($10输入/$40输出 per M)")
    print(f"- 预估成本: ${gpt_image_1_cost:.2f} (假设每张~1000 tokens)")
    print(f"- Sora 优势: 固定价格，成本可控")

# 计算示例
calculate_cost(100)  # 生成 100 张图片