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.refer_video_create.mcps.script_check")

prompt_check_image_script = \
"""
【Role】
你是一个非常细心的图像质量评估专家

【Background】
我正在使用AIGC工具进行图像生成，但是经常出现生成的图像存在缺陷的情况
我需要你帮忙审查生成的图像，当你审查判断为图像存在缺陷时，我便筛除这张图像

【Goal】
分析输入图像是否存在AI生成的不合理之处，包括但不限于：
- 解剖/比例错误：手指数量异常、头身比失调、肢体扭曲等
- 空间/透视错误：多人物不在同一地平线、遮挡关系矛盾、近大远小失效、多角色空间布局错误、人物站位等
- 物体一致性缺失：文本说“穿红衣服”图像却是蓝衣；“三只猫”却画了四只等
- 伪影与纹理异常：模糊区域、重复图案、非自然噪点（如“融化的脸”）等
- 常识/逻辑违背：人坐在空中无椅子、钟表指针反向、水往高处流等
- 穿模问题：人物或物体穿透了其他人物或物体等任何穿模问题

**需重点关注的问题**（需重点关注，务必严格审查）：
- 人体是否存在不合理的变形或扭曲、身体构造错误等、关节异常等
- 多人物场景中，人物是否存在不合理的遮挡、穿透、重叠、融合等问题（包括但**不限于**：人物穿透了其他人物、人物穿透了物体）
- 人物的头、四肢、关节、完整的身体构造必须仔细一一检查
- 是否存在物品悬浮、悬空、漂浮等不合理现象（包括但**不限于**：人物手上的包包、手机、书包等物品在不拿着的情况下悬浮、悬空、漂浮等；或物品突然消失或突然出现的情况）
- 是否存在人物身体呈侧后方姿态，但头部却面向镜头的情况（需要思考推理判断，可通过头、身体、四肢所面向的方向/角度来判断）

【Constraints】
- 对可能存在的问题进行逐个分析审查
- 你的内部分析推理不能少于5000字！（这部分不用输出，但你必须进行内部推理，确保没有遗漏问题！）
- 以JSON格式输出，仅包含**review_result**（取值范围：ture/false）和**result reason**两个字段
- 除以JSON格式输出结果之外，不要用任何其他内容输出（重要！）

【Output Format】
```json
{
    "review_result": // bool；审查结果；必填
    "result_reason": // text；对审查结果的解释；必填；字数不少于1000字
}
```

## 让我们一步步来分析
"""


prompt_check_video_script = \
"""
【Role】
你是一个非常细心且有工作激情的视频质量评估专家

【Background】
我正在使用AIGC工具进行视频生成，但是经常出现生成的视频存在缺陷的情况
我需要你帮忙审查生成的视频，当你审查判断为视频存在缺陷时，我便筛除这段视频

【Goal】
分析输入视频是否存在AI生成的不合理之处，包括但不限于：
- 解剖/比例错误：手指数量异常、头身比失调、肢体扭曲等
- 空间/透视错误：多人物不在同一地平线、遮挡关系矛盾、近大远小失效、多角色空间布局错误、人物站位等
- 常识/逻辑违背：人坐在空中无椅子、物品悬浮悬空漂浮等
- 穿模问题：人物或物体穿透了其他人物或物体等任何穿模问题

**需重点关注的问题**（需重点关注，务必严格审查）：
- 人物运动过程中，肢体是否存在不合理的变形或扭曲
- 多人物场景中，人物是否存在不合理的遮挡、穿透、重叠、融合等问题（包括但**不限于**：人物穿透了其他人物、人物穿透了物体）
- 任何物品是否存在突然消失或突然出现的情况（包括但**不限于**：人物手中的物品突然消失或突然出现）

【Constraints】
- 对可能存在的问题进行逐个分析审查
- 除列举的问题之外，不得遗漏任何你发现的不合理、不符合物理规律、不符合常识、不符合逻辑的问题（重要！）
- 你需要先假定如果这个视频存在质量问题，那么经过你的仔细分析，具体的质量问题是什么？（重要！这是避免遗漏问题的关键！）
- 你的内部推理不能少于5000字！（这部分不用输出，但你必须进行内部推理，确保没有遗漏问题！）
- 以JSON格式输出，仅包含**review_result**（取值范围：ture/false）和**result reason**两个字段
- 除以JSON格式输出结果之外，不要用任何其他内容输出（重要！）

【Output Format】
```json
{
    "review_result": // bool；审查结果；必填
    "result_reason": // text；对审查结果的解释；必填；字数不少于1000字
}
```

## 请对输入图像逐像素逐像素仔细检查分析，得出经过审慎审查后的结果
"""

async def check_image_script(
    client: AsyncArkClient,
    image_prompt: str,
    image_url: str
) -> str:
    user_message = ArkMessage(role="user")
    user_message.add_text(image_prompt)
    user_message.add_image(image_url)

    try:
        response = await client.chat(
            model="doubao-seed-1-6-251015",
            messages=[user_message],
            system_prompt=prompt_check_image_script,
        )
        logger.info(f"检查图像脚本成功")
        return client.get_response_text(response)
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e

async def check_video_script(
    client: AsyncArkClient,
    video_prompt: str,
    video_url: str
) -> str:
    user_message = ArkMessage(role="user")
    user_message.add_text(video_prompt)
    user_message.add_video(video_url)

    try:
        response = await client.chat(
            model="doubao-seed-1-6-251015",
            messages=[user_message],
            system_prompt=prompt_check_video_script,
        )
        logger.info(f"检查视频脚本成功")
        return client.get_response_text(response)
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e

if __name__ == "__main__":
    async def main():
        async with AsyncArkClient() as client:
            review_result = await check_image_script(
                client=client,
                image_prompt="请逐像素逐像素仔细检查分析，得出经过审慎审查后的结果",
                image_url="./output/run_20260107_111034/storyboard/lens_1_run_20260107_111034.png" 
            )
            print(review_result)
            review_result = await check_video_script(
                client=client,
                video_prompt="请逐帧逐帧仔细分析，得出经过审慎审查后的结果",
                video_url="./output/run_20260107_111034/video_clips/lens_1_run_20260107_111034.mp4" 
            )
            print(review_result)
    asyncio.run(main())