"""
AGoT: Adaptive Graph of Thoughts — Core Algorithm

Implements the AGoT algorithm (Pandey et al., 2025):

  AGoT(q, h, G_h^p):
    V_h ← ∅, E_h ← ∅, F_h ← ∅        # Initialize nested graph
    G_h ← (V_h, E_h, F_h)
    d ← |h|                            # Current depth
    for l = 0, ..., l_max-1:           # Iterate layers
      if l = 0 and d = 0:  thoughts ← T_∅(q)
      elif l = 0:          thoughts ← T_0(q, G_h)
      else:                thoughts, edges ← T_e(q, G_h)
      if thought is final: return Eval(thought, G_h)
      for thought in thoughts:
        if is_complex(thought) and d < d_max:
          AGoT(thought, h', G_h)       # Recursive decomposition
        else:
          Eval(thought, G_h)
    return Φ(G_h)                      # Synthesize final answer

6 Agent types: T_∅, T_0, T_e, C, Eval, Φ
"""

from __future__ import annotations
import json
import os
import sys
import time
from pathlib import Path
from typing import Optional, Callable

# Add project root to path
_project_root = Path(__file__).resolve().parent.parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from src.agot.graph_utils import AGoTGraph
from src.agot.prompts import (
    T_EMPTY_SYSTEM, T_EMPTY_PROMPT,
    T_ZERO_SYSTEM, T_ZERO_PROMPT,
    T_EXTEND_SYSTEM, T_EXTEND_PROMPT,
    C_SYSTEM, C_PROMPT,
    EVAL_SYSTEM, EVAL_PROMPT,
    PHI_SYSTEM, PHI_PROMPT,
    get_agot_prompts,
    format_graph_state,
)
from src.llm_client import LLMClient, make_client


# ── Config ───────────────────────────────────────────────────────────

AGOT_DEFAULT_CONFIG = {
    "model": "qwen2.5:32b",
    "max_depth": 2,           # d_max: maximum recursion depth
    "max_layers": 3,          # l_max: layers per depth level
    "max_nodes_per_layer": 3, # n_max: max nodes generated per layer
    "temperature": 0.3,
    "self_terminate": True,   # stop early if a final thought emerges
}


# ── AGoT Engine ──────────────────────────────────────────────────────

