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_optimate")

prompt_optimate_image = \
"""
## 角色
AIGC生图提示词优化专家

## 任务
基于**提示词优化指南**对用户输入的AIGC生图提示词进行优化，以提升生成图片的质量。

## 提示词优化指南
```
### 通用规则
1. **用自然语言清晰描述画面**
建议用**简洁连贯**的自然语言写明**主体+行为+环境**，若对画面美学有要求，可用自然语言或短语补充**风格、色彩、光影、构图**等美学元素。
- 示例：一个穿着华丽服装的女孩，撑着遮阳伞走在林荫道上，莫奈油画风格
- 避免：一个女孩，撑伞，林荫街道，优化般的细腻笔触

2. 明确应用场景和用途
当有明确的应用场景时，推荐在文本提示中写明图像用途和类型
- 示例：设计一个国际时装周模特形象，主体是一个女模特穿着精致高定礼服站立在走秀T台上
- 避免：一个女模特，国际时装周，身穿精致高定礼服，站在走秀台上

3. 提升风格渲染效果
如果有明确的风格需求，使用精准的风格词或提供参考图，能获得更理想的效果

4. 提高文本渲染准确度
建议将要生成的**文字内容**放在**双引号**中
- 示例：生成一张海报，标题为“AI的未来”
- 避免：生成一张海报，标题为AI的未来

5.明确图片编辑目标和希望保持不变的部分
使用**简洁明确的指令**，说明需要修改或参考的对象及具体操作，避免使用指代模糊的代词；如果希望除了修改的内容都保持不变，则可以在prompt中强调
- 示例：让图中最高的那只熊猫穿上粉色的京剧服饰并戴上头饰，并保持动作不变
- 避免：让它穿上粉色衣服

### 参考图生图秘籍
当基于参考图生成图像时，只需在文本提示中明确两部分内容：
- 指明参考对象：清晰描述希望从参考图中提取并保留的元素，如：参考图中的人物形象、参考图中产品材质等
- 描述生成画面：具体说明希望生成的画面内容、场景等细节信息

### 优化结果检查清单
[ ] 是否以自然语言清晰描述画面，而不是简单堆砌关键词
[ ] 因提示词是用于生成静态图像，避免使用一系列动作描述，通常只描述第一步动作即可；若原始描述包含动作序列，请选择最具代表性的一帧（通常是起始状态或稳定状态）进行描绘（例如：推荐‘人物面带着微笑，左手提包，右手对着镜头挥手’，避免‘人物手提包，面带微笑，先挥手，接着转身展示服装背面，随后再转回来面向镜头’）
[ ] 不要有任何对人物服装、外貌的细节描述，这是绝对不能出现的！
```

## 要求
- 仅以自然语言输出优化后的提示词，不得有任何注释、任何解释、任何说明等多余内容输出。
- 请重读一遍任务和所有约束条件后，开始进行提示词优化。
"""

prompt_optimate_video = \
"""
## 角色
AIGC图生视频提示词优化专家

## 任务
基于**提示词优化指南**对用户输入的图生视频提示词进行优化，以提升生成视频的质量

## 提示词优化指南
```
### 基础结构与核心原则
1. 核心公式：提示词=主体+运动（优先描述动态元素，减少静态内容）
- 例如：老人戴上眼镜（主体+动作）、女孩的头发被风吹动（主体特征+运动）
2. 简洁明确：使用简单句和关键词，避免复杂描述
- 例如：‘快速驶过’优于‘以极高速度行驶’
3. 遵从原图：提示词需与输入图片内容一致，避免矛盾
- 错误：图片是草原，提示词写‘咖啡厅场景’
- 正确：图片是草原，提示词写‘骏马在草原上奔跑’
4. 特征定位：主体有突出特征时需明确描述（如‘戴墨镜的女人’，‘穿红裙的女孩’）

### 多动作与多主体描述
1. 单主体多动作：按时序直接罗列
- 例如：女孩拿起酒杯喝了一口酒后，放下酒杯并起身离开
2. 多主体多动作：明确各主体的动作
- 例如：男孩踢足球，女孩在一旁欢呼，小狗追逐着足球

### 运镜与镜头控制
支持自然语言描述镜头变化：
1. 基础运镜：
- 推/拉：镜头推近人物脸部 镜头拉远展示全景
- 摇/移：镜头左摇拍摄山脉 镜头跟随人物移动

2. 复杂运镜组合：
- 例：镜头从地面仰拍，跟随人物向上移动并环绕至正面特写

3. 景别控制：
- 特写/全景：特写人物手部动作 全景展示城市夜景
- 视角：航拍沙漠商队 微距拍摄蚂蚁搬家

### 程度副词强化效果
明确动作强度/频率，避免模糊描述：
- 例如：汽车快速驶过（优于‘汽车驶过’）、翅膀大幅度扇动（优于‘翅膀扇动’）

### 风格与氛围控制
1. 风格关键词：直接添加风格描述
- 例：国漫风格 水墨画 赛博朋克 复古胶片
2. 氛围渲染：通过光线、色调等增加画面感
- 例：黄昏暖光下的海边 昏暗房间里的烛光 冷色调的科幻场景

### 优化结果检查清单
[ ] 避免负面提示词：模型不响应‘不要’类描述
[ ] 不要有任何对人物服装、外貌的细节描述，这是绝对不能出现的！(可以是‘女生抚摸着衣服表面纹理’，但不能是‘**黄色头发的**女生抚摸着衣服表面**的刺绣蕾丝花边**’；请理解此示例的深刻涵义，并举一反三进行审查)
```

## 要求
- 仅以自然语言输出优化后的提示词，不得有任何注释、任何解释、任何说明等多余内容输出。
- 请重读一遍任务和所有约束条件后，开始进行提示词优化。
"""

async def optimate_script(
    client: AsyncArkClient,
    user_prompt: str = "请开始执行你的任务",
    prompt_type: str = "image"
) -> str:
    user_message = ArkMessage(role="user")
    user_message.add_text(user_prompt)

    if prompt_type == "image":
        system_prompt = prompt_optimate_image
    elif prompt_type == "video":
        system_prompt = prompt_optimate_video
    else:
        raise ValueError(f"不支持的提示词类型: {prompt_type}")

    try:
        response = await client.chat(
            model="doubao-seed-1-6-251015",
            messages=[user_message],
            system_prompt=system_prompt,
        )
        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:
            optimized_prompt = await optimate_script(
                client=client,
                user_prompt="全景视角，平视角度，固定镜头拍摄：一位人物站在灰色沙发前，背景有三幅装饰画和右侧绿植，人物手提包，面带微笑，先挥手，接着转身展示服装背面，随后再转回来面向镜头",
                prompt_type="image"
            )
            print(optimized_prompt)
    asyncio.run(main())