import streamlit as st
import os
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_community.llms import Tongyi

# ---- 環(huán)境變量輸入?yún)^(qū)（可根據(jù)需求修改為硬編碼） ----
st.set_page_config(page_title="Qwen 聊天機器人", page_icon="??", layout="centered")

with st.sidebar:
    st.title("配置參數(shù)")
    dashscope_api_key = st.text_input("DASHSCOPE_API_KEY", type="password", value=os.environ.get("DASHSCOPE_API_KEY", ""))
    langsmith_api_key = st.text_input("LANGSMITH_API_KEY（可選）", type="password", value=os.environ.get("LANGSMITH_API_KEY", ""))
    langsmith_project = st.text_input("LANGSMITH_PROJECT（可選）", value=os.environ.get("LANGSMITH_PROJECT", ""))
    langsmith_endpoint = st.text_input("LANGSMITH_ENDPOINT（可選）", value=os.environ.get("LANGSMITH_ENDPOINT", "https://api.smith.langchain.com"))
    st.markdown("> *如無需 LangSmith，可留空*")
    st.markdown("---")
    st.markdown("by [你的名字/團隊]")

# 設(shè)置環(huán)境變量
if dashscope_api_key:
    os.environ["DASHSCOPE_API_KEY"] = dashscope_api_key
if langsmith_api_key:
    os.environ["LANGSMITH_API_KEY"] = langsmith_api_key
    os.environ["LANGSMITH_TRACING"] = "true"
if langsmith_project:
    os.environ["LANGSMITH_PROJECT"] = langsmith_project
if langsmith_endpoint:
    os.environ["LANGSMITH_ENDPOINT"] = langsmith_endpoint

# ---- 對話狀態(tài)定義 ----
class State(TypedDict):
    messages: Annotated[list, add_messages]

# ---- 狀態(tài)機組裝 ----
@st.cache_resource(show_spinner="正在加載 Qwen 模型…")
def get_graph():
    graph_builder = StateGraph(State)
    llm = Tongyi(
        model="qwen-max",
        streaming=True,
        temperature=0.7,
        top_p=0.8,
        api_key=os.environ.get("DASHSCOPE_API_KEY")
    )
    def chatbot(state: State):
        return {"messages": [llm.invoke(state["messages"])]}
    graph_builder.add_node("chatbot", chatbot)
    graph_builder.add_edge(START, "chatbot")
    graph_builder.add_edge("chatbot", END)
    return graph_builder.compile()

graph = get_graph()

# ---- Streamlit 聊天 session 狀態(tài)維護 ----
if "history" not in st.session_state:
    st.session_state.history = []

st.title("?? Qwen 聊天機器人（LangGraph 版）")

with st.form(key="chat-form", clear_on_submit=True):
    user_input = st.text_area("請輸入你的問題：", value="", placeholder="和機器人聊點什么吧…", height=80)
    submitted = st.form_submit_button("發(fā)送")

def stream_graph_updates(user_input: str, history: list):
    yield from graph.stream({"messages": history + [{"role": "user", "content": user_input}]})

# 聊天歷史展示
for msg in st.session_state.history:
    if msg["role"] == "user":
        st.chat_message("user").write(msg["content"])
    else:
        st.chat_message("assistant").write(msg["content"])

# 新問題處理
if submitted and user_input.strip():
    st.session_state.history.append({"role": "user", "content": user_input})
    response_placeholder = st.empty()
    full_response = ""
    # 流式顯示
    with st.spinner("Assistant 正在思考..."):
        for event in stream_graph_updates(user_input, st.session_state.history[:-1]):
            for value in event.values():
                resp = value["messages"][-1]
                # 僅展示assistant內(nèi)容
                # if resp["role"] == "assistant":
                #     full_response = resp
                response_placeholder.chat_message("assistant").write(resp)
    # 更新歷史
    if full_response:
        st.session_state.history.append({"role": "assistant", "content": full_response})

# 清空歷史按鈕
st.sidebar.button("清空對話歷史", on_click=lambda: st.session_state.pop("history", None))

st.info("提示：本機器人支持多輪上下文記憶，支持流式輸出。輸入 'quit'/'exit'/'q' 可結(jié)束對話（但不會關(guān)閉頁面）。")