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 langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace

from tools import search_tool, weather_info_tool, hub_stats_tool, guest_info_tool
from retriever import docs

from langchain_ollama import ChatOllama
import os
HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN is None:
    raise RuntimeError("⚠️ 没有找到 HF_TOKEN，请先在 Spaces 的 Variables and secrets 添加。")

# 生成聊天界面，包括工具
#本地加载ollama模型
# llm = ChatOllama(model="gpt-oss:20b", request_timeout=120.0)
# chat_with_tools = llm.bind_tools(tools)
#使用远程推理服务器
tools = [guest_info_tool, search_tool, weather_info_tool, hub_stats_tool]
# 初始化 Hugging Face 模型
# 生成聊天界面，包括工具
llm = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen-7B-Instruct",
    huggingfacehub_api_token=HF_TOKEN,
)

chat = ChatHuggingFace(llm=llm, verbose=True)
tools = [guest_info_tool]
chat_with_tools = chat.bind_tools(tools)

# 生成 AgentState 和 Agent 图
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

def assistant(state: AgentState):
    return {
        "messages": [chat_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 节点
    # 否则直接响应
    tools_condition,
)
builder.add_edge("tools", "assistant")
alfred = builder.compile()

#示例1
# response = alfred.invoke({"messages": "Tell me about 'Lady Ada Lovelace' and translate output to chinese."})

# print("🎩 Alfred's Response:")
# print(response['messages'][-1].content)

#示例2
# response = alfred.invoke({"messages": "What's the weather like in Tokyo tonight? Will it be suitable for our fireworks display?and translate output to chinese."})

# print("🎩 Alfred's Response:")
# print(response['messages'][-1].content)

#示例 3：给 AI 研究者留下深刻印象
# response = alfred.invoke({"messages": "One of our guests is from Qwen. What can you tell me about their most popular model?请用中文回答"})

# print("🎩 Alfred's Response:")
# print(response['messages'][-1].content)

#示例 4：组合多工具应用
# response = alfred.invoke({"messages":"我需要与“尼古拉·特斯拉博士”讨论最近在无线能源方面的进展。你能帮我为这次对话做准备吗？"})

# print("🎩 Alfred's Response:")
# print(response['messages'][-1].content)

#高级功能：对话记忆

# 首次交互
response = alfred.invoke({"messages": [HumanMessage(content="Tell me about 'Lady Ada Lovelace'. What's her background and how is she related to me?请用中文回答")]})


print("🎩 Alfred's Response:")
print(response['messages'][-1].content)
print("以下是第二次对话内容")

# 二次交互（引用首次内容）
response = alfred.invoke({"messages": response["messages"] + [HumanMessage(content="What projects is she currently working on?请用中文回答")]})

print("🎩 Alfred's Response:")
print(response['messages'][-1].content)