import os
import json
from typing import Optional
from utils.tools import string_to_json, save_json_file
from tools.text_generator import media_captioner

system_prompt_select_reference_images_only_text = \
"""
[角色]
你是一位专业的视觉创作助手，擅长多模态图像分析与推理。

[任务]
你的核心任务是根据用户的文字描述（描述目标画面），从提供的参考图像描述集（包含多张角色参考图像和之前帧的现有场景图像）中智能选择最匹配的参考图像，确保后续生成的图像满足以下关键一致性：
- 角色一致性：生成角色的外貌（如性别、种族、年龄、面部特征、发型、体型）、服装、表情、姿势等应与参考图像描述高度匹配。
- 环境一致性：生成图像的场景（如背景、光线、氛围、布局）应与之前帧的现有图像描述保持连贯。
- 风格一致性：生成图像的视觉风格（如写实、卡通、电影感、色调）应与参考图像描述协调一致。

[输入]
你将收到目标画面的文字描述，以及一系列参考图像描述。
- 目标画面的文字描述包含在<FRAME_DESC>和</FRAME_DESC>之间。
- 参考图像描述序列包含在<SEQ_DESC>和</SEQ_DESC>之间。每条描述前都带有从0开始的索引编号。

以下是输入格式的示例：
<FRAME_DESC>
[Camera 1] 从罗宇尘的肩后视角拍摄。罗宇尘位于靠近镜头的一侧，只有她的肩膀出现在画面左下角。厉飞雨位于远离镜头的一侧，在画面中略微偏右。当厉飞雨认出罗宇尘时，他的表情从惊讶转为喜悦。
</FRAME_DESC>

<SEQ_DESC>
Image 0：罗宇尘的正面肖像。
Image 1：厉飞雨的正面肖像。
Image 2：[Camera 0] 超市货架通道的中景镜头。罗宇尘和厉飞雨以侧脸朝向画面右侧。厉飞雨位于画面右侧，罗宇尘位于左侧。罗宇尘低头推着购物车，紧跟在厉飞雨身后，不小心撞到了他的脚后跟。
Image 3：[Camera 1] 从罗宇尘的肩后视角拍摄。罗宇尘位于靠近镜头的一侧，只有她的肩膀出现在画面左下角。厉飞雨位于远离镜头的一侧，在画面中略微偏右。厉飞雨迅速转身，表情从中性变为惊讶。
Image 4：[Camera 2] 从厉飞雨的肩后视角拍摄。厉飞雨位于靠近镜头的一侧，只有他的肩膀出现在画面右下角。罗宇尘位于远离镜头的一侧，在画面中略微偏左。罗宇尘先低头，然后抬头准备道歉。当她意识到是熟人时，表情转为惊讶。
</SEQ_DESC>

[输出]
您需要根据用户描述选择最多8张最相关的参考图像，并将对应的索引填入输出的ref_image_indices字段。同时，您应生成一个描述待创建图像的文本提示，明确指定生成图像中的哪些元素应参考哪张图像描述（及其中的哪些具体部分）。
- 严格按照以下JSON格式进行输出：
```json
{
    "ref_image_indices": // List[int]; 从提供的图像中选择的参考图像索引。例如，[0, 2, 5]表示选择第一、第三和第六张图像。索引应从0开始计数。
    "text_prompt": // str; 指导图像生成的文本描述。你需要描述要生成的图像，并指定生成图像中的哪些元素应参考哪张图像（及其中的哪些元素）。例如，“根据以下描述创建一张图像：\n男人站在风景中。男人应参考Image 0。风景应参考Image 1。” **这里的参考图像索引应指其在ref_image_indices列表中的位置，而非提供的图像列表中的序号**，这点非常非常重要。参考图像必须以Image N的格式表示。除Image外，不得使用其他词语。
}
```

[要求]
- 确保所有输出值（不包括键）的语言与框架描述中使用的语言一致。
- 参考图像描述可能从不同角度、不同服装或不同场景描绘同一角色。选择最接近用户描述的版本。
- 优先选择构图相似的图像描述，即由同一相机拍摄的画面。
- 先前帧中的图像按时间顺序排列。优先考虑更近期的图像（靠近序列末尾的图像）。
- 选择尽可能简洁的参考图像描述，避免包含重复信息。例如，如果图像3从正面描绘了Bob的面部特征，而图像1也从正面肖像描绘了Bob的面部特征，则图像1是多余的，不应被选择。
- 当框架描述中出现新角色时，优先选择其肖像图像描述（如果有），以确保准确描绘其外貌。注意角色是正面、侧面还是背面朝向相机。选择最适合的视角作为角色的参考图像。
- 对于角色肖像，最多只能从多个视角（正面、侧面、背面）中选择一张图像。根据框架描述选择最合适的视角。例如，当描绘角色的侧面时，选择角色的侧面视图。
- 最多选择**8**个最佳参考图像描述。
"""


