# -*- coding: utf-8 -*-
import logging
import os
import subprocess

from Manim.Demos.ManimKit import SYS_PROMPT, generate_code_with_llm

# 更详细地控制日志输出
logger = logging.getLogger('T0')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)

if __name__ == '__main__':
    try:
        # 首次生成唯一时间戳
        # timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')[:-3]  # 保留到毫秒
        # 后续写死时间戳
        timestamp = "20250804_084910_884"

        skeleton_basename = f"Skeleton/Skeleton_{timestamp}.py"
        scene_class_name = f"Skeleton_{timestamp}"

        # 教师输入的提示词数组（简化版）
        teacher_prompts = [
            {
                "id": 1,
                "task": "生成片头标题动画",
                "requirements": [
                    "黑色背景",
                    "使用大标题文本'东师理想数学动画AI制作平台作品',字号=60，白色",
                    "添加副标题'正弦与余弦的关系'，字号=36，蓝色",
                    "主标题+副标题淡入2秒，停留3秒，淡出2秒",
                    "最后清屏1秒"
                ]
            },
            {
                "id": 2,
                "task": "生成可运行的 Manim 场景骨架",
                "requirements": [
                    "黑色背景",
                    "建立完整四象限坐标系（x∈[-3.5π,3.5π]，y∈[-1.5,1.5]）",
                    "不画任何曲线，只显示坐标轴,坐标轴要记得带箭头"
                ]
            },
            {
                "id": 3,
                "task": "在上一步基础上追加",
                "requirements": [
                    "用绿色画 y = sin(x)",
                    "用 MathTex 在 (-3π,1) 位置标 'sin(x)'",
                    "动画：坐标轴淡入 1s → 曲线从左到右扫出 3s"
                ]
            },
            {
                "id": 4,
                "task": "在上一步基础上追加",
                "requirements": [
                    "用红色画 y = cos(x)",
                    "用 MathTex 在 (3π,-1) 位置标 'cos(x)'",
                    "动画：TransformFromCopy 把正弦变成余弦 3s",
                    "最后在屏幕上方用黄色 MathTex 写出 cos(x) = sin(x + π/2)"
                ]
            },
            {
                "id": 5,
                "task": "在上一步基础上追加",
                "requirements": [
                    "黄色虚线竖线随 x 移动",
                    "绿色圆点跟踪 sin(x)",
                    "红色圆点跟踪 cos(x)",
                    "ValueTracker 从 -3π 匀速扫到 3π，耗时 6s",
                    "其余元素保持不变"
                ]
            }
        ]

        # 系统补充信息（教师无需关心）
        system_supplements = {
            "base_info": f"文件名：{skeleton_basename}，场景类：{scene_class_name}",
            "output_requirement": "输出：完整代码 + `if __name__ == \"__main__\":`",
            "continue_requirement": f"输出：完整替换 {skeleton_basename} 的新代码",
            "video_config": "帧率：30fps，输出目录：output/",
        }

        # 设置保存路径
        skeleton_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), skeleton_basename)

        # 显示所有可用任务
        print("可用任务列表:")
        for prompt_info in teacher_prompts:
            print(f"任务 {prompt_info['id']}: {prompt_info['task']}")
            for req in prompt_info['requirements']:
                print(f"  - {req}")
            print()

        # 让用户选择要执行的任务ID
        while True:
            try:
                selected_id = int(input("请输入要执行的任务ID (1-5): "))
                if 1 <= selected_id <= 5:
                    break
                else:
                    print("无效的ID，请输入1-5之间的数字。")
            except ValueError:
                print("请输入有效的数字。")

        # 查找用户选择的任务
        selected_prompt = None
        for prompt_info in teacher_prompts:
            if prompt_info['id'] == selected_id:
                selected_prompt = prompt_info
                break

        # 处理用户选择的任务
        prompt_id = selected_prompt["id"]
        task = selected_prompt["task"]
        requirements = selected_prompt["requirements"]
        logger.info(f"处理任务 {prompt_id}: {task}...")

        # 构建教师输入部分
        teacher_part = f"任务：{task}\n要求：\n"
        for req in requirements:
            teacher_part += f"- {req}\n"

        # 构建完整提示词
        if prompt_id == 1:
            # 第一个任务，包含基础信息和输出要求
            full_prompt = f"""
            {teacher_part}
            {system_supplements['base_info']}
            {system_supplements['output_requirement']}            
            重要配置要求:
            - 请设置输出目录为output/
            - 设置帧率为30fps
            - 确保生成MP4格式视频
            - 添加标题淡入淡出动画以确保生成视频而非静态图像
            """
        else:
            # 后续任务，检查前置文件是否存在
            if os.path.exists(skeleton_path):
                # 构建继续要求
                full_prompt = f"""
                {teacher_part}
                {system_supplements['continue_requirement']}                

                {skeleton_basename} 内容：
                {open(skeleton_path, 'r', encoding='utf-8').read()}
                重要配置要求:
                - 请保持输出目录为output/
                - 保持帧率为30fps
                - 确保生成MP4格式视频
                """
            else:
                logger.error(f"错误: 任务 {prompt_id} 依赖于前置任务生成的 {skeleton_basename} 文件，但该文件不存在。")
                logger.info("请先执行任务 1 生成基础文件，然后再执行后续任务。")
                exit(1)

        # 格式化系统提示
        formatted_prompt = SYS_PROMPT.format(XuQiu=full_prompt,Preview=True)

        # 调用函数生成代码
        generate_code_with_llm(prompt=formatted_prompt, save_path=skeleton_path)

        logger.info(f"任务 {prompt_id} 处理完成！生成的文件为：{skeleton_basename}")

        # 执行生成的Skeleton文件
        logger.info(f"正在执行 {skeleton_basename} ...")
        try:
            # 定义清晰的输出目录
            current_dir = os.path.dirname(os.path.abspath(__file__))
            output_dir = os.path.join(current_dir, "output")
            os.makedirs(output_dir, exist_ok=True)

            # 构建执行命令，根据preview参数决定是否添加-p选项
            command = ['python', skeleton_path, '-o', output_dir]

            # 执行文件并通过命令行参数指定输出目录
            subprocess.run(
                command,
                check=True
            )
            logger.info(f"视频已生成到: {output_dir}\\videos\\1080p30\\Skeleton_{timestamp}.mp4")
            # TODO
            # 这个生成的视频路径对于研发来讲很重要，需要返回给前端，前端展示给用户查看效果

        except subprocess.CalledProcessError as e:
            logger.error(f"执行文件时出错: {str(e)}")
        except Exception as e:
            logger.error(f"执行过程中发生未知错误: {str(e)}")

    except Exception as e:
        logger.error(f"程序执行出错: {str(e)}")