import os
import asyncio
from typing import Optional, Dict
from taskflow import FileIOHandler
from api_modules.ark_client_async import AsyncArkClient
from api_modules.ark_client import ArkMessage, APIError
from taskflow import get_logger

io_handler = FileIOHandler()
logger = get_logger("examples.video_create.mcps.refer_image")

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

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

[输入]
你将收到目标画面的文字描述，以及一系列参考图像描述。
- 目标画面的文字描述包含在<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>

[如何输出优秀的text_prompt？]
- 用自然语言清晰描述画面：用简洁连贯的自然语言写明**主体**+**行为**+**环境**+**主体空间位置关系**+**拍摄角度**+**构图**；如果有人物主体，需要补充人物的情绪、动作、表情等细节。
- 明确图片编辑目标和希望保持不变的部分：使用简洁明确的指令，说明需要修改或参考的对象及具体操作，避免使用指代模糊的代词

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

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

async def select_refer_image(
    client: AsyncArkClient,
    frame_description: str,
    image_text_pairs: list[tuple[str, str]]
) -> Dict:
    user_prompt = f"<FRAME_DESC>\n{frame_description}\n</FRAME_DESC>"
    user_prompt += "\n<SEQ_IMAGES>\n"
    for i, (_, image_text) in enumerate(image_text_pairs):
        user_prompt += f"Image {i}：{image_text}\n"
    user_prompt += "</SEQ_IMAGES>"

    user_message = ArkMessage(role="user")
    user_message.add_text(user_prompt)

    try:
        response = await client.chat(
            model="doubao-seed-1-6-251015",
            messages=[user_message],
            system_prompt=prompt_select_refer_image,
        )
        logger.info(f"选择参考图像成功")
        refer_image = client.get_response_text(response)
        refer_image = io_handler.string_to_json(refer_image)
        
        refer_image_path_and_text_pairs = [
            image_text_pairs[i] for i in refer_image["ref_image_indices"]
        ]
        text_prompt = refer_image["text_prompt"]
        result = {
            "reference_image_path_and_text_pairs": refer_image_path_and_text_pairs,
            "text_prompt": text_prompt,
        }
        return result

    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e

if __name__ == "__main__":

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

    async def main():
        async with AsyncArkClient() as client:
            result = await select_refer_image(client=client, frame_description=frame_description, image_text_pairs=image_text_pairs)
            print(result)
    
    asyncio.run(main())