"""
Leonardo AI 集成方案 - 比 Stability AI 便宜 50%
API 文档: https://docs.leonardo.ai/
"""
import os
import httpx
import base64
from typing import Optional

LEONARDO_API_KEY = os.getenv("LEONARDO_API_KEY")
LEONARDO_API_URL = "https://cloud.leonardo.ai/api/rest/v1"


async def generate_with_leonardo(prompt: str, size: str = "1024x1024") -> bytes:
    """
    调用 Leonardo AI API 生成图片
    成本: ~$0.001-0.003/张 (比 Stability AI 便宜 50%)
    """
    if not LEONARDO_API_KEY:
        raise ValueError("LEONARDO_API_KEY not set")
    
    width, height = map(int, size.split('x'))
    
    async with httpx.AsyncClient() as client:
        # Step 1: 创建生成任务
        response = await client.post(
            f"{LEONARDO_API_URL}/generations",
            headers={
                "Authorization": f"Bearer {LEONARDO_API_KEY}",
                "Content-Type": "application/json"
            },
            json={
                "prompt": prompt,
                "width": width,
                "height": height,
                "num_images": 1,
                "modelId": "6b645e3a-d64f-4341-a370-20de11a56288",  # SDXL 模型
                "guidance_scale": 7,
                "num_inference_steps": 30
            },
            timeout=30.0
        )
        
        if response.status_code != 200:
            raise Exception(f"Leonardo API Error: {response.text}")
        
        generation_id = response.json()["sdGenerationJob"]["generationId"]
        
        # Step 2: 轮询任务状态（简化版，实际应循环检查）
        import asyncio
        await asyncio.sleep(10)  # 等待生成完成
        
        # Step 3: 获取生成结果
        result = await client.get(
            f"{LEONARDO_API_URL}/generations/{generation_id}",
            headers={"Authorization": f"Bearer {LEONARDO_API_KEY}"}
        )
        
        if result.status_code != 200:
            raise Exception(f"Failed to get generation result: {result.text}")
        
        data = result.json()
        image_url = data["generations_by_pk"]["generated_images"][0]["url"]
        
        # Step 4: 下载图片
        img_response = await client.get(image_url)
        return img_response.content


# 成本对比 (生成 1000 张图片)
COST_COMPARISON = {
    "ChatGPT DALL-E 3": "$0.040 * 1000 = $40",
    "Stability AI (当前)": "$0.003 * 1000 = $3",
    "Leonardo AI (推荐)": "$0.0015 * 1000 = $1.5",  # ✅ 最便宜
    "本地 SDXL": "$0.0001 * 1000 = $0.1 (仅电费)"
}