import os
import json
import asyncio
from typing import Optional
from itertools import product
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.storyboard_create")


system_prompt_design_storyboard = \
"""
[角色]
你是一位专业的分镜脚本艺术家，具备以下核心技能：
- 剧本分析：能够快速解读剧本文字，准确识别场景设定、角色动作、对白、情绪以及叙事节奏。
- 视觉化能力：擅长将文字描述转化为视觉画面，包括构图、光影和空间布局。
- 分镜绘制：精通电影语言，如镜头类型（特写、中景、全景等）、摄影角度（俯拍、平视等）、摄像机运动（推拉、摇移等）以及镜头转场方式。
- 叙事连贯性：能够确保分镜序列逻辑流畅，突出关键情节，并保持情感表达的一致性。
- 技术知识：熟悉基本的分镜格式和行业标准，例如使用编号镜头和简洁明了的画面说明。

[任务]
你的任务是根据用户提供的剧本（仅包含一个场景）设计一套完整的分镜脚本。分镜脚本应以文字形式呈现，清晰展示每个镜头的视觉元素和叙事流程，帮助用户直观地想象该场景。

[输入]
用户将提供以下输入内容：
- 剧本：一段完整的单场剧本，包含对白、动作描述和场景设定。该剧本仅聚焦于一个场景，无需处理多场景之间的转场。剧本内容位于<SCRIPT>和</SCRIPT>标签之间。
- 角色列表：列出每位角色的基本信息，例如姓名、性格特征、外貌特征（如相关）。角色列表位于<CHARACTERS>和</CHARACTERS>标签之间。
- 用户要求（可选）：位于 <USER_REQUIREMENT> 和 </USER_REQUIREMENT> 标签之间，可能包括：
    - 目标受众（例如儿童、青少年、成人）；
    - 分镜风格（例如写实、卡通、抽象）；
    - 期望的镜头数量（例如"不超过10个镜头"）；
    - 其他具体指示（例如强调角色的动作）。

[输出]
严格以JSON格式输出（必须严格遵守JSON格式，不得有任何其他内容），每个元素代表本场景下完整分镜脚本中的一个分镜脚本，例如：
```json
{
    "storyboard":[
        {
            "idx": // 镜头编号，从0开始
            "is_last": // 是否是镜头序列中的最后一个镜头；如果是，则故事结束，不会再有更多镜头序列。布尔值
            "cam_idx": // 本镜头所属机位索引，从0开始
            "visual_desc": // 对镜头的生动详细视觉描述，通过文字传达丰富的视觉信息。描述中的角色标识符必须与角色列表中的匹配，并用尖括号括起来（例如<罗宇尘>、<林婉>）。应描述所有可见的角色。如果有对话，请写下对话内容），当遇到一些对话时，您应该用：''符号写入视觉内容描述和角色特征（例如<罗宇尘>（男性，20多岁，德州口音因军事精确性而软化，自信且充满活力。）说：'起落架已收起。襟翼正在转换。飞行路径稳定。您可以爬升了。'）。
            "audio_desc": // 镜头中音频的详细描述。例如：[音效] 环境声音（超时背景噪音，购物车轮子滚动声）、[说话者] 罗宇尘（开心）：'你好呀！'、None（表示没有任何声音）
        },
        // 更多分镜脚本序列
    ]
}
```

[要求]
- 确保所有输出内容（除键名外）所使用的语言与剧本中的语言一致。
- 每个镜头必须具有明确的叙事目的——例如建立场景环境、展现角色关系或突出角色反应。
- 有意识地运用电影语言：用特写表现情绪，用全景交代环境，通过多样的角度引导观众注意力。
- 在设计新镜头时，首先考虑是否可以沿用已有的摄像机位置；仅当镜头景别、角度或焦点发生显著变化时，才引入新的摄像机位置。若摄像机进行了大幅度运动，则该位置此后不可再使用。
- 在视觉描述和对白发言者字段中，角色名称必须与角色列表保持一致。在视觉描述中，角色名需用尖括号括起（例如 <罗宇尘>），但在对白或发言者字段中则不加括号。
- 描述视觉元素时，必须指明其在画面中的具体位置。例如：“角色A位于画面左侧，面朝右方，面前有一张桌子；该桌子位于画面中心偏左的位置。” 不得包含画面中不可见的元素，例如：若门紧闭，则不可描述门后的人物。
- 视觉描述中避免出现不安全内容（如暴力、歧视等）。必要时可采用间接手法，如通过声音或暗示性画面表现，并对敏感元素进行替代处理（例如用番茄酱代替血液）。
- 每个镜头中，每位角色最多分配一句对白。每句对白应对应一个独立镜头。
- 每个镜头的描述必须独立完整，不得引用其他镜头内容。
- 当镜头聚焦于某角色时，需说明具体聚焦的身体部位（如面部、手部等）。
- 描述角色时，必须标明其朝向（例如“面朝左”、“背对镜头”等）。
"""

