from tools import TOOL_MAP
from llm import llm_call
import re
from typing import List, Dict
from models import GraphRAGResponse

SYSTEM_PROMPT = """
You are an educational GraphRAG assistant for schools.

You MUST use tools when answering knowledge questions.

Available tool:
- query_knowledge_graph

Format:
Action: query_knowledge_graph
Action Input: <question>

OR:

Final Answer: <answer>
"""

def parse_action(text: str):
    action = re.search(r"Action:\s*(\w+)", text)
    input_ = re.search(r"Action Input:\s*(.+)", text)

    if action and input_:
        return action.group(1), input_.group(1)

    return None

def parse_final(text: str):
    m = re.search(r"Final Answer:\s*(.+)", text, re.DOTALL)
    return m.group(1).strip() if m else None

def react_agent(user_message: str, history: List[Dict], G=None, max_steps: int = 5):

    messages = (
        [{"role": "system", "content": SYSTEM_PROMPT}]
        + history
        + [{"role": "user", "content": user_message}]
    )

    last = ""

    for _ in range(max_steps):

        reply = llm_call(messages)
        last = reply

        final = parse_final(reply)
        if final:
            return GraphRAGResponse(
                question=user_message,
                answer=final,
                evidence=""
            )

        action = parse_action(reply)

        if action:

            tool_name, tool_input = action

            if tool_name == "query_knowledge_graph":
                result = TOOL_MAP[tool_name](tool_input, G)

                messages.append({"role": "assistant", "content": reply})
                messages.append({"role": "user", "content": f"Observation: {result}"})

                return GraphRAGResponse(
                    question=user_message,
                    answer=result["answer"],
                    evidence=result["evidence"]
                )

    return GraphRAGResponse(
        question=user_message,
        answer=last,
        evidence=""
)