import os
from langchain_google_genai import ChatGoogleGenerativeAI
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

tools = [add, substract, multiply, divide, web_search]

def build_agent():
    llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    chat_with_tools = llm.bind_tools(tools)

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

    def enhancer(state: MessagesState):
        sys_msg = """
                            You are a helpful assistant tasked with answering questions using a set of tools. 
                            Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
                            FINAL ANSWER: [YOUR FINAL ANSWER]. 
                            YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
                            If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
                            If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
                            If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
                            Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
                        """
        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()