human_prompt_design_storyboard = \
"""
<SCRIPT>
{script_str}
</SCRIPT>

<CHARACTERS>
{characters_str}
</CHARACTERS>

<USER_REQUIREMENT>
{user_requirement_str}
</USER_REQUIREMENT>
"""


system_prompt_decompose_visual_description = \
"""
[角色]
你是一位专业的视觉文本分析师，精通电影语言与镜头叙事。你的专长在于将一段完整的镜头描述精准地拆解为三个核心组成部分：静态的起始画面、静态的结束画面，以及连接两者之间的动态运动过程。

[任务]
你的任务是严格且深入地将用户提供的镜头视觉文本描述拆解并重写为以下三个独立部分：
- 起始画面描述：描述镜头最开始时的静态画面。聚焦于构图元素、角色初始姿态、环境布局、光影、色彩及其他静态视觉特征。
- 运动过程描述：描述从起始画面到结束画面之间发生的所有动态变化。包括摄像机运动（例如：固定、推进、拉远、横摇、跟拍、俯仰等）以及画面内元素的运动（例如：角色移动、物体位移、光影变化等）。这是整个描述中最具动态性的部分。在描述角色的运动和变化时，不得直接使用角色姓名，而应通过其外部特征（尤其是显著的衣着特征等）来指代该角色。
- 结束画面描述：描述镜头结束时的静态画面。同样关注静态构图，但必须体现因摄像机运动或画面内元素移动、角色动作所导致的最终状态。

[输入]
你将收到一段镜头的视觉文本描述，其中通常隐含或明确包含起始状态、运动过程和结束状态的信息。
此外，你还将收到一份潜在角色列表，每个角色包含一个标识符及其显著特征。
- 视觉描述位于 <VISUAL_DESC> 与 </VISUAL_DESC> 标签之间。
- 角色列表位于 <CHARACTERS> 与 </CHARACTERS> 标签之间。

[输出]
严格按照以下JSON格式进行输出（必须严格遵守JSON格式，不得有任何其他内容）：
```json
{
    "ff_desc": // 镜头第一帧的详细描述，捕捉初始的视觉元素和构图。
    "ff_vis_char_idxs": // 镜头第一帧中可见角色的索引列表，对应于输入中提供的角色列表。例如：[0]、[0,1]
    "lf_desc": // 对镜头最后一帧的详细描述，捕捉其最终的视觉元素与构图。
    "lf_vis_char_idxs": // 镜头最后一帧中可见角色的索引列表，对应于输入中提供的角色列表。例如：[0]、[0,1]
    "motion_desc": // 镜头的运动描述，描述镜头内的动态变化（包括摄像机运动和画面内元素的移动）。例如：从半身镜头推近至特写。罗宇尘（留着胡子，穿着白色T恤）对着镜头微笑。
    "variation_type": // 表示第一帧和最后一帧之间的变化程度。可选值有"large"、"medium"和"small"。
    "variation_reason": // 给出以上变化程度的合理解释。例如：与第一帧相比，最后一帧中出现了一个新角色，而构图没有发生显著变化。因此变化程度中等。
}
```

[指导原则]
- 确保所有输出内容（除键名外）所使用的语言与剧本中的语言一致。
- 起始画面和结束画面的描述必须是纯粹的“静态快照”，不得包含正在进行的动作（例如：“他正要站起来”是不可接受的；应表述为：“他坐在椅子上，身体略微前倾”）。
- 在运动过程描述中，必须清晰区分摄像机运动与画面内元素的运动。尽可能准确地使用专业电影术语（如：推轨镜头、横摇、变焦等）来描述摄像机运动，使用可视化语言（如：女生举起右手、男生向右转过身去、书本从书桌上掉下来等）描述画面内元素的运动。
- 在运动过程描述中，不得直接使用角色姓名指代角色；而应通过角色可见的外部特征进行指代。例如，“罗宇尘正在行走”不可接受，应表述为“一位短发、身穿绿色连衣裙的青年男生正在行走”。
- 结束画面因初始画面中的运动过程而产生，因此结束画面描述必须在逻辑上与起始画面描述及运动过程描述保持一致。运动过程中描述的所有动作、姿态和位置变化都应在结束画面的静态图像中有所体现。
- 若输入描述对某些细节表述模糊，可根据上下文做出合理推断和补充，以确保三个部分完整流畅；但核心要素必须严格遵循输入文本。
- 使用准确、简洁且专业的描述性语言。避免使用过于文学化的修辞（如隐喻或情绪化修饰），聚焦于可被视觉化呈现的信息。
- 与输入的视觉描述类似，起始画面和结束画面的描述应包含镜头类型、拍摄角度、构图等细节。
- 镜头内部的变化可分为以下三类（注意：这是单个镜头内的变化，而非镜头之间的切换）：
(1) “large”变化：通常指夸张的过渡镜头，即构图与焦点发生显著改变，例如从全景平滑过渡到特写。此类变化通常伴随显著的摄像机运动（如穿越城市上空的无人机视角）。
(2) “medium”变化：常涉及新角色的引入，或已有角色从背对镜头转为正面朝向镜头。
(3) “small”变化：通常指细微变化，例如角色表情变化、已有角色的姿态或位置变化（如行走、坐下、站起），以及适度的摄像机运动（如横摇、俯仰、跟拍）。
- 描述角色时，必须标明其朝向（例如“面朝左”、“背对镜头”等）。
- 第一个镜头必须以尽可能宽广的景别建立整体场景环境。
- 尽可能减少摄像机位置的数量。
"""

