import streamlit as st
import os
from dotenv import load_dotenv
from DeepseekAi import get_response_stream  # 导入流式函数

# 第一步：加载.env配置（必须放在最前面）
load_dotenv()
API_KEY = os.getenv('api_key')
BASE_URL = os.getenv('base_url')

# 检查配置完整性
if not API_KEY or not BASE_URL:
    st.error("❌ 配置缺失：请检查.env文件中的 api_key 和 base_url 是否填写正确！")
    st.stop()

# 页面标题
st.title("智聊机器人")

# 初始化会话状态
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "你好，我是智聊机器人，有什么可以帮助你的么？"}]

# 显示历史对话
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 处理用户输入
user_input = st.chat_input("请输入你的问题：")

if user_input:
    # 1. 添加用户消息到会话状态并显示
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.write(user_input)

    # 2. 流式输出AI回复
    full_response = ""  # 存储完整回复
    with st.chat_message("assistant"):
        # 初始化思考状态（赋值给status变量，用于后续修改）
        with st.status(label="🤔 正在思考中...", expanded=True) as status:
            response_placeholder = st.empty()  # 只创建一次占位符

            # 逐段获取流式内容
            for chunk in get_response_stream(st.session_state.messages, api_key=API_KEY, base_url=BASE_URL):
                full_response += chunk
                # 实时更新占位符，加光标模拟打字效果
                response_placeholder.markdown(full_response + "|")

            # 核心修改：思考完成后更新状态文本
            status.update(label="✅ 思考结束", state="complete", expanded=False)

        # 3. 最终展示完整回复（保留你原有逻辑，未修改DOM渲染）
        st.markdown(full_response)

    # 4. 把完整回复存入会话状态
    st.session_state.messages.append({"role": "assistant", "content": full_response})