import os
from dotenv import load_dotenv
from langchain.tools import Tool
from langchain.utilities import SerpAPIWrapper
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import AnyMessage, HumanMessage
from langchain_groq import ChatGroq
from typing import TypedDict, Annotated

# Load API keys
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")

# --- أدوات ---

def calculator_tool_func(query: str):
    try:
        result = str(eval(query, {"__builtins__": {}}))
        return result
    except Exception:
        return "تعذر حساب التعبير المدخل."

calculator_tool = Tool(
    name="Calculator",
    func=calculator_tool_func,
    description="Performs simple arithmetic calculations."
)

def weather_tool_func(location: str):
    return f"حالة الطقس في {location}: مشمس، درجة الحرارة 25 درجة مئوية."

weather_tool = Tool(
    name="Weather",
    func=weather_tool_func,
    description="Provides weather information for a given location."
)

serpapi = SerpAPIWrapper(serpapi_api_key=serpapi_api_key)
SerpAPI_tool = Tool(
    name="WebSearch",
    func=serpapi.run,
    description="Searches the web for recent information."
)

# --- تعريف موديل الذكاء الاصطناعي والوكيل ---

tools = [SerpAPI_tool, calculator_tool, weather_tool]
llm = ChatGroq(model="deepseek-r1-distill-llama-70b", groq_api_key=groq_api_key)
llm_with_tools = llm.bind_tools(tools)

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

def assistant(state: AgentState):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

builder = StateGraph(AgentState)
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")

# 👇 وكيل باسم ninu يمكن استيراده من أي ملف آخر
ninu = builder.compile()

# --- دالة تشغيل مستقلة (اختياري) ---

def run_ninu(query: str):
    conversation = []

    intro_prompt = """
You are a general AI assistant with access to the following tools:
1. WebSearch: to search for general or recent information from the internet.
2. Calculator: to perform arithmetic calculations.
3. Weather: to provide weather information.
Think step by step, and decide which tools to use to answer the query.
Respond with:
FINAL ANSWER: [your final answer]
Only include what the user asked for in the final answer.
"""
    conversation.append(HumanMessage(content=intro_prompt))
    conversation.append(HumanMessage(content=query))

    response = ninu.invoke({"messages": conversation})
    return response["messages"][-1].content