"""
AGoT Prompt Templates

Defines the 6 LLM Agent types from the AGoT paper (Pandey et al., 2025):

  Agent | Role
  ------+-----
  T_∅   | Generate initial thoughts from empty graph
  T_0   | Generate initial thoughts for a sub-graph from parent context
  T_e   | Generate new thoughts + edges given existing graph state
  C     | Complexity classifier — decides if a thought needs recursive decomposition
  Eval  | Evaluate a thought or produce the final answer
  Φ     | Synthesize the final answer from the graph
"""

# ── T_∅: Generate Initial Thoughts (from empty graph) ─────────────
# Used at the top level when no prior reasoning exists.

T_EMPTY_SYSTEM = """You are a reasoning assistant that decomposes problems into initial sub-problems.
Given a question, generate a set of distinct, parallel initial approaches or sub-questions that would help solve it. Be thorough but concise."""

T_EMPTY_PROMPT = """Decompose the following question into {n_thoughts} distinct initial approaches or sub-questions:

Question: {question}

Each approach should tackle a different aspect of the problem. Return them as a numbered list.

1."""

# ── T_0: Generate Initial Thoughts for a Sub-graph ────────────────
# Used when recursing into a sub-problem (d > 0).

T_ZERO_SYSTEM = """You are a reasoning assistant solving a sub-problem within a larger reasoning task.
Given the parent context and the specific sub-problem, generate focused initial approaches."""

T_ZERO_PROMPT = """Parent problem context:
{parent_context}

The current sub-problem to solve is:
{sub_problem}

Generate {n_thoughts} focused approaches to solve this sub-problem.

1."""

# ── T_e: Generate New Thoughts + Edges ────────────────────────────
# Used within a layer to extend the graph.

T_EXTEND_SYSTEM = """You are a reasoning assistant that extends a graph of thoughts.
Given the current state of reasoning, generate new thoughts and specify how they connect to existing ones.
Use existing nodes as stepping stones — build on previous insights, refine, or challenge them."""

T_EXTEND_PROMPT = """Problem: {question}

Current reasoning graph (nodes and their contents):
{nodes_str}

Existing connections:
{edges_str}

Generate {n_thoughts} new thoughts that advance the reasoning. For each thought:
- State the thought content
- List which existing node IDs it connects FROM (can be multiple)

Additional guidance:
- Build on the most promising existing ideas
- Cross-validate between different lines of reasoning
- Identify and reconcile contradictions

Format each as:
Thought: <content>
Connects from: <node_ids>

Thought 1:"""

# ── C: Complexity Classifier ──────────────────────────────────────
# Decides whether a thought needs recursive decomposition.

C_SYSTEM = """You are a complexity classifier. Determine whether a given reasoning step or sub-problem is simple enough to solve directly, or complex enough to warrant recursive decomposition.

Respond with only one word: "SIMPLE" or "COMPLEX"."""

C_PROMPT = """Is the following reasoning step COMPLEX (needs further decomposition) or SIMPLE (can be solved directly)?

Step: {thought}

Consider:
- Does it involve multiple sub-steps?
- Does it require integration of diverse knowledge?
- Can it be answered in a single, straightforward pass?

Answer (SIMPLE or COMPLEX):"""

# ── Eval: Evaluate a Thought ──────────────────────────────────────
# Evaluates a thought's contribution to solving the problem.

EVAL_SYSTEM = """You are a thought evaluator. Assess how well a given reasoning step contributes to solving the original problem. Then either produce a sub-answer or indicate that more reasoning is needed.

Output format:
Score: <1-10>
Assessment: <brief assessment>
Sub-answer: <partial answer if applicable, or "NONE">"""

EVAL_PROMPT = """Problem: {question}

Reasoning step to evaluate:
{thought}

Assess this step's contribution:
1. Is it logically valid? (1-10)
2. Does it make progress toward the answer? (1-10)
3. What partial result does it establish?

Score:"""

# ── Φ: Synthesize Final Answer ────────────────────────────────────
# Aggregates all validated nodes into the final answer.

PHI_SYSTEM = """You are a synthesis assistant. Given a graph of reasoning steps about a problem, synthesize the final answer. Integrate information from all nodes, resolve any conflicts, and produce a definitive answer.

Put your final answer in the format: \\boxed{answer}"""

PHI_PROMPT = """Based on the following reasoning graph, synthesize the final answer to the original question.

Question: {question}

Reasoning nodes (each with evaluation scores):
{nodes_str}

Key connections between ideas:
{edges_str}

Synthesize all information and provide the final answer in \\boxed{{}}."""


# ── Chinese Versions (中文版 AGoT 提示) ────────────────────────────

T_EMPTY_SYSTEM_ZH = """你是一个将问题分解为初始子问题的推理助手。
给定一个问题，生成一组不同的、并行的初始思路或子问题来帮助解决它。要全面但简洁。"""

T_EMPTY_PROMPT_ZH = """将以下问题分解为 {n_thoughts} 个不同的初始思路或子问题：

问题：{question}

每个思路应该解决问题的不同方面。请用编号列表返回。

1."""

T_ZERO_SYSTEM_ZH = """你正在一个更大的推理任务中解决一个子问题。
给定父问题的上下文和当前的子问题，生成有针对性的初始思路。"""

