from typing import TypedDict, Annotated, Sequence
import operator
import re
from langgraph.graph import StateGraph, END
from .ai_tools import Calculator, DocRetriever, WebSearcher

class AgentState(TypedDict):
    input: str
    context: Annotated[Sequence[str], operator.add]
    last_tool: str
    output: str

class GaiaGraph:
    def __init__(self, model, tokenizer, tools):
        self.model = model
        self.tokenizer = tokenizer
        self.tools = tools
        self.tool_map = {tool.name: tool for tool in tools}
        self.graph = self._build_graph()
    
    def _build_graph(self):
        graph = StateGraph(AgentState)
        
        graph.add_node("agent", self._agent_node)
        graph.add_node("tool", self._tool_node)
        graph.set_entry_point("agent")
        
        graph.add_edge("agent", "tool")
        graph.add_conditional_edges(
            "tool",
            self._route_action,
            {"continue": "agent", "end": END}
        )
        
        return graph.compile()
    
    def _agent_node(self, state: AgentState) -> dict:
        tool_list = "\n".join([f"- {t.name}: {t.description}" for t in self.tools])
        prompt = f"""<|system|>
You're an expert problem solver. Use these tools when needed:
{tool_list}

Respond ONLY in this format:
Thought: <your reasoning>
Action: <tool_name or 'FINISH'>
Action Input: <input for tool>
</s>
<|user|>
{state['input']}
Context: {state['context']}
</s>
<|assistant|>
"""
        
        response = self.model(
            prompt,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.2,
            pad_token_id=self.tokenizer.eos_token_id
        )[0]['generated_text']
        
        # Extract tool call
        action_match = re.search(r"Action: (\w+)", response)
        action_input_match = re.search(r"Action Input: (.+?)\n", response, re.DOTALL)
        
        if action_match and action_input_match:
            tool_name = action_match.group(1)
            tool_input = action_input_match.group(1).strip()
            return {
                "last_tool": tool_name,
                "tool_input": tool_input,
                "thought": response
            }
        else:
            return {"last_tool": "FINISH", "output": response}
    
    def _tool_node(self, state: AgentState) -> dict:
        if state["last_tool"] == "FINISH":
            return {"output": state.get("output", "No output generated")}
        
        tool = self.tool_map.get(state["last_tool"])
        if not tool:
            return {"context": f"Error: Unknown tool {state['last_tool']}"}
        
        result = tool.run(state["tool_input"])
        return {"context": f"Tool {tool.name} returned: {result}"}
    
    def _route_action(self, state: AgentState) -> str:
        return "end" if state["last_tool"] == "FINISH" else "continue"
    
    def run(self, input: str) -> str:
        state = {"input": input, "context": [], "last_tool": "", "output": ""}
        for step in self.graph.stream(state):
            for node, value in step.items():
                if node == "__end__":
                    return value["output"]
        return "Execution completed without output"