from typing import List  # 用于类型注解，增强代码可读性

import os
import base64
from openai import OpenAI
from PIL import Image
from io import BytesIO

client = OpenAI(api_key="sk-y0T7HiUdqap9xeag1d2dCaD30c1a425681200a3f71A9E1F3", base_url="https://api.laozhang.ai/v1")


##标准迁移流程提示词
DEFAULT_PROMPT = (
    "请将图1中的物体进行渲染。\n"
    "要求：\n"
    "- 严格保持图1中原始模型的结构和几何形状，不得变形、扭曲或重构；\n"
    "- 图像的材质、色彩和纹理应完全匹配图2中的视觉风格；\n"
    "- 输出结果需呈现图2的艺术风格，同时完整保留图1中的所有结构细节。\n"
    "- 仅编辑图1中所提供的掩码区域。\n"
    "- 输出图像的分辨率是1024*1024。"
)


def combine_images_edit(img_path1: str, img_path2: str, img_path3: str = None, output_path: str = "combined.png"):
    """调用 GPT-Image-1 模型合成两张图像，并保存结果到 output_path

    Args:
        img_path1: 初始图像路径
        img_path2: 参考图像路径
        img_path3: 掩码图像路径（可选，如果为None则不使用掩码）
        output_path: 输出图像路径
    """
    try:
        # 准备图像列表
        images = [
            open(img_path1, "rb"),
            open(img_path2, "rb"),
        ]

        # 根据是否有掩码选择不同的API调用方式
        if img_path3 and os.path.exists(img_path3):
            # 使用掩码的编辑模式
            print("🎭 使用掩码进行图像合成...")
            result = client.images.edit(
                model="gpt-image-1",
                image=images,
                mask=open(img_path3, "rb"),
                prompt=DEFAULT_PROMPT,  # 使用带掩码的提示词
                size="1024x1024"
            )
        else:
            # 不使用掩码的编辑模式（仍然使用edit方法，但不传入mask）
            print("🎨 不使用掩码进行图像合成...")
            result = client.images.edit(
                model="gpt-image-1",
                image=images,
                prompt=DEFAULT_PROMPT,  # 使用不带掩码的提示词
                size="1024x1024"
            )

        b64 = result.data[0].b64_json
        if not b64:
            raise RuntimeError("未返回图像数据")
        image_bytes = base64.b64decode(b64)
        with open(output_path, "wb") as f:
            f.write(image_bytes)
        print(f"✅ 合成图像已保存到：{output_path}")
        return True

    except Exception as e:
        print(f"❌ 图像合成失败: {e}")
        return False