human_prompt_decompose_visual_description = \
"""
<VISUAL_DESC>
{visual_desc}
</VISUAL_DESC>

<CHARACTERS>
{characters_str}
</CHARACTERS>
"""

async def design_storyboard(
    client: AsyncArkClient,
    script: str,
    characters: str,
    user_requirement: Optional[str] = None
):
    user_prompt = human_prompt_design_storyboard.format(
        script_str=script,
        characters_str=characters,
        user_requirement_str=user_requirement
    )

    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_design_storyboard,
        )
        logger.info(f"设计分镜成功")
        storyboard = client.get_response_text(response)
        storyboard = io_handler.string_to_json(storyboard)
        return storyboard
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e

async def decompose_visual_description(
    client: AsyncArkClient,
    visual_desc: str,
    characters: str
):
    user_prompt = human_prompt_decompose_visual_description.format(
        visual_desc=visual_desc,
        characters_str=characters
    )
    
    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_decompose_visual_description,
        )
        logger.info(f"拆解视觉描述成功")
        decomposed_visual_desc = client.get_response_text(response)
        decomposed_visual_desc = io_handler.string_to_json(decomposed_visual_desc)
        return decomposed_visual_desc
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e

async def create_storyboard_for_single_scene(
    client: AsyncArkClient,
    script: str,
    characters: str,
    user_requirement: Optional[str] = None,
    global_semaphore: Optional[asyncio.Semaphore] = None,
    scene_idx: Optional[int] = None
):
    """
    处理单个场景的分镜设计（优化版本：使用全局信号量控制并发）
    
    Args:
        client: AsyncArkClient 实例
        script: 场景剧本
        characters: 角色信息
        user_requirement: 用户要求
        global_semaphore: 全局信号量，用于控制所有API请求的总并发数
        scene_idx: 场景索引（用于日志）
    """
    # 判断script是否为str类型，如果不是，则先转换为str类型
    if not isinstance(script, str):
        script = str(script)
    if not isinstance(characters, str):
        characters = str(characters)
    if not isinstance(user_requirement, str):
        user_requirement = str(user_requirement)

    scene_log_prefix = f"场景 {scene_idx}" if scene_idx is not None else "场景"

    try:
        # 使用全局信号量控制并发
        if global_semaphore:
            async with global_semaphore:
                storyboard = await design_storyboard(
                    client=client,
                    script=script,
                    characters=characters,
                    user_requirement=user_requirement
                )
        else:
            storyboard = await design_storyboard(
                client=client,
                script=script,
                characters=characters,
                user_requirement=user_requirement
            )
        logger.info(f"{scene_log_prefix} 设计分镜成功")
    except Exception as e:
        logger.error(f"{scene_log_prefix} 设计分镜失败: {e}")
        raise e

    # 拆解每个分镜的视觉描述（异步并行执行）
    total_items = len(storyboard["storyboard"])
    logger.info(f"{scene_log_prefix} 开始异步并行处理 {total_items} 个分镜项的视觉描述...")

    async def process_item(item):
        """处理单个分镜项的异步函数"""
        item_idx = item.get("idx", "unknown")
        logger.info(f"{scene_log_prefix} 处理分镜项 {item_idx}...")
        
        # 使用全局信号量控制并发
        if global_semaphore:
            async with global_semaphore:
                decomposed_visual_desc = await decompose_visual_description(
                    client=client,
                    visual_desc=item["visual_desc"],
                    characters=characters
                )
        else:
            decomposed_visual_desc = await decompose_visual_description(
                client=client,
                visual_desc=item["visual_desc"],
                characters=characters
            )
        
        logger.info(f"{scene_log_prefix} 分镜项 {item_idx} 处理完成")
        return item, decomposed_visual_desc

    # 创建所有任务并立即并发执行（不等待其他场景）
    tasks = [process_item(item) for item in storyboard["storyboard"]]
    
    completed_count = 0
    try:
        # 并发执行所有分镜项的处理（使用全局信号量控制总并发数）
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果和异常
        for result in results:
            if isinstance(result, Exception):
                logger.error(f"{scene_log_prefix} 处理分镜项时发生错误: {result}")
                raise result
            
            item, decomposed_visual_desc = result
            item |= decomposed_visual_desc
            completed_count += 1
        
        logger.info(f"{scene_log_prefix} 进度: {completed_count}/{total_items} 个分镜项已完成")
    except Exception as e:
        logger.error(f"{scene_log_prefix} 处理分镜项时发生错误: {e}")
        raise

    logger.info(f"{scene_log_prefix} 所有 {total_items} 个分镜项的视觉描述处理完成")
    return storyboard


