from openai import OpenAI

client = OpenAI()

def generate_image(prompt: str, size: str = "1024x1024",
                   quality: str = "standard") -> str:
    """生成图片并返回 URL"""
    response = client.images.generate(
        model="dall-e-3",
        prompt=prompt,
        size=size,
        quality=quality,
        n=1,
    )
    return response.data[0].url

# 使用
url = generate_image("一只在东京街头漫步的赛博朋克风格猫，霓虹灯效果")
print(url)

response = client.images.generate(
    model="dall-e-3",
    prompt="A serene Japanese garden with cherry blossoms",
    size="1024x1792",         # 1024x1024 | 1024x1792 | 1792x1024
    quality="hd",             # standard | hd
    style="natural",          # natural | vivid
    n=1,                      # DALL-E 3 只能是 1
    response_format="url",    # url | b64_json
)

import httpx

def save_image(url: str, path: str):
    """下载并保存图片"""
    response = httpx.get(url)
    with open(path, "wb") as f:
        f.write(response.content)
    print(f"已保存: {path}")

# 也可以用 b64_json 格式
import base64

def generate_and_save(prompt: str, path: str):
    response = client.images.generate(
        model="dall-e-3",
        prompt=prompt,
        response_format="b64_json",
        size="1024x1024",
    )
    image_data = base64.b64decode(response.data[0].b64_json)
    with open(path, "wb") as f:
        f.write(image_data)
    print(f"已保存: {path}")

# DALL-E 2 支持单次生成多张
response = client.images.generate(
    model="dall-e-2",
    prompt="a white siamese cat sitting on a windowsill",
    size="1024x1024",       # 256x256 | 512x512 | 1024x1024
    n=3,                     # 最多 10 张
    response_format="url",
)

for i, img in enumerate(response.data):
    print(f"图片 {i+1}: {img.url}")

def edit_image(image_path: str, mask_path: str, prompt: str) -> str:
    """编辑图片的指定区域"""
    response = client.images.edit(
        model="dall-e-2",
        image=open(image_path, "rb"),
        mask=open(mask_path, "rb"),
        prompt=prompt,
        size="1024x1024",
        n=1,
    )
    return response.data[0].url

# 使用：在遮罩区域添加内容
url = edit_image(
    image_path="original.png",
    mask_path="mask.png",       # 白色区域为编辑区域
    prompt="在空白处添加一只蝴蝶",
)

def create_variation(image_path: str, n: int = 3) -> list[str]:
    """生成图片变体"""
    response = client.images.create_variation(
        model="dall-e-2",
        image=open(image_path, "rb"),
        n=n,
        size="1024x1024",
    )
    return [img.url for img in response.data]

import asyncio
from openai import AsyncOpenAI

async_client = AsyncOpenAI()

async def generate_batch(prompts: list[str]) -> list[str]:
    """批量生成图片"""
    async def gen(prompt):
        response = await async_client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size="1024x1024",
            quality="standard",
        )
        return response.data[0].url

    tasks = [gen(p) for p in prompts]
    return await asyncio.gather(*tasks)

# 使用
prompts = [
    "未来城市的天际线",
    "海底珊瑚礁的水下世界",
    "火星上的第一个人类基地",
]

urls = asyncio.run(generate_batch(prompts))
for prompt, url in zip(prompts, urls):
    print(f"{prompt}: {url}")

def generate_from_conversation(user_request: str) -> str:
    """用 GPT 优化 DALL-E prompt"""
    # 第一步：GPT 优化 prompt
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "你是一个DALL-E prompt专家。将用户的简单描述转化为详细的英文图片生成prompt。只输出prompt文本。"},
            {"role": "user", "content": user_request},
        ],
        max_tokens=300,
    )

    optimized_prompt = response.choices[0].message.content
    print(f"优化后的 prompt: {optimized_prompt}")

    # 第二步：生成图片
    image_response = client.images.generate(
        model="dall-e-3",
        prompt=optimized_prompt,
        size="1024x1024",
        quality="hd",
    )

    return image_response.data[0].url

url = generate_from_conversation("一只可爱的机器猫在弹吉他")