import os
import re
from openai import OpenAI

# 加载LLM配置的函数
class LLMConfig:
    def __init__(self, api_key, model_name, base_url):
        self.api_key = api_key
        self.model_name = model_name
        self.base_url = base_url

def load_llm_config():
    """
    加载LLM配置
    
    返回:
        LLMConfig: 包含API密钥、模型名称和基础URL的配置对象
    """
    try:
        from Config.Config import ALY_LLM_API_KEY, ALY_LLM_MODEL_NAME, ALY_LLM_BASE_URL
        return LLMConfig(
            api_key=ALY_LLM_API_KEY,
            model_name=ALY_LLM_MODEL_NAME,
            base_url=ALY_LLM_BASE_URL
        )
    except Exception as e:
        print(f"加载配置时出错: {str(e)}")
        return None

SYS_PROMPT = """ 
        你是一位资深 Manim 动画工程师，专注于教育场景。请遵循以下规范：
        1. 仅输出 Python 代码，不解释思路。
        2. 使用 Manim 社区版 v0.18.0 语法。
        3. 严格按提示的步骤进行编码，不要发挥、扩展。
        4. 所有公式用 MathTex，文字用 Text（更好地支持中文），颜色统一用 Manim 内置调色板。
        5. 关键坐标用常量定义在类顶部，例如 LEFT_MARGIN = 3.2。
        6. 每行不超过 88 字符，函数名用 snake_case。
        7. 动画时长 3-8 秒 / 步，默认 ease-in-out。
        8. 请严格按照以下模板：坐标轴范围、颜色、字体统一使用 Manim 社区版默认主题。
        9. 代码的最后，参考下面的代码添加相应内容:
        if __name__ == "__main__":
            # 导入必要的模块
            import sys
            import os
            from manim import *
            from manim.utils.rate_functions import ease_in_out_sine

            # 设置 UTF-8 编码
            os.environ['PYTHONIOENCODING'] = 'utf-8'
            sys.stdout.reconfigure(encoding='utf-8')
            sys.stderr.reconfigure(encoding='utf-8')

            # 全局指定 ctex 模板（一次性）
            config.tex_template = TexTemplateLibrary.ctex
            # 设置输出格式为MP4
            config.output_format = "mp4"
            # 设置输出目录
            config.media_dir = "./output/"
            # 使用临时配置渲染场景（配置只在with块内有效）
            with tempconfig({{
                "quality": "high_quality",
                "background_color": BLACK,
                "preview": {Preview},
                "pixel_height": 1080,
                "pixel_width": 1920,
                "frame_rate": 30
            }}):
                # 实例化场景类
                scene = Skeleton()
                # 执行渲染流程（包含文件生成和预览）                
                scene.render()
        10. 【输出要求】
            (1). 完整可运行 Python 脚本
            (2). 包含 `if __name__ == "__main__":` 可直接运行。  
            (3). 代码内只出现中文注释，方便二次修改。  
        本任务的具体需求如下：
        {XuQiu}
        """

def generate_code_with_llm(prompt, save_path, config=None):
    """
    调用大模型生成代码并保存
    
    参数:
        prompt: 提示词
        save_path: 代码保存路径
        config: 可选的配置对象，如果未提供则自动加载
    
    返回:
        bool: 是否成功保存代码
    """
    try:
        # 如果未提供配置，则加载默认配置
        if config is None:
            config = load_llm_config()
            if config is None:
                print("无法加载配置，生成代码失败")
                return False

        # 初始化OpenAI客户端
        client = OpenAI(
            api_key=config.api_key,
            base_url=config.base_url
        )

        reasoning_content = ""  # 定义完整思考过程
        answer_content = ""  # 定义完整回复
        is_answering = False  # 判断是否结束思考过程并开始回复

        # 创建聊天完成请求
        completion = client.chat.completions.create(
            model=config.model_name,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text",
                         "text": prompt},
                    ],
                },
            ],
            stream=True,
        )

        print("\n" + "=" * 20 + "思考过程" + "=" * 20 + "\n")

        for chunk in completion:
            # 如果chunk.choices为空，则打印usage
            if not chunk.choices:
                print("\nUsage:")
                print(chunk.usage)
            else:
                delta = chunk.choices[0].delta
                # 打印思考过程
                if hasattr(delta, 'reasoning_content') and delta.reasoning_content != None:
                    print(delta.reasoning_content, end='', flush=True)
                    reasoning_content += delta.reasoning_content
                else:
                    # 开始回复
                    if delta.content != "" and is_answering is False:
                        print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                        is_answering = True
                    # 打印回复过程
                    print(delta.content, end='', flush=True)
                    answer_content += delta.content

        # 处理并保存生成的代码
        return save_code_to_file(answer_content, save_path)
    except Exception as e:
        print(f"调用大模型时出错: {str(e)}")
        return False


def save_code_to_file(content, file_path):
    try:
        # 提取Python代码块
        code_match = re.search(r'```python(.*?)```', content, re.DOTALL)
        if code_match:
            code_content = code_match.group(1).strip()
        else:
            # 如果没有代码块，尝试提取所有内容
            code_content = content.strip()

        # 替换全角符号为半角
        code_content = code_content.replace('，', ',').replace('。', '.').replace('；', ';')
        code_content = code_content.replace('（', '(').replace('）', ')').replace('【', '[').replace('】', ']')

        # 添加UTF-8编码声明
        code_content = f"# -*- coding: utf-8 -*-.py\n{code_content}"

        # 确保目录存在
        os.makedirs(os.path.dirname(file_path), exist_ok=True)

        # 写入文件
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(code_content)

        print(f"代码已成功保存到: {file_path}")
        return True
    except Exception as e:
        print(f"保存代码时出错: {str(e)}")
        return False