import os
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.prebuilt import tools_condition, ToolNode
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.tools import tool
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper
from langchain_core.messages import SystemMessage, HumanMessage

# Lade Umgebungsvariablen (Google API Key)
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# === Tools definieren ===
@tool
def multiply(a: int, b: int) -> int:
    """Multiplies two numbers."""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Adds two numbers."""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtracts two numbers."""
    return a - b

@tool
def divide(a: int, b: int) -> float:
    """Divides two numbers."""
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulo(a: int, b: int) -> int:
    """Returns the remainder of dividing two numbers."""
    return a % b

@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return the result."""
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    return "\n\n".join(doc.page_content for doc in search_docs)

@tool
def arxiv_search(query: str) -> str:
    """Search Arxiv for academic papers about a query."""
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    return "\n\n".join(doc.page_content[:1000] for doc in search_docs)

@tool
def web_search(query: str) -> str:
    """Perform a DuckDuckGo web search."""
    wrapper = DuckDuckGoSearchAPIWrapper(max_results=5)
    results = wrapper.run(query)
    return results

# === System Prompt definieren ===
system_prompt = SystemMessage(content=(
    "You are an expert assistant. You MUST answer precisely, factually, and accurately. "
    "If you do not know the answer, use the available tools such as Wikipedia Search, Arxiv Search, "
    "or Web Search to find the correct information. "
    "If a math operation is needed, use the calculation tools. "
    "Do NOT invent answers. Only return answers you are confident in."
))

# === LLM definieren ===
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key=GOOGLE_API_KEY,
    temperature=0,
    max_output_tokens=2048,
    system_message=system_prompt,
)

# === Tools in LLM einbinden ===
tools = [multiply, add, subtract, divide, modulo, wiki_search, arxiv_search, web_search]
llm_with_tools = llm.bind_tools(tools)

# === Nodes für LangGraph ===
def assistant(state: MessagesState):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# === LangGraph bauen ===
builder = StateGraph(MessagesState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")

# === Agent Executor ===
agent_executor = builder.compile()