T_ZERO_PROMPT_ZH = """父问题上下文：
{parent_context}

当前要解决的子问题是：
{sub_problem}

生成 {n_thoughts} 个有针对性的思路来解决这个子问题。

1."""

T_EXTEND_SYSTEM_ZH = """你是一个扩展思维图的推理助手。
给定当前的推理状态，生成新的想法并说明它们如何与已有想法连接。
利用已有的节点作为跳板——在已有的洞察上构建、精炼或质疑它们。"""

T_EXTEND_PROMPT_ZH = """问题：{question}

当前推理图（节点及其内容）：
{nodes_str}

已有连接：
{edges_str}

生成 {n_thoughts} 个能推进推理的新想法。对每个想法：
- 说明想法内容
- 列出它从哪些已有节点 ID 连接而来（可以是多个）

额外指导：
- 在最有希望的想法上继续构建
- 对不同推理路径进行交叉验证
- 识别并调合矛盾

格式：
想法：<内容>
从节点连接：<节点ID>

想法1："""

C_SYSTEM_ZH = """你是一个复杂度分类器。判断给定的推理步骤或子问题是简单到可以直接解决，还是复杂到需要递归分解。

只回复一个词："SIMPLE" 或 "COMPLEX"（简答或复杂）。"""

C_PROMPT_ZH = """以下推理步骤是 COMPLEX（需要进一步分解）还是 SIMPLE（可以直接解决）？

步骤：{thought}

考虑：
- 是否涉及多个子步骤？
- 是否需要整合多个领域的知识？
- 能否通过单次推理直接回答？

答案（SIMPLE 或 COMPLEX）："""

EVAL_SYSTEM_ZH = """你是一个思维评估器。评估给定的推理步骤对解决原问题的贡献。然后给出子答案或指示需要更多推理。

输出格式：
得分：<1-10>
评估：<简要评估>
子答案：<部分答案（如适用），或"无">"""

EVAL_PROMPT_ZH = """问题：{question}

要评估的推理步骤：
{thought}

评估此步骤的贡献：
1. 逻辑是否有效？（1-10分）
2. 是否在向答案推进？（1-10分）
3. 它建立了什么部分结果？

得分："""

PHI_SYSTEM_ZH = """你是一个综合助手。给定关于一个问题的推理图，综合出最终答案。整合来自所有节点的信息，解决任何冲突，给出确定的答案。

最终答案请放在 \\boxed{答案} 中。"""

PHI_PROMPT_ZH = """基于以下推理图，综合出原始问题的最终答案。

问题：{question}

推理节点（含评估分数）：
{nodes_str}

关键连接：
{edges_str}

综合所有信息，在 \\boxed{{}} 中给出最终答案。"""


# ── Language Selector ──────────────────────────────────────────────

def is_chinese_task(task_name: str) -> bool:
    """Detect whether the task uses Chinese prompts."""
    return task_name.lower().startswith("c_") or task_name.lower() in ("cmath", "c_aime")

def get_agot_prompts(task_name: str = "") -> dict:
    """Return the appropriate AGoT prompt set based on task language."""
    zh = is_chinese_task(task_name)
    return {
        "T_EMPTY_SYSTEM": T_EMPTY_SYSTEM_ZH if zh else T_EMPTY_SYSTEM,
        "T_EMPTY_PROMPT": T_EMPTY_PROMPT_ZH if zh else T_EMPTY_PROMPT,
        "T_ZERO_SYSTEM": T_ZERO_SYSTEM_ZH if zh else T_ZERO_SYSTEM,
        "T_ZERO_PROMPT": T_ZERO_PROMPT_ZH if zh else T_ZERO_PROMPT,
        "T_EXTEND_SYSTEM": T_EXTEND_SYSTEM_ZH if zh else T_EXTEND_SYSTEM,
        "T_EXTEND_PROMPT": T_EXTEND_PROMPT_ZH if zh else T_EXTEND_PROMPT,
        "C_SYSTEM": C_SYSTEM_ZH if zh else C_SYSTEM,
        "C_PROMPT": C_PROMPT_ZH if zh else C_PROMPT,
        "EVAL_SYSTEM": EVAL_SYSTEM_ZH if zh else EVAL_SYSTEM,
        "EVAL_PROMPT": EVAL_PROMPT_ZH if zh else EVAL_PROMPT,
        "PHI_SYSTEM": PHI_SYSTEM_ZH if zh else PHI_SYSTEM,
        "PHI_PROMPT": PHI_PROMPT_ZH if zh else PHI_PROMPT,
    }


# ── Utility ───────────────────────────────────────────────────────

def format_graph_state(nodes: list, edges: list) -> tuple:
    """Convert AGoT graph state to formatted strings for prompts."""
    nodes_str = "\n".join(
        f"  [{n['id']}] (depth={n.get('depth', 0)}, score={n.get('score', '?'):>4}) {n['content'][:150]}"
        for n in nodes
    )
    edges_str = "\n".join(
        f"  {e['from']} → {e['to']}" for e in edges
    ) if edges else "  (no connections yet)"
    return nodes_str, edges_str
