import os

from langchain.tools import Tool
from langgraph.graph import StateGraph, END, MessagesState
from langchain.agents import AgentExecutor, initialize_agent
from langchain.memory import ConversationBufferMemory
from langchain_community.tools import DuckDuckGoSearchRun
#from langchain.schema import SystemMessage
#from langchain.schema.messages import HumanMessage
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
GAIA_SYSTEM_PROMPT = """
You are an advanced autonomous agent designed to pass the GAIA benchmark.
You must:

- Think step-by-step before answering.
- Use available tools (e.g., Wikipedia, arXiv) to gather accurate information.
- Cite scientific sources when applicable.
- Reason before deciding what tool to use.
- When unsure, search or say "I need more information."

Available tools:
- Wikipedia for factual and general knowledge.
- arXiv for up-to-date scientific papers.
- You may call tools multiple times before giving a final answer.

Think aloud and explain your process clearly.
"""
#google_api_key = os.environ.get("GOOGLE_API_KEY")
from langchain_google_genai import ChatGoogleGenerativeAI
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key="",
    temperature=0.2,
    max_output_tokens=1024
)
messages = [
    SystemMessage(content=GAIA_SYSTEM_PROMPT)  ,
    HumanMessage(content="Get started, messages on the way.")
]

response =  llm.invoke(messages)  

def calculator_tool(input: str) -> str:
    try:
        return str(eval(input))
    except Exception as e:
        return f"Error: {e}"

calc_tool = Tool(name="Calculator", func=calculator_tool, description="Perform basic math.")
search_tool = DuckDuckGoSearchRun()

import wikipedia
def search_wikipedia(query: str) -> str:
    try:
        return wikipedia.summary(query, sentences=3)
    except Exception as e:
        return f"Wikipedia Error: {str(e)}"

wikipedia_tool = Tool(
    name="Wikipedia Search",
    func=search_wikipedia,
    description="Use this to search and summarize articles from Wikipedia."
)

import arxiv
def search_arxiv(query: str) -> str:
    try:
        search = arxiv.Search(
            query=query,
            max_results=1,
            sort_by=arxiv.SortCriterion.Relevance
        )
        for result in search.results():
            return f"{result.title}\n\n{result.summary}\n\nPDF: {result.pdf_url}"
        return "No relevant arXiv papers found."
    except Exception as e:
        return f"arXiv Error: {str(e)}"

arxiv_tool = Tool(
    name="arXiv Search",
    func=search_arxiv,
    description="Search for scientific papers on arXiv.org"
)


class GaiaState:
    input: str
    intermediate_steps: list = []
    output: str = ""

# Initialize tools and agent
#llm = ChatOpenAI(temperature=0,model_name="gpt-4", system_message=SystemMessage(content=GAIA_SYSTEM_PROMPT))
tools = [calc_tool, search_tool, arxiv_tool, wikipedia_tool]  # Add more tools as needed

agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent_type="openai-functions",
    verbose=True
)

def run_agent(state: MessagesState) -> MessagesState:
    # Let the agent respond based on the full message history
    result = agent.invoke(state["messages"])

    # Add the AI's response to the message history
    state["messages"].append(AIMessage(content=result.content))

    return state   
    
graph = StateGraph(MessagesState)

graph.add_node("run_agent", run_agent)
graph.set_entry_point("run_agent")
graph.add_edge("run_agent", END)

# Compile the graph
app = graph.compile()