"""
通义千问图片编辑模块

使用通义千问图片编辑模型进行图片编辑
"""

import os
import json
import base64
import mimetypes
import requests
import dashscope
from dashscope import MultiModalConversation

# 配置API地址（中国北京地域）
# 若使用新加坡地域的模型，需将url替换为：https://dashscope-intl.aliyuncs.com/api/v1
dashscope.base_http_api_url = 'https://dashscope.aliyuncs.com/api/v1'

# API Key配置
# 获取API Key：https://help.aliyun.com/zh/model-studio/get-api-key
api_key = "sk-04b63960983445f980d85ff185a17876"

def download_image(image_url, save_path='output.png'):
    """
    下载图片到本地
    
    Args:
        image_url: 图片URL
        save_path: 保存路径
    """
    try:
        response = requests.get(image_url, stream=True, timeout=300)
        response.raise_for_status()
        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"图像已成功下载到: {save_path}")
    except requests.exceptions.RequestException as e:
        print(f"图像下载失败: {e}")


def encode_file(file_path):
    """
    将图片文件编码为base64格式的data URI
    
    Args:
        file_path: 图片文件路径
    
    Returns:
        base64编码的data URI字符串
    
    Raises:
        ValueError: 如果文件格式不支持
        IOError: 如果读取文件失败
    """
    mime_type, _ = mimetypes.guess_type(file_path)
    if not mime_type or not mime_type.startswith("image/"):
        raise ValueError("不支持或无法识别的图像格式")

    try:
        with open(file_path, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
        return f"data:{mime_type};base64,{encoded_string}"
    except IOError as e:
        raise IOError(f"读取文件时出错: {file_path}, 错误: {str(e)}")



def qwen_edit(image_path, prompt, save_path):
    """
    使用通义千问图片编辑模型编辑单张图片
    
    Args:
        image_path: 输入图片路径
        prompt: 编辑提示词
        save_path: 保存路径
    """
    image = encode_file(image_path)

    messages = [
        {
            "role": "user",
            "content": [
                {"image": image},
                {"text": prompt}
            ]
        }
    ]

    # 调用图片编辑模型
    # qwen-image-edit-plus支持输出1-6张图片，此处设置为1张
    response = MultiModalConversation.call(
        api_key=api_key,
        model="qwen-image-edit-plus-2025-10-30",
        messages=messages,
        stream=False,
        n=1,
        watermark=False,
        negative_prompt="不要随意减少纽扣或者随意增加纽扣",
        prompt_extend=False,
        # 仅当输出图像数量n=1时支持设置size参数，否则会报错
        # size="2048*1024",
    )

    if response.status_code == 200:
        for i, content in enumerate(response.output.choices[0].message.content):
            print(f"输出图像{i+1}的URL: {content['image']}")
            download_image(content['image'], save_path)
    else:
        print(f"HTTP返回码：{response.status_code}")
        print(f"错误码：{response.code}")
        print(f"错误信息：{response.message}")
        print("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code")


def qwen_edit_mutil(image_path_list, prompt, save_path):
    """
    使用通义千问图片编辑模型编辑多张图片
    
    Args:
        image_path_list: 输入图片路径列表
        prompt: 编辑提示词
        save_path: 保存路径
    """
    images = [{"image": encode_file(image_path)} for image_path in image_path_list]

    messages = [
        {
            "role": "user",
            "content": images + [{"text": prompt}]
        }
    ]

    # 调用图片编辑模型
    response = MultiModalConversation.call(
        api_key=api_key,
        model="qwen-image-edit-plus-2025-10-30",
        messages=messages,
        stream=False,
        n=1,
        watermark=False,
        negative_prompt="不要随意减少纽扣或者随意增加纽扣",
        prompt_extend=False,
    )

    if response.status_code == 200:
        for i, content in enumerate(response.output.choices[0].message.content):
            print(f"输出图像{i+1}的URL: {content['image']}")
            download_image(content['image'], save_path)
    else:
        print(f"HTTP返回码：{response.status_code}")
        print(f"错误码：{response.code}")
        print(f"错误信息：{response.message}")
        print("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code")


if __name__ == "__main__":
    image_path = [r"H:\data\线稿图\S1261C003.jpg",r"C:\Users\PC\Desktop\企业微信截图_17636312608996.png"]
    prompt = "提取出图1里面整条裤子的平铺精修图，保持原比例，保留衣服的细节,不要随意增加纽扣,图2是图1的细节的补充参考图（图2仅供参考），要严格保留纽扣的数量"
    qwen_edit_mutil(image_path,prompt,save_path=r"H:\data\线稿图\S1261C003_S1261C003_concatenated.jpg")