"""
ReAct Agent powered by Groq (llama-3.3-70b-versatile).
Runs a Thought → Action → Observation loop until it reaches a Final Answer.
"""

import os
import re
import json
from groq import Groq
from tools import TOOLS

# ── System prompt ─────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are a precise, detail-oriented AI assistant solving questions from the GAIA benchmark.
You answer step-by-step using a ReAct loop: Thought, Action, Observation, repeat, then Final Answer.

Available tools:
{tool_descriptions}

FORMAT — always follow this exactly:
Thought: <your reasoning about what to do next>
Action: <tool_name>
Action Input: <input to the tool, as plain text>

After you receive an Observation, continue with another Thought/Action or:
Final Answer: <your concise, exact answer>

Rules:
- Use tools whenever you need external facts, calculations, or files.
- Final Answer must be concise and exact — no extra sentences.
- If the answer is a number, give just the number (with units if asked).
- If the answer is a list, comma-separate it.
- Never make up facts. If unsure, search again.
- Do NOT include the text "FINAL ANSWER:" in capitals — just write "Final Answer:".
"""

def _build_tool_descriptions() -> str:
    lines = []
    for name, meta in TOOLS.items():
        lines.append(f"- {name}: {meta['description']}")
    return "\n".join(lines)


FILLED_SYSTEM = SYSTEM_PROMPT.format(tool_descriptions=_build_tool_descriptions())


class GAIAAgent:
    def __init__(self, max_steps: int = 8):
        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        self.model = "llama-3.3-70b-versatile"
        self.max_steps = max_steps

    def _call_llm(self, messages: list) -> str:
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=0.0,
            max_tokens=1024,
        )
        return response.choices[0].message.content.strip()

    def _parse_action(self, text: str):
        """
        Returns (tool_name, tool_input) if an Action is present, else None.
        """
        action_match = re.search(r"Action:\s*(.+)", text)
        input_match = re.search(r"Action Input:\s*([\s\S]+?)(?=\nThought:|\nAction:|\nObservation:|\nFinal Answer:|$)", text)
        if action_match and input_match:
            tool = action_match.group(1).strip()
            tool_input = input_match.group(1).strip()
            return tool, tool_input
        return None

    def _parse_final_answer(self, text: str):
        match = re.search(r"Final Answer:\s*([\s\S]+)", text, re.IGNORECASE)
        if match:
            return match.group(1).strip()
        return None

    def _run_tool(self, tool_name: str, tool_input: str) -> str:
        if tool_name not in TOOLS:
            return f"Unknown tool '{tool_name}'. Available tools: {', '.join(TOOLS.keys())}"
        try:
            return TOOLS[tool_name]["fn"](tool_input)
        except Exception as e:
            return f"Tool error: {e}"

    def __call__(self, question: str, task_id: str = "") -> str:
        """Run the ReAct loop and return the final answer string."""
        messages = [
            {"role": "system", "content": FILLED_SYSTEM},
            {"role": "user", "content": (
                f"Task ID: {task_id}\n\nQuestion: {question}\n\n"
                "Begin your ReAct loop. If the question mentions an attached file, "
                f"use get_task_file with task_id='{task_id}' first."
            )},
        ]

        for step in range(self.max_steps):
            llm_output = self._call_llm(messages)
            messages.append({"role": "assistant", "content": llm_output})

            # Check for final answer first
            final = self._parse_final_answer(llm_output)
            if final:
                return final

            # Check for action
            action_result = self._parse_action(llm_output)
            if action_result:
                tool_name, tool_input = action_result
                observation = self._run_tool(tool_name, tool_input)
                obs_msg = f"Observation: {observation}"
                messages.append({"role": "user", "content": obs_msg})
            else:
                # LLM didn't produce a valid action or final answer — nudge it
                messages.append({
                    "role": "user",
                    "content": (
                        "Please continue. Either use a tool (Action + Action Input) "
                        "or provide your Final Answer."
                    )
                })

        # Exhausted steps — ask for best guess
        messages.append({
            "role": "user",
            "content": "You've reached the step limit. Provide your best Final Answer now.",
        })
        last = self._call_llm(messages)
        final = self._parse_final_answer(last)
        return final if final else last.strip()