async def create_storyboard(
    script: dict,
    characters: str,
    user_requirement: Optional[str] = None,
    max_concurrent_requests: int = 50
):
    """
    创建分镜脚本（优化版本：完全异步并发处理）
    
    优化策略：
    1. 外层循环：并发处理所有场景的分镜设计
    2. 内层循环：每个场景内部并发处理所有分镜项的视觉描述拆解
    3. 使用全局信号量控制所有API请求的总并发数，防止服务器限流
    
    Args:
        script: 剧本字典，包含多个场景
        characters: 角色信息字符串
        user_requirement: 用户要求
        max_concurrent_requests: 最大并发请求数（默认50，可根据服务器调整）
    
    Returns:
        包含所有场景分镜脚本的字典
    """
    logger.info(f"开始设计分镜（共 {len(script['script'])} 个场景）...")
    
    # 创建全局信号量，控制所有API请求的总并发数
    # 这包括：场景级别的设计分镜请求 + 所有分镜项的视觉描述拆解请求
    global_semaphore = asyncio.Semaphore(max_concurrent_requests)
    
    async def handle_scene(scene_idx: int, scene_script: str):
        """处理单个场景的异步函数"""
        logger.info(f"开始处理场景 {scene_idx}...")
        try:
            scene_storyboard = await create_storyboard_for_single_scene(
                client=client,
                script=scene_script,
                characters=characters,
                user_requirement=user_requirement,
                global_semaphore=global_semaphore,
                scene_idx=scene_idx
            )
            logger.info(f"场景 {scene_idx} 的分镜设计完成")
            return scene_idx, scene_storyboard
        except Exception as e:
            logger.error(f"场景 {scene_idx} 处理失败: {e}")
            raise
    
    async with AsyncArkClient() as client:
        # 并发处理所有场景（外层并发）
        scene_tasks = [
            handle_scene(idx, scene_script) 
            for idx, scene_script in enumerate(script["script"])
        ]
        
        logger.info(f"🚀 并发启动 {len(scene_tasks)} 个场景的处理任务（最大并发请求数: {max_concurrent_requests}）...")
        
        # 并发执行所有场景，并保持顺序
        scene_results = await asyncio.gather(*scene_tasks, return_exceptions=True)
        
        # 检查是否有异常
        failed_scenes = []
        for idx, result in enumerate(scene_results):
            if isinstance(result, Exception):
                logger.error(f"❌ 场景 {idx} 处理失败: {result}")
                failed_scenes.append(idx)
        
        if failed_scenes:
            raise Exception(f"以下场景处理失败: {failed_scenes}")
        
        # 按场景索引排序，确保结果顺序正确
        scene_results.sort(key=lambda x: x[0])
        scene_storyboard_list = [result[1] for result in scene_results]
        
        logger.info(f"✅ 所有 {len(script['script'])} 个场景的分镜设计完成")
        logger.info(f"✅ 分镜设计完成")
        
        return {"storyboard": scene_storyboard_list}

if __name__ == "__main__":  
    async def main():
        storyboard = await create_storyboard(
            script=script,
            characters=characters,
            user_requirement=user_requirement
        )
        io_handler.write_json(storyboard, "./output/storyboard.json")
    
    asyncio.run(main())