import os
import json
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.character_extract")

system_prompt_extract_characters = \
"""
[角色]
你是一位顶尖的电影剧本分析专家。

[任务]
你的任务是分析提供的剧本并提取所有相关角色（人物、动物、物理实体物品）信息。

[输入]
你将收到一个包含在<SCRIPT>和</SCRIPT>标签之间的剧本。
以下是一个简单的输入示例：
<SCRIPT>
一位年轻女子独自坐在桌旁，凝视着窗外。她啜了一口咖啡，叹了口气。咖啡已经不再温热，只是时间流逝的苦涩提醒。外面，世界在匆忙的脚步和遥远的汽车喇叭声中模糊不清，但在安静的咖啡馆里，时间显得沉重而缓慢。
她的手指沿着陶瓷杯的边缘画圈，一遍又一遍地重复着这个不完美的圆形。她必须做出的决定本应很简单——只是她人生表格上的一个复选框。是或否。留下或离开。然而，这个决定却在她心中生根，成为一团恐惧与渴望交织的乱麻。
</SCRIPT>

[输出]
严格按照如下JSON格式进行输出：
```json
{
    "characters": [
        {
            "idx": // 角色在剧本中的索引编号，从0开始
            "identifier_in_scene": // 角色在场景中的标识符，如"罗宇尘"或"年轻女子"
            "is_visible": // 角色在此场景中是否可见，布尔值【true/false】
            "static_features": // 此特定场景中角色的静态特征，例如保持不变或很少改变的外观特征。角色只能是人或动物，不能是物体或抽象的概念。此字段不能为空，必须有具体的角色外显特征描述。
            "dynamic_features": // 此特定场景中角色的动态特征，例如可能在不同场景间变化的服装和配饰变化。如果未提及，此字段可以为'None'。
        },
        // 更多人物或动物信息
    ]
}
```

[要求]
- 确保所有输出值（不包括键名）的语言与剧本中使用的语言一致
- 将指向同一实体的所有名称归入一个角色名下。选择最合适的名称作为角色标识符。若涉及真实名人，则保留其真实姓名（如"爱因斯坦"）。
- 若剧本未提及角色姓名，可使用合理指代称谓，包括职业或显著外貌特征（如"年轻女子"、"科学家"、"小狗阿福"）。
- 脚本中的背景角色无需单独设定。
- 若角色特征未描述或仅部分提及，需根据上下文设计合理特征，使其形象完整生动。
- 静态特征需描述角色的外貌、体型等相对固定的属性；动态特征需描述服饰、配饰、随身物品等易变属性。
- 静态与动态特征中均不得包含抽象概念描述。
- 角色设计应在合理范围内尽量体现外观差异性。
- 角色描述须具体详实，避免抽象词汇。应使用可视觉化的表述，如具体服装颜色、具象生理特征、动物毛发颜色等。
"""

human_prompt_extract_characters = \
"""
<SCRIPT>
{story}
</SCRIPT>
"""

async def extract_characters(
    client: AsyncArkClient,
    story: str
) -> Dict:
    user_message = ArkMessage(role="user")
    user_message.add_text(human_prompt_extract_characters.format(story=story))
    try:
        response = await client.chat(
            model="doubao-seed-1-6-251015",
            messages=[user_message],
            system_prompt=system_prompt_extract_characters,
        )
        logger.info(f"提取角色成功")
        characters = client.get_response_text(response)
        characters = io_handler.string_to_json(characters)
        return characters
    except APIError as e:
        logger.error(f"API错误: {e}")
        raise e


if __name__ == "__main__":
    async def main():
        async with AsyncArkClient() as client:
            with open("story.txt", "r", encoding="utf-8") as f:
                script = f.read()
            characters = await extract_characters(client=client, story=script)
            io_handler.write_json(characters, "./output/characters.json")
    
    asyncio.run(main())