import os
from langchain_groq import ChatGroq
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_core.messages import SystemMessage, HumanMessage
from tools import *
from typing import TypedDict, Optional

tools = [add, multiply, divide, subtract, search_wikipedia, duckduckgo_search, search_papers]

def build_agent():
    llm = ChatGroq(model="llama3-70b-8192", temperature=0) # deepseek-r1-distill-llama-70b
    chat_with_tools = llm.bind_tools(tools)

    def assistant(state: MessagesState):
        response = chat_with_tools.invoke(state["messages"])
        return {
            "messages": [response]
        }

    def enhancer(state: MessagesState):
        system_prompt = """You are a reasoning assistant that can answer complex questions using external tools. You must follow a structured reasoning format and always output a final answer.
                            Always follow this step-by-step structure:
                            - Thought: Describe your reasoning.
                            - Action:  Whenever necessary or recommended, call a tool, using this format → Action: tool_name(input).
                                       If not, use your reasoning for answering the question, using this format → Action: reasoning.
                            - Observation: Describe what the tool returned.
                            - Repeat Thought → Action → Observation as needed until you are SURE you can answer.
                            When you are ready to answer, say:
                            FINAL ANSWER: [your final answer here]
                            Formatting rules for FINAL ANSWER:
                            - If a number: no commas, no units unless explicitly asked.
                            - If a string: no articles, no abbreviations, digits in plain text.
                            - If a comma-separated list: apply the above rules to each element.
                            - Only include the final answer, do not include any reasoning steps."""

        sys_msg = SystemMessage(content=system_prompt)
        
        return {
                "messages": [sys_msg] + state["messages"]
            }

    ## The graph
    builder = StateGraph(MessagesState)
    
    # Define nodes: these do the work
    builder.add_node("enhancer", enhancer)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    
    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "enhancer")
    builder.add_edge("enhancer", "assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message requires a tool, route to tools
        # Otherwise, provide a direct response
        tools_condition,
    )
    builder.add_edge("tools", "assistant")

    return builder.compile()