def folder_ref_image_style_transfer(initial_img_path: str,
                               mask_img_path: str,
                               reference_dir: str,
                               output_dir: str,
                               use_mask: bool = True,
                               ) -> List[str]:
    """
    批量将参考图像文件夹中的所有图像风格迁移到初始图像的指定区域（基于掩码）

    Args:
        initial_img_path: 初始图像路径（结构保持不变的图像）
        mask_img_path: 掩码图像路径（指定需要修改的区域）
        reference_dir: 参考图像文件夹路径（包含所有待迁移的风格图像）
        output_dir: 输出图像保存文件夹路径
        use_mask: 是否使用掩码进行区域限制（True: 使用掩码，False: 全图迁移）
        default_prompt: 使用掩码时的提示词
        no_mask_prompt: 不使用掩码时的提示词

    Returns:
        成功生成的输出图像路径列表
    """
    # 验证输入文件和文件夹是否存在
    if not os.path.isfile(initial_img_path):
        raise FileNotFoundError(f"初始图像不存在: {initial_img_path}")

    if use_mask and not os.path.isfile(mask_img_path):
        raise FileNotFoundError(f"掩码图像不存在: {mask_img_path}")

    if not os.path.isdir(reference_dir):
        raise NotADirectoryError(f"参考图像文件夹不存在: {reference_dir}")

    # 创建输出文件夹（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 获取参考文件夹中所有图像文件（简单过滤常见图像格式）
    image_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff')
    reference_files = [f for f in os.listdir(reference_dir)
                       if f.lower().endswith(image_extensions)
                       and os.path.isfile(os.path.join(reference_dir, f))]

    if not reference_files:
        print(f"⚠️ 参考文件夹 {reference_dir} 中未找到任何图像文件")
        return []

    # 存储成功生成的输出路径
    successful_outputs = []

    # 遍历所有参考图像并执行风格迁移
    for ref_filename in reference_files:
        ref_img_path = os.path.join(reference_dir, ref_filename)
        print(f"\n📌 开始处理参考图像: {ref_filename}")

        # 生成输出文件名（保留原始参考图像名称并添加后缀）
        base_name, ext = os.path.splitext(ref_filename)
        output_filename = f"transferred_{base_name}{ext}"
        output_path = os.path.join(output_dir, output_filename)

        # 执行单张图像的风格迁移
        try:
            # 使用 with 语句自动管理文件关闭，避免资源泄露
            with open(initial_img_path, "rb") as init_img, \
                    open(ref_img_path, "rb") as ref_img:

                images = [init_img, ref_img]
                prompt = DEFAULT_PROMPT

                # 调用API进行图像编辑
                if use_mask:
                    with open(mask_img_path, "rb") as mask_img:  # 此处正确定义 mask_img
                        print("🔍 使用掩码模式进行区域风格迁移...")
                        result = client.images.edit(
                            model="gpt-image-1",
                            image=images,
                            mask=mask_img,
                            prompt=prompt,
                            size="1024x1024"
                        )
                else:
                    print("🔍 使用全图模式进行风格迁移...")
                    result = client.images.edit(
                        model="gpt-image-1",
                        image=images,
                        prompt=prompt,
                        size="1024x1024"
                    )

            # 处理API返回结果
            b64_data = result.data[0].b64_json
            if not b64_data:
                raise RuntimeError("API返回结果中未包含图像数据")

            # 保存输出图像
            with open(output_path, "wb") as f:
                f.write(base64.b64decode(b64_data))

            print(f"✅ 处理成功，结果保存至: {output_path}")
            successful_outputs.append(output_path)

        except Exception as e:
            print(f"❌ 处理 {ref_filename} 失败: {str(e)}")
            continue

    print(f"\n📊 批量处理完成。成功生成 {len(successful_outputs)}/{len(reference_files)} 张图像")
    return successful_outputs


#生成无缝纹理提示词
DEFAULT_PROMPT_2 = (
    "Generate a seamless, tileable fabric texture inspired by the main material in this image.\n"
    "Do NOT include any clothing shape, object outline, stitching, seams or folds.\n"
    "Only retain the base material — its color, surface pattern, and texture.\n"
    "Output must:\n"
    "1. Be 1024x1024, horizontally and vertically tileable;\n"
    "2. Contain no lighting gradients or shadows;\n"
    "3. Appear like a flat, uniform fabric swatch with consistent roughness."
)

#单张图片生成无缝纹理
def seamless_albedo_generate(input_path,output_dir, mask_path=None, prompt=DEFAULT_PROMPT_2):
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "gpt_image.png")

    try:
        with open(input_path, "rb") as img_file:
            if mask_path:
                with open(mask_path, "rb") as mask_file:
                    result = client.images.edit(
                        model="gpt-image-1",
                        image=[img_file],
                        mask=[mask_file],
                        prompt=prompt,
                        size="1024x1024"
                    )
            else:
                # 无掩码，调用编辑接口但不传掩码参数
                result = client.images.edit(
                    model="gpt-image-1",
                    image=[img_file],
                    prompt=prompt,
                    size="1024x1024"
                )

        b64 = result.data[0].b64_json
        if not b64:
            raise RuntimeError("API 未返回有效图像数据")

        img = Image.open(BytesIO(base64.b64decode(b64)))
        img.save(output_path)

        print(f"✅ 图像已保存: {output_path}")
        return output_path

    except Exception as e:
        print(f"❌ 生成图像失败: {e}")
        return None


"""
输出的文件目录在out_put目录下
"""