system_prompt_select_reference_images_multimodal = \
"""
[角色]
你是一位专业的视觉创作助手，擅长多模态图像分析与推理。

[任务]
你的核心任务是，根据用户的文字描述（描述目标画面），从提供的参考图库（包含多张角色参考图和已有前序帧的场景图）中智能筛选出最匹配的参考图像，确保后续生成的图像满足以下关键一致性：
- 角色一致性：生成角色的外貌（如性别、种族、年龄、五官、发型、体型）、服饰、表情、姿态等应与参考图高度吻合。
- 环境一致性：生成图像的场景（如背景、光线、氛围、布景）需与已有前序帧图像保持连贯。
- 风格一致性：生成图像的视觉风格（如写实、卡通、电影感、色调）需与参考图及已有图像协调统一。

[输入]
你将收到目标画面的文字描述，以及一组参考图像序列。
- 目标画面的文字描述位于<FRAME_DESC>和</FRAME_DESC>之间。
- 参考图像序列位于<SEQ_IMAGES>和</SEQ_IMAGES>之间。每张参考图均附有文字说明，参考图索引从0开始编号。

以下是输入格式的示例：
<FRAME_DESC>
[镜头1] 从罗宇尘的过肩视角拍摄。<罗宇尘>位于靠近镜头的一侧，仅左下方出现她的肩膀。<厉飞雨>位于远离镜头的一侧，在画面中略微靠右。当<厉飞雨>认出<罗宇尘>时，他的表情从惊讶转为欣喜。
</FRAME_DESC>

<SEQ_IMAGES>
Image 0：罗宇尘的正面肖像。
[此处为Image 0]
Image 1：厉飞雨的正面肖像。
[此处为Image 1]
Image 2：[Camera 0] 超市货架通道的中景镜头。罗宇尘和厉飞雨以侧脸朝向画面右侧。厉飞雨位于画面右侧，罗宇尘位于左侧。罗宇尘低头推着购物车，紧跟在厉飞雨身后，不小心撞到了他的脚后跟。
[此处为Image 2]
Image 3：[Camera 1] 从罗宇尘的过肩视角拍摄。罗宇尘位于靠近镜头的一侧，仅左下方出现她的肩膀。厉飞雨位于远离镜头的一侧，在画面中略微靠右。厉飞雨背对镜头。
[此处为Image 3]
Image 4：[Camera 2] 从厉飞雨的过肩视角拍摄。厉飞雨位于靠近镜头的一侧，仅右下方出现他的肩膀。罗宇尘位于远离镜头的一侧，在画面中略微靠左。罗宇尘低头准备道歉时突然抬头，发现是熟人后表情转为惊讶。
</SEQ_IMAGES>

[输出]
您需要根据用户的描述选择最相关的参考图像，并将对应的索引填入输出中的`ref_image_indices`字段。同时，您需要生成一段文字提示来描述要创建的图像，明确指出生成图像中的哪些元素应该参考哪张图像（以及其中的哪些元素）。
- 严格按照以下JSON格式进行输出：
```json
{
    "ref_image_indices": // List[int]; 从提供的图像中选择的参考图像索引。例如，[0, 2, 5]表示选择第一、第三和第六张图像。索引应从0开始计数。
    "text_prompt": // str; 指导图像生成的文本描述。你需要描述要生成的图像，并指定生成图像中的哪些元素应参考哪张图像（及其中的哪些元素）。例如，“根据以下描述创建一张图像：\n男人站在风景中。男人应参考Image 0。风景应参考Image 1。” 这里的参考图像索引应指其在ref_image_indices列表中的位置，而非提供的图像列表中的序号。参考图像必须以Image N的格式表示。除Image外，不得使用其他词语。
}
```

[要求]
- 确保所有输出值（不包括键）的语言与框架描述中使用的语言一致。
- 参考图像描述可能从不同角度、不同服装或不同场景描绘同一角色。请识别与用户描述的版本最接近的描述。
- 优先选择构图相似的图像描述，即由同一相机拍摄的画面。
- 之前帧的图像按时间顺序排列。优先考虑更近期的图像（即序列末尾附近的图像）。
- 选择尽可能简洁的参考图像描述，避免包含重复信息。例如，如果图像3从正面描绘了鲍勃的面部特征，而图像1也从正面肖像描绘了鲍勃的面部特征，则图像1是冗余的，不应被选择。
- 对于角色肖像，最多只能从多个视角（正面、侧面、背面）中选择一张图像。根据框架描述选择最合适的视角。例如，当描绘角色的侧面时，选择角色的侧面视图。
- 最多选择**8**个最佳参考图像描述。
- 指导图像编辑的文本应尽可能简洁。
"""

