import os
import json
from typing import Optional
from utils.tools import string_to_json, save_json_file
from tools.text_generator import media_captioner

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>
{script}
</SCRIPT>
"""

class CharacterExtractor:
    def __init__(self):
        pass

    def extract_characters(self, script: str):
        response = media_captioner.generate_text_understanding(
            system_prompt=system_prompt_extract_characters,
            user_prompt=human_prompt_extract_characters.format(script=script)
        )

        response = string_to_json(response)
        save_json_file(response, "characters.json")
        return response

character_extractor = CharacterExtractor()

if __name__ == "__main__":

    with open("story.txt", "r", encoding="utf-8") as f:
        script = f.read()
    characters = character_extractor.extract_characters(script)
    print(characters)