def folder_seamless_albedo_generate(input_dir, output_dir, mask_path=None, prompt=DEFAULT_PROMPT_2, ):
    """
    批量处理文件夹中的所有图片，生成无缝纹理

    Args:
        input_dir (str): 输入图片文件夹路径
        mask_path (str, optional): 掩码图片路径，默认为None
        prompt (str, optional): 生成提示词，默认为DEFAULT_PROMPT
        output_dir (str, optional): 输出文件夹路径，默认为"output/gpt_image"

    Returns:
        list: 成功生成的图片路径列表
    """
    # 验证输入文件夹
    if not os.path.isdir(input_dir):
        raise NotADirectoryError(f"输入文件夹不存在: {input_dir}")

    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有图片文件
    image_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff')
    image_files = [
        f for f in os.listdir(input_dir)
        if f.lower().endswith(image_extensions)
           and os.path.isfile(os.path.join(input_dir, f))
    ]

    if not image_files:
        print(f"⚠️ 输入文件夹 {input_dir} 中未找到图片文件")
        return []

    # 处理每张图片
    successful_outputs = []
    for image_name in image_files:
        input_path = os.path.join(input_dir, image_name)
        base_name, _ = os.path.splitext(image_name)
        output_path = os.path.join(output_dir, f"seamless_{base_name}.png")

        print(f"\n🔄 正在处理: {image_name}")

        try:
            with open(input_path, "rb") as img_file:
                if mask_path and os.path.isfile(mask_path):
                    with open(mask_path, "rb") as mask_file:
                        result = client.images.edit(
                            model="gpt-image-1",
                            image=[img_file],
                            mask=[mask_file],
                            prompt=prompt,
                            size="1024x1024"
                        )
                else:
                    # 无掩码，调用编辑接口但不传掩码参数
                    result = client.images.edit(
                        model="gpt-image-1",
                        image=[img_file],
                        prompt=prompt,
                        size="1024x1024"
                    )

            b64 = result.data[0].b64_json
            if not b64:
                raise RuntimeError("API 未返回有效图像数据")

            img = Image.open(BytesIO(base64.b64decode(b64)))
            img.save(output_path)

            print(f"✅ 已生成无缝纹理: {output_path}")
            successful_outputs.append(output_path)

        except Exception as e:
            print(f"❌ 处理失败: {e}")
            continue

    print(f"\n📊 批量处理完成。成功生成 {len(successful_outputs)}/{len(image_files)} 张无缝纹理")
    return successful_outputs


# if __name__ == "__main__":

#     input_image = "./test_files/dress_combine.png"
#     mask_image = None  # 你也可以置 None 测试无掩码情况
#
#     input_dir = r"D:\AppData\PycharmProjects\system-test\data\image\seamless\dress"
#     output_dir = r"D:\AppData\PycharmProjects\system-test\output\seamless_texture\dress"
#     input_mask = r"D:\AppData\PycharmProjects\system-test\data\image\seamless\mask\dress.png"
#     overwrite_images_with_mask(input_dir, input_mask)
#
#     seamless_albedo_generate_batch(input_dir, output_dir)
#     # seamless_albedo_generate(input_image, mask_path=mask_image)
#     # generate_image_from_prompt(input_image, mask_path=None)
#
# ##测试
# if __name__ == "__main__":
#     # 使用你提供的相对路径..跳到上一级目录进行
# #     img1 = "./data/image/bag2.png"  # 初始图像
# #     img2 = "./data/image/bag_r.png"  # 参考图像
# #     img3 = "./data/image/mask_bag.png"  # 掩码图像
# #
# #     # 测试1：使用掩码的图像合成
# #     output1 = "./output/gpt_image/bag_t.png"  ##迁移图像的位置
# #     # print("=== 测试1：使用掩码的图像合成 ===")
# # ## combine_images_edit(img1, img2, img3, output1)
# #
# # # # 测试2：不使用掩码的图像合成
# # # output2 = "../output/gpt_imgae/chair_t2.png"
# # # print("\n=== 测试2：不使用掩码的图像合成 ===")
# # # combine_images_edit(img1, img2, None, output2)
# #
# # ##输入路径
# # INITIAL_IMG = r"D:\AppData\PycharmProjects\system-test\data\image\initial_image\bag2.png"
# # MASK_IMG = r"D:\AppData\PycharmProjects\system-test\data\image\mask\bag2.png"
# # REFERENCE_DIR = r"D:\AppData\PycharmProjects\system-test\data\image\ref\bag2"
# # OUTPUT_DIR = r"D:\AppData\PycharmProjects\system-test\output\gpt_image\bag2"
# # # 重命名
# # rename_files_in_folder(REFERENCE_DIR)
# #
# # os.makedirs(OUTPUT_DIR, exist_ok=True)
# # print("🚀 启动批量图像风格迁移任务（带掩码）")
# # successful = batch_image_style_transfer(
# #     initial_img_path=INITIAL_IMG,
# #     mask_img_path=MASK_IMG,
# #     reference_dir=REFERENCE_DIR,
# #     output_dir=OUTPUT_DIR,
# #     use_mask=True
# # )
# #
# # if successful:
# #     print("\n🎉 成功生成的图像路径：")
# #     for path in successful:
# #         print(f"- {path}")
# # else:
# #     print("\n❌ 未生成任何图像，请检查输入参数和API连接")