import os
from typing import List, TypedDict, Annotated, Optional
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from tools.searchTools_lg import wiki_search, mini_web_search, arvix_search

class AgentState(TypedDict):
    input_file: Optional[str]
    messages: Annotated[list[AnyMessage], add_messages]


 # add toools:
tools = [
        wiki_search,
        mini_web_search,
        arvix_search,

    ]

def agent(state:AgentState):
    tools_description = """
    wiki_search(query: str) -> str:
    Search Wikipedia for a query and return maximum 2 results.
    
    Args:
        query: The search query.

    Search Tavily for a query and return maximum 3 results.
    
    mini_web_search(query: str) -> str:
    Args:
        query: The search query

    arvix_search(query: str) -> str:
    Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query.
    """
    sys_message = SystemMessage(content=f"""You are a helpful AI agent that can use tools to answer questions.
    You can use the following tools:{tools_description}
    PLEASE FOLLOW THE INSTRUCTIONS FOR ANSWERING CAREFULLY:
    Your answer should follow the 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.                                
    """)

   


    # LLM model and tools
    vision_llm = ChatOpenAI(model="gpt-4o")
    llm = ChatOpenAI(model="gpt-4o")
    llm_withtools = llm.bind_tools(tools, parallel_tool_calls = False)

    return {
    "input_file": state["input_file"],
    "messages": [llm_withtools.invoke([sys_message]+ state["messages"])]
    }


def build_agents():

    builder = StateGraph(AgentState)
    builder.add_node("agent",agent)
    builder.add_node("tools",ToolNode(tools))

    builder.add_edge(START, "agent")
    builder.add_conditional_edges(
        "agent",
        tools_condition,
    )
    builder.add_edge("tools", "agent")
    react_graph = builder.compile()

    return react_graph