from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langgraph.checkpoint.memory import InMemorySaver  
from tool import DuckDuckGoSearchRun,web_search_tool,latest_news_tool, get_weather_tool as weather_info_tool, hub_stats_tool

from langchain_huggingface import HuggingFaceEndpoint



model = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen2.5-7B-Instruct",
    task="text-generation",
    max_new_tokens=256
)

tools=[weather_info_tool,web_search_tool,weather_info_tool,hub_stats_tool]
model_with_tool = model.bind_tools(tools)

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


def assistant(state: AgentState):
    return {
        "messages": [model_with_tool.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')

checkpointer = InMemorySaver()

alfred = builder.compile(checkpointer=checkpointer)
thread_config = {"configurable": {"thread_id": "1"}}



def run_agent(question: str):

    response = alfred.invoke(
        {"messages":[HumanMessage(content=question)]},
        {"configurable":{"thread_id":"evaluation"}}
    )

    answer = response["messages"][-1].content

    return answer