import ast
import re
from openai import OpenAI

client = OpenAI()


# ============================================================
# Planner：接收问题，输出 Python list 格式的步骤列表
# ============================================================

class Planner:
    def __init__(self):
        self.system_prompt = """你是一个任务规划专家。
将用户的问题拆解成清晰的执行步骤，以 Python list 格式输出。

要求：
1. 每个步骤是一个字符串，描述具体要做什么
2. 步骤要足够细，每步只做一件事
3. 最后一步必须是"汇总前面所有结果，给出最终答案"
4. 只输出 Python list，不要有任何其他文字

输出示例：
["步骤1：计算...", "步骤2：根据步骤1的结果计算...", "步骤3：汇总前面所有结果，给出最终答案"]"""

    def plan(self, question: str) -> list[str]:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": question},
            ],
            temperature=0,
        )
        raw_output = response.choices[0].message.content.strip()

        # 用 ast.literal_eval 安全解析，避免 eval 执行任意代码
        try:
            steps = ast.literal_eval(raw_output)
            if isinstance(steps, list) and all(isinstance(s, str) for s in steps):
                return steps
        except (ValueError, SyntaxError):
            pass

        # 解析失败：尝试用正则提取列表内容
        items = re.findall(r'"([^"]+)"', raw_output)
        if items:
            return items

        # 兜底：把原始输出作为单步计划
        return [raw_output, "汇总前面所有结果，给出最终答案"]


# ============================================================
# Executor：逐步执行，维护 history 字符串积累上下文
# ============================================================

class Executor:
    def __init__(self):
        self.system_prompt = """你是一个严谨的执行者。
你会收到：原始问题、完整计划、当前要执行的步骤、以及之前步骤的执行结果。
只执行当前步骤，给出该步骤的结果，不要超前执行其他步骤。
输出要简洁，只写计算过程和结果，不要废话。"""

    def execute_step(
        self,
        question: str,
        full_plan: list[str],
        current_step: str,
        step_index: int,
        history: str,
    ) -> str:
        plan_text = "\n".join(f"{i+1}. {s}" for i, s in enumerate(full_plan))

        user_content = f"""原始问题：{question}

完整计划：
{plan_text}

之前步骤的执行结果：
{history if history else "（这是第一步，暂无历史结果）"}

当前执行步骤（第 {step_index + 1} 步）：{current_step}

请执行当前步骤，给出结果："""

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": user_content},
            ],
            temperature=0,
        )
        return response.choices[0].message.content.strip()


# ============================================================
# PlanAndSolveAgent：组合 Planner 和 Executor
# ============================================================

class PlanAndSolveAgent:
    def __init__(self):
        self.planner = Planner()
        self.executor = Executor()

    def run(self, question: str) -> str:
        print(f"\n问题：{question}")
        print("=" * 60)

        # 规划阶段
        print("\n[规划阶段]")
        steps = self.planner.plan(question)
        for i, step in enumerate(steps):
            print(f"  {i + 1}. {step}")

        # 执行阶段
        print("\n[执行阶段]")
        history = ""

        for i, step in enumerate(steps):
            print(f"\n步骤 {i + 1}：{step}")
            result = self.executor.execute_step(
                question=question,
                full_plan=steps,
                current_step=step,
                step_index=i,
                history=history,
            )
            print(f"结果：{result}")

            # 追加到历史，供后续步骤使用
            history += f"\n步骤 {i + 1}（{step}）的结果：\n{result}\n"

        # 最后一步的结果就是最终答案
        return result


# ============================================================
# 测试
# ============================================================

if __name__ == "__main__":
    agent = PlanAndSolveAgent()

    question = (
        "小明周一买了 15 个苹果，周二又买了周一数量的 2 倍，"
        "周三吃掉了 5 个。请问小明现在一共有多少个苹果？"
        "并计算周一和周二平均每天新增多少个苹果？"
    )

    answer = agent.run(question)
    print(f"\n最终答案：{answer}")

问题：小明周一买了 15 个苹果...
============================================================

[规划阶段]
  1. 计算周一购买的苹果数量
  2. 计算周二购买的苹果数量（周一的2倍）
  3. 计算周三之后剩余的苹果数量（扣除吃掉的5个）
  4. 计算周一和周二平均每天新增苹果数
  5. 汇总前面所有结果，给出最终答案

[执行阶段]

步骤 1：计算周一购买的苹果数量
结果：周一购买 15 个苹果。

步骤 2：计算周二购买的苹果数量（周一的2倍）
结果：周二购买 15 × 2 = 30 个苹果。

步骤 3：计算周三之后剩余的苹果数量
结果：总计 15 + 30 = 45 个，吃掉 5 个后剩余 40 个苹果。

步骤 4：计算周一和周二平均每天新增苹果数
结果：周一新增 15 个，周二新增 30 个，平均 (15 + 30) / 2 = 22.5 个。

步骤 5：汇总所有结果，给出最终答案
结果：小明现在共有 40 个苹果。周一和周二平均每天新增 22.5 个苹果。