"""
Qwen 图像生成示例
使用 DashScope API 进行文生图 (通义千问 Qwen-Image)

模型文档:
https://help.aliyun.com/zh/model-studio/developer-reference/qwen-image

支持的模型:
- qwen-image-plus: 通用文生图模型
- qwen-image-max: 高质量文生图模型

支持的分辨率 (Qwen-Image):
- 1328*1328 (1:1)
- 1664*928 (16:9, 默认)
- 928*1664 (9:16)
- 1472*1104 (4:3)
- 1104*1472 (3:4)

注意: 需要安装依赖
    uv add dashscope python-dotenv requests
"""

import os
from http import HTTPStatus
from pathlib import Path
from urllib.parse import unquote, urlparse

import dashscope
import requests
from dashscope import ImageSynthesis
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 设置 API 基础 URL (北京地域)
dashscope.base_http_api_url = "https://dashscope.aliyuncs.com/api/v1"


def text_to_image_sync(
    prompt: str,
    output_path: str = "output.png",
    size: str = "1328*1328",
    prompt_extend: bool = True,
) -> str:
    """
    同步方式生成图像

    Args:
        prompt: 图像描述文本
        output_path: 输出图像路径
        size: 图像尺寸，支持 1328*1328, 1664*928, 928*1664, 1472*1104, 1104*1472
        prompt_extend: 是否开启 prompt 智能改写 (默认开启，会额外耗时 3-5 秒)

    Returns:
        生成的图像 URL
    """
    api_key = os.getenv("DASHSCOPE_API_KEY")

    print("⏳ 正在生成图像，请稍候...")

    response = ImageSynthesis.call(
        api_key=api_key,
        model="qwen-image-plus",  # 或 qwen-image-max
        prompt=prompt,
        negative_prompt=" ",  # 反向提示词
        n=1,  # 生成图片数量
        size=size,  # 图片尺寸
        prompt_extend=prompt_extend,  # 智能改写 prompt
        watermark=False,  # 不添加水印
    )

    if response.status_code == HTTPStatus.OK:
        result = response.output.results[0]
        image_url = result.url

        print("✅ 图像生成成功!")
        print(f"📷 图像 URL: {image_url}")

        # 如果开启了 prompt_extend，显示改写后的 prompt
        if hasattr(result, "actual_prompt") and result.actual_prompt:
            print(f"📝 改写后的 Prompt: {result.actual_prompt[:100]}...")

        # 下载并保存图片
        img_response = requests.get(image_url)
        with open(output_path, "wb") as f:
            f.write(img_response.content)
        print(f"💾 已保存到: {Path(output_path).absolute()}")

        return image_url
    else:
        print(f"❌ 生成失败: {response.code} - {response.message}")
        return ""


def text_to_image_async(
    prompt: str,
    output_path: str = "output_async.png",
    size: str = "1328*1328",
) -> str:
    """
    异步方式生成图像 (推荐用于生产环境)

    注意: qwen-image-max 模型不支持异步接口

    Args:
        prompt: 图像描述文本
        output_path: 输出图像路径
        size: 图像尺寸

    Returns:
        生成的图像 URL
    """
    import time

    api_key = os.getenv("DASHSCOPE_API_KEY")

    # 提交异步任务
    response = ImageSynthesis.async_call(
        api_key=api_key,
        model="qwen-image-plus",
        prompt=prompt,
        negative_prompt=" ",
        n=1,
        size=size,
        prompt_extend=True,
        watermark=False,
    )

    if response.status_code != HTTPStatus.OK:
        print(f"❌ 提交任务失败: {response.code} - {response.message}")
        return ""

    task_id = response.output.task_id
    print(f"📋 任务已提交, ID: {task_id}")

    # 轮询等待任务完成 (推荐: 前 30 秒每 3 秒查询一次，之后拉长间隔)
    max_wait_time = 120  # 最长等待 2 分钟
    elapsed = 0
    interval = 3

    while elapsed < max_wait_time:
        status = ImageSynthesis.fetch(task_id, api_key=api_key)

        if status.status_code != HTTPStatus.OK:
            print(f"❌ 查询失败: {status.code} - {status.message}")
            return ""

        task_status = status.output.task_status
        print(f"⏳ 任务状态: {task_status}")

        if task_status == "SUCCEEDED":
            result = status.output.results[0]
            image_url = result.url

            print("✅ 图像生成成功!")
            print(f"📷 图像 URL: {image_url}")

            # 下载并保存图片
            img_response = requests.get(image_url)
            with open(output_path, "wb") as f:
                f.write(img_response.content)
            print(f"💾 已保存到: {Path(output_path).absolute()}")

            return image_url

        elif task_status in ("FAILED", "UNKNOWN"):
            print(f"❌ 任务失败: {status.output.code} - {status.output.message}")
            return ""

        time.sleep(interval)
        elapsed += interval

        # 30 秒后拉长轮询间隔
        if elapsed > 30:
            interval = 5

    print("❌ 任务超时")
    return ""


if __name__ == "__main__":
    # 示例 1: 同步方式生成图像
    print("=" * 50)
    print("示例 1: 同步方式生成图像 (Qwen-Image)")
    print("=" * 50)

    prompt_1 = "一只可爱的橘猫坐在咖啡杯旁边，阳光透过窗户洒在桌面上，温馨的下午茶时光"
    text_to_image_sync(prompt_1, "cat_coffee.png")

    print()

    # 示例 2: 异步方式生成图像
    print("=" * 50)
    print("示例 2: 异步方式生成图像 (Qwen-Image)")
    print("=" * 50)

    prompt_2 = "未来科幻城市夜景，霓虹灯闪烁，飞行汽车穿梭在高楼大厦之间，赛博朋克风格"
    text_to_image_async(prompt_2, "cyberpunk_city.png")