import os
import asyncio
import json
from typing import Optional, Dict, Any, List
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.camera_tree")


system_prompt_camera_tree = \
"""
[角色]  
你是一位专业的视频剪辑专家，擅长多机位镜头分析与场景结构建模。你深谙影视语言，能够理解景别（如全景、中景、特写）与内容包含关系，并能根据镜头描述推断机位间的层级结构。  

[任务]  
你的任务是分析输入的机位数据，构建"机位树"。该树状结构表示父机位内容包含子机位内容的关系。具体而言，你需要为每个机位识别其父机位（若存在），并确定依赖镜头索引（即父机位素材中包含子机位内容的具体镜头）。若某机位无父机位，则输出None。  

[输入]  
输入为一系列机位数据，序列由<CAMERA_SEQ>和</CAMERA_SEQ>包裹。  
每个机位包含该机位拍摄的镜头序列，由<CAMERA_N>和</CAMERA_N>包裹，其中N为机位索引。  

以下为输入格式示例：  

<CAMERA_SEQ>  
<CAMERA_0>  
shot 0：街道中景。爱丽丝和鲍勃相向而行。  
shot 2：街道中景。爱丽丝和鲍勃相拥。  
</CAMERA_0>  
<CAMERA_1>  
shot 1：爱丽丝面部特写。她认出鲍勃时表情从惊讶转为欣喜。  
</CAMERA_1>  
</CAMERA_SEQ>  

[输出]
严格遵循以下JSON格式输出：
```json
{
    "camera_tree": [
        {
            "parent_cam_idx": // 父机位的索引。如果机位没有父级（例如根机位），则设置为**None**。例如：0、1、None。
            "parent_shot_idx": // 依赖镜头的索引。如果机位没有父级（例如根机位），则设置为**None**。例如：0、3、None。
            "reason": // 选择父机位的原因。如果机位没有父级，应解释为什么它是根机位。例如：父机位的视野涵盖了子机位的视野（从中景到特写）
            "is_parent_fully_covers_child": // 父机位是否完全覆盖子机位的内容。如果机位没有父级，则设置为**None**。例如：True、False、None。
            "missing_info": // 子镜头中父镜头未涵盖的缺失元素。如果父镜头完全覆盖子镜头，则设置为**None**。例如：罗宇尘的正面视角、None。
        },
        // 更多机位的父机位信息
    ]
}
```

[要求]
- 所有输出值（不包括键名）的语言必须与输入语言保持一致。
- 内容包容性检查：父机位应尽可能在特定画面中完全包含子机位内容（例如父中景双人镜头应涵盖子过肩反打镜头）。通过对比关键词（如角色、动作、场景）分析镜头描述，确保父镜头视场能覆盖子镜头。
- 过渡流畅度优先：优先选择更大景别作为父机位，例如全景→中景或中景→特写。相邻父子节点的景别差异应尽可能小，严禁直接从远景跳切到特写（除非绝对必要）。
- 时间邻近性：每个机位由其对应的首个画面描述确定父机位位置，父机位的画面索引应尽可能接近子机位的首个画面索引。
- 逻辑一致性：机位树必须无环，避免循环依赖。若某镜头被多个潜在父机位包含，则选择最佳匹配（基于景别和内容）。若无合适父机位则输出None。
- 当缺乏更广视角时，选择视场重叠最大的镜头作为父镜头（信息重合度最高者），或正反打镜头可互为父子。当两个机位可互为父子时，索引较小者作为索引较大者的父机位。
- 仅允许存在一个无父机位的根机位。
- 描述镜头缺失元素时，需仔细比对父子镜头细节。例如父镜头是角色A与B侧身相对的中景，子镜头是角色A的正脸特写时，需注明子镜头缺失角色A的正面视角信息。
- 首个机位必须作为机位树的根节点。
- **camera_tree**中每个元素代表一个机位的父机位信息；如果机位没有父级（例如根机位），则设置为None。列表的长度应与机位数量相同。
"""

human_prompt_camera_tree = \
"""
<CAMERA_SEQ>
{camera_seq_str}
</CAMERA_SEQ>
"""


async def create_camera_tree(
    client: AsyncArkClient,
    storyboard: Dict,
) -> Dict[str, Any]:

    cameras = []
    for shot in storyboard["storyboard"]:
        if shot["cam_idx"] not in [camera["idx"] for camera in cameras]:
            cameras.append({"idx": shot["cam_idx"], "active_shot_idxs": [shot["idx"]]})
        else:
            cameras[shot["cam_idx"]]["active_shot_idxs"].append(shot["idx"])

    camera_seq_str = ""
    for cam in cameras:
        camera_seq_str += f"<CAMERA_{cam['idx']}>\n"
        for shot_idx in cam["active_shot_idxs"]:
            camera_seq_str += f"Shot {shot_idx}: {storyboard['storyboard'][shot_idx]['visual_desc']}\n"
        camera_seq_str += f"</CAMERA_{cam['idx']}>\n"

    user_prompt = human_prompt_camera_tree.format(camera_seq_str=camera_seq_str)

    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=system_prompt_camera_tree,
        )
        logger.info(f"创建机位树成功")
        camera_tree = client.get_response_text(response)
        camera_tree = io_handler.string_to_json(camera_tree)
        for idx, item in enumerate(camera_tree["camera_tree"]):
            item["active_shot_idxs"] = cameras[idx]["active_shot_idxs"]
        storyboard |= camera_tree
        return storyboard
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e


if __name__ == "__main__":

    json_file = "./output/run_20251215_164241/step4_storyboard.json"
    with open(json_file, "r", encoding="utf-8") as f:
        shot_descriptions = json.load(f)

    shot = shot_descriptions["storyboard"][0]

    async def main():
        async with AsyncArkClient() as client:
            camera_tree = await create_camera_tree(client=client, storyboard=shot)
            print(camera_tree)
    
    asyncio.run(main())