# app.py — Entry point for HF Spaces
import os
import gradio as gr
from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph.message import add_messages
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_groq import ChatGroq
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.tools import tool
import requests

os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")

# ─── LLM ─────────────────────────────────
llm = ChatGroq(
    model="llama-3.1-8b-instant",
    max_tokens=300,
    temperature=0
)

# ─── Tools ───────────────────────────────
@tool
def search(query: str) -> str:
    """Search the web for current information."""
    return DuckDuckGoSearchRun().run(query)

@tool
def calculator(first_num: float, second_num: float, operation: str) -> dict:
    """Perform basic arithmetic. Operations: add, sub, mul, div"""
    ops = {"add": first_num + second_num, "sub": first_num - second_num,
           "mul": first_num * second_num}
    if operation == "div":
        return {"result": "Division by zero" if second_num == 0 else first_num / second_num}
    return {"result": ops.get(operation, f"Unknown operation: {operation}")}

@tool
def get_stock_price(symbol: str) -> dict:
    """Fetch latest stock price for a symbol like AAPL or TSLA."""
    url = f"https://www.alphavantage.co/query?function=GLOBAL_QUOTE&symbol={symbol}&apikey={os.getenv('STOCKS_API_KEY')}"
    return requests.get(url).json()

tools = [search, calculator, get_stock_price]
llm_with_tools = llm.bind_tools(tools)

# ─── State ───────────────────────────────
class ChatState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

# ─── Graph ───────────────────────────────
def chat_node(state: ChatState):
    """LLM node that may answer or request a tool call."""
    messages = state['messages']
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

graph = StateGraph(ChatState)
graph.add_node("chat_node", chat_node)
graph.add_node("tools", ToolNode(tools))
graph.add_edge(START, "chat_node")
graph.add_conditional_edges("chat_node", tools_condition)
graph.add_edge("tools", "chat_node")
agent = graph.compile()

# ─── Gradio UI ───────────────────────────
def respond(message, history):
    # Convert Gradio history → LangChain messages
    messages = []
    for user_msg, bot_msg in history:
        messages.append(HumanMessage(content=user_msg))
        if bot_msg:
            messages.append(AIMessage(content=bot_msg))
    messages.append(HumanMessage(content=message))

    result = agent.invoke({"messages": messages})
    return result["messages"][-1].content

demo = gr.ChatInterface(
    fn=respond,
    title="🔍 AI Research Agent",
    description="Ask me anything — I can search the web and do calculations!",
    examples=[
        "What is LangGraph?",
        "What's happening in AI news today?",
        "Calculate 128 multiplied by 37",
    ],
    cache_examples=False   # 🔥 ADD THIS
)

if __name__ == "__main__":
    demo.launch()