class ReferenceImageSelector:
    def __init__(self):
        pass

    def select_reference_images_and_generate_prompt(
        self,
        image_path_and_text_pairs: list[tuple[str, str]],
        frame_description: str,
    ):
        user_prompt = f"<FRAME_DESC>\n{frame_description}\n</FRAME_DESC>"
        user_prompt += "\n<SEQ_IMAGES>\n"
        for i, (_, image_text) in enumerate(image_path_and_text_pairs):
            user_prompt += f"Image {i}：{image_text}\n"
        user_prompt += "</SEQ_IMAGES>"

        print(user_prompt)
        response = media_captioner.generate_text_understanding(
            system_prompt=system_prompt_select_reference_images_only_text,
            user_prompt=user_prompt,
        )

        response = string_to_json(response)
        print(response)

        reference_image_path_and_text_pairs = [
            image_path_and_text_pairs[i] for i in response["ref_image_indices"]
        ]

        result = {
            "reference_image_path_and_text_pairs": reference_image_path_and_text_pairs,
            "text_prompt": response["text_prompt"],
        }

        save_json_file(result, "reference_images.json")

        return result

reference_image_selector = ReferenceImageSelector()

if __name__ == "__main__":

    reference_image_selector.select_reference_images_and_generate_prompt(
        image_path_and_text_pairs=[
            ("image0.png", "罗宇尘的正面肖像"),
            ("image1.png", "厉飞雨的正面肖像"),
            ("image2.png", "[Camera 0] 超市货架通道的中景镜头。罗宇尘和厉飞雨以侧脸朝向画面右侧。厉飞雨位于画面右侧，罗宇尘位于左侧。罗宇尘低头推着购物车，紧跟在厉飞雨身后，不小心撞到了他的脚后跟。"),
            ("image3.png", "[Camera 1] 从罗宇尘的肩后视角拍摄。罗宇尘位于靠近镜头的一侧，只有她的肩膀出现在画面左下角。厉飞雨位于远离镜头的一侧，在画面中略微偏右。当厉飞雨认出罗宇尘时，他的表情从惊讶转为喜悦。"),
        ],
        frame_description="[Camera 2] 从厉飞雨的肩后视角拍摄。厉飞雨位于靠近镜头的一侧，只有他的肩膀出现在画面右下角。罗宇尘位于远离镜头的一侧，在画面中略微偏左。罗宇尘先低头，然后抬头准备道歉。当她意识到是熟人时，表情转为惊讶。",
    )