class AGoTEngine:
    """Orchestrates the AGoT algorithm using local LLM calls."""

    def __init__(self, config: dict = None, task_name: str = ""):
        self.config = {**AGOT_DEFAULT_CONFIG, **(config or {})}
        self.client = make_client(
            model=self.config["model"],
            temperature=self.config["temperature"],
            max_tokens=1024,
        )
        self.prompts = get_agot_prompts(task_name)
        self.stats = {
            "total_calls": 0,
            "total_tokens": 0,
            "start_time": None,
            "end_time": None,
            "node_count": 0,
            "edge_count": 0,
            "recursion_count": 0,
            "layers_used": [],
        }

    def solve(self, question: str, task_name: str = "") -> dict:
        """
        Run AGoT on a question.

        Returns:
          {"answer": str, "graph": AGoTGraph, "stats": dict}
        """
        self.stats["start_time"] = time.time()
        graph = AGoTGraph(depth=0)

        try:
            self._agot_recursive(question, graph, depth=0)
        except Exception as e:
            print(f"  [AGoT Error] {e}")

        # Fallback if AGoT couldn't produce an answer
        if graph.final_answer is None:
            if graph.nodes:
                # Give phi one more chance with all nodes
                answer = self._agent_phi(question, graph)
                graph.final_answer = answer or self._fallback_cot(question, task_name)
            else:
                graph.final_answer = self._fallback_cot(question, task_name)

        self.stats["end_time"] = time.time()
        self.stats["elapsed"] = round(self.stats["end_time"] - self.stats["start_time"], 2)
        self.stats["node_count"] = graph.node_count
        self.stats["edge_count"] = graph.edge_count

        return {
            "answer": graph.final_answer,
            "graph": graph,
            "stats": dict(self.stats),
        }

    # ── Core Recursive Algorithm ────────────────────────────────

    def _agot_recursive(self, question: str, graph: AGoTGraph, depth: int):
        """AGoT(q, h, G_h^p) — the core recursive algorithm."""
        l_max = self.config["max_layers"]
        d_max = self.config["max_depth"]
        n_max = self.config["max_nodes_per_layer"]
        self_terminate = self.config["self_terminate"]

        for layer in range(l_max):
            self.stats["layers_used"].append((depth, layer))

            # ── Generate Thoughts ─────────────────────────────
            if layer == 0 and depth == 0:
                # T_∅: Initial thoughts from empty graph
                thoughts = self._agent_T_empty(question, n_max)
            elif layer == 0 and depth > 0:
                # T_0: Initial thoughts for sub-graph
                parent_context = self._format_parent_context(graph)
                thoughts = self._agent_T_zero(question, parent_context, n_max)
            else:
                # T_e: Generate new thoughts + edges
                thoughts_with_edges = self._agent_T_extend(question, graph, n_max)

            # ── Process thoughts ──────────────────────────────
            for thought_data in thoughts:
                # Check for early termination
                if self_terminate and self._is_final_thought(thought_data):
                    answer = self._agent_eval_final(thought_data, question)
                    graph.final_answer = answer
                    return

                # Add node
                thought_content = thought_data if isinstance(thought_data, str) else thought_data.get("content", str(thought_data))
                node_id = graph.add_node(thought_content, layer=layer)

                # Add edges if provided (T_e format)
                if isinstance(thought_data, dict) and "connects_from" in thought_data:
                    for src in thought_data["connects_from"]:
                        if graph.get_node(src):
                            graph.add_edge(src, node_id, label="extends")

                # ── Complexity Check ──────────────────────────
                if depth < d_max:
                    is_complex = self._agent_complexity(thought_content)
                    graph.update_node(node_id, is_complex=is_complex)

                    if is_complex:
                        self.stats["recursion_count"] += 1
                        sub_graph = graph.get_or_create_subgraph(node_id)
                        self._agot_recursive(thought_content, sub_graph, depth + 1)
                        continue

                # ── Evaluate ──────────────────────────────────
                score = self._agent_eval(thought_content, question)
                graph.update_node(node_id, score=score)

        # ── After all layers: synthesize final answer ────────
        if graph.final_answer is None and graph.nodes:
            answer = self._agent_phi(question, graph)
            if answer:
                graph.final_answer = answer

    # ── Agent Implementations ──────────────────────────────────

    def _agent_T_empty(self, question: str, n: int) -> list:
        """T_∅: Generate initial thoughts from empty graph."""
        prompt = self.prompts["T_EMPTY_PROMPT"].format(question=question, n_thoughts=n)
        response = self._call_llm(prompt, self.prompts["T_EMPTY_SYSTEM"])
        return self._parse_numbered_list(response)

    def _agent_T_zero(self, sub_problem: str, parent_context: str, n: int) -> list:
        """T_0: Generate initial thoughts for sub-graph."""
        prompt = self.prompts["T_ZERO_PROMPT"].format(
            parent_context=parent_context,
            sub_problem=sub_problem,
            n_thoughts=n,
        )
        response = self._call_llm(prompt, self.prompts["T_ZERO_SYSTEM"])
        return self._parse_numbered_list(response)

    def _agent_T_extend(self, question: str, graph: AGoTGraph, n: int) -> list:
        """T_e: Generate new thoughts + edges."""
        nodes_str, edges_str = format_graph_state(graph.nodes, graph.edges)
        prompt = self.prompts["T_EXTEND_PROMPT"].format(
            question=question,
            nodes_str=nodes_str,
            edges_str=edges_str,
            n_thoughts=n,
        )
        response = self._call_llm(prompt, self.prompts["T_EXTEND_SYSTEM"])
        return self._parse_extended_thoughts(response, graph)

    def _agent_complexity(self, thought: str) -> bool:
        """C: Complexity classifier."""
        prompt = self.prompts["C_PROMPT"].format(thought=thought[:200])
        response = self._call_llm(prompt, self.prompts["C_SYSTEM"]).strip().upper()
        return "COMPLEX" in response

    def _agent_eval(self, thought: str, question: str) -> float:
        """Eval: Score a thought."""
        prompt = self.prompts["EVAL_PROMPT"].format(question=question, thought=thought[:300])
        response = self._call_llm(prompt, self.prompts["EVAL_SYSTEM"])
        return self._parse_score(response)

    def _agent_eval_final(self, thought: str, question: str) -> str:
        """Eval when a thought looks final — produce sub-answer."""
        return self._agent_eval_phi(question, thought)

    def _agent_eval_phi(self, question: str, final_thought: str) -> str:
        """Φ-like: produce final answer from a terminal thought."""
        ps = self.prompts
        zh = ps["T_EMPTY_SYSTEM"].startswith("你")
        if zh:
            prompt = f"问题：{question}\n\n最终推理步骤：{final_thought}\n\n基于此推理，在 \\boxed{{答案}} 中给出最终答案。"
        else:
            prompt = f"Question: {question}\n\nFinal reasoning step: {final_thought}\n\nBased on this reasoning, provide the final answer in \\boxed{{answer}}."
        response = self._call_llm(prompt, self.prompts["PHI_SYSTEM"])
        return self._extract_final_answer(response)

    def _agent_phi(self, question: str, graph: AGoTGraph) -> str:
        """Φ: Synthesize final answer from graph."""
        nodes_str, edges_str = format_graph_state(graph.nodes, graph.edges)
        prompt = self.prompts["PHI_PROMPT"].format(
            question=question,
            nodes_str=nodes_str,
            edges_str=edges_str,
        )
        response = self._call_llm(prompt, self.prompts["PHI_SYSTEM"])
        return self._extract_final_answer(response)

    def _fallback_cot(self, question: str, task_name: str = "") -> str:
        """Simple CoT fallback if AGoT errors out."""
        from src.baseline.benchmark import get_prompts as get_baseline_prompts, extract_answer
        bp = get_baseline_prompts(task_name)
        prompt = bp["cot_template"].format(question=question)
        response = self._call_llm(prompt, bp["cot_system"])
        return extract_answer(response) or ""

    # ── Parsing Helpers ────────────────────────────────────────

    def _parse_numbered_list(self, text: str) -> list:
        """Parse a numbered list into clean strings."""
        items = []
        for line in text.split("\n"):
            line = line.strip()
            if not line:
                continue
            # Match "1. content" or "1) content" or "- content"
            import re
            m = re.match(r'^(?:\d+[\.\)]|[-*])\s*(.*)', line)
            if m:
                items.append(m.group(1))
        return items[:self.config["max_nodes_per_layer"]]

    def _parse_extended_thoughts(self, text: str, graph: AGoTGraph) -> list:
        """Parse T_e output which includes connection info."""
        items = []
        current_thought = None
        current_connects = []

        for line in text.split("\n"):
            line = line.strip()
            if not line:
                continue

            import re

            # Match "Thought: content"
            tm = re.match(r'Thought\s*\d*\s*:\s*(.*)', line, re.IGNORECASE)
            if tm:
                if current_thought:
                    items.append({
                        "content": current_thought,
                        "connects_from": current_connects,
                    })
                current_thought = tm.group(1)
                current_connects = []
                continue

            # Match "Connects from: 1, 2, 3" or "Connects from: 1"
            cm = re.match(r'Connects?\s*(?:from)?\s*:\s*(.*)', line, re.IGNORECASE)
            if cm:
                ids = re.findall(r'\d+', cm.group(1))
                current_connects = [int(i) for i in ids if graph.get_node(int(i))]

        # Don't forget the last one
        if current_thought:
            items.append({
                "content": current_thought,
                "connects_from": current_connects,
            })

        return items[:self.config["max_nodes_per_layer"]]

    def _parse_score(self, text: str) -> float:
        """Extract a numeric score from Eval output."""
        import re
        m = re.search(r'(?:Score|score)\s*:?\s*(\d+(?:\.\d+)?)', text)
        if m:
            try:
                return min(10.0, max(0.0, float(m.group(1))))
            except ValueError:
                pass
        # Try just the first number
        nums = re.findall(r'\d+(?:\.\d+)?', text)
        if nums:
            try:
                return min(10.0, max(0.0, float(nums[0])))
            except ValueError:
                pass
        return 5.0

    def _extract_final_answer(self, text: str) -> str:
        """Extract \\boxed{answer} from synthesis output."""
        from src.baseline.benchmark import extract_answer
        ans = extract_answer(text)
        return ans or ""

    def _is_final_thought(self, thought: str) -> bool:
        """Heuristic: does this thought look like a final answer?"""
        if isinstance(thought, dict):
            content = thought.get("content", "")
        else:
            content = thought

        # Check for answer-like patterns
        import re
        if re.search(r'\\boxed\{', content):
            return True
        if re.search(r'(?:therefore|thus|hence|so the answer is|final answer)', content.lower()):
            # Be cautious — only if it contains a concrete answer
            if re.search(r'(?:=|=|=|\d+)', content):
                return True
        return False

    def _format_parent_context(self, graph: AGoTGraph) -> str:
        """Format parent graph state as context string."""
        if not graph.nodes:
            return "(No prior reasoning)"
        parts = [f"[{n['id']}] {n['content'][:100]}" for n in graph.nodes[-3:]]
        return "\n".join(parts)

    def _call_llm(self, prompt: str, system: str) -> str:
        """Make an LLM call and track stats."""
        self.stats["total_calls"] += 1
        resp = self.client.generate_with_metadata(prompt, system)
        content = resp["message"]["content"]
        self.stats["total_tokens"] += self.client.count_tokens(content)
        return content


# ── Convenience runner ──────────────────────────────────────────────

def run_agot(
    question: str,
    model: str = "qwen2.5:32b",
    max_depth: int = 2,
    max_layers: int = 3,
    verbose: bool = True,
    task_name: str = "",
) -> dict:
    """Run AGoT on a single question."""
    engine = AGoTEngine({
        "model": model,
        "max_depth": max_depth,
        "max_layers": max_layers,
    }, task_name=task_name)
    result = engine.solve(question, task_name)

    if verbose:
        print(f"\n{'='*60}")
        print(f"AGoT Result")
        print(f"{'='*60}")
        print(f"Question: {question[:80]}...")
        print(f"Answer:   {result['answer']}")
        print(f"Stats:    {json.dumps(result['stats'], indent=2)}")
        print(f"Graph:    {result['graph']}")

    return result


# ── CLI Entry Point ─────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run AGoT reasoning")
    parser.add_argument("--model", default="qwen2.5:32b")
    parser.add_argument("--question", default=None, help="Question to reason about")
    parser.add_argument("--interactive", action="store_true", help="Interactive mode")
    parser.add_argument("--max-depth", type=int, default=2)
    parser.add_argument("--max-layers", type=int, default=3)
    parser.add_argument("--verbose", action="store_true", default=True)
    args = parser.parse_args()

    if args.interactive:
        print("AGoT Interactive Mode. Type 'quit' to exit.")
        while True:
            q = input("\nQuestion: ").strip()
            if q.lower() in ("quit", "exit", "q"):
                break
            if q:
                run_agot(q, args.model, args.max_depth, args.max_layers, verbose=True)
    elif args.question:
        run_agot(args.question, args.model, args.max_depth, args.max_layers, verbose=True)
    else:
        # Demo question
        run_agot(
            "Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning "
            "and bakes muffins with four every day. She sells the remaining eggs at the "
            "farmers' market daily for $2 per dozen. How much money does she make per day?",
            args.model, args.max_depth, args.max_layers,
        )
