import os
from dotenv import load_dotenv
from langchain_deepseek import ChatDeepSeek

# 加载环境变量
load_dotenv(override=True)

# ---------- 基础配置 ----------
MODEL_NAME = "deepseek-v4-flash"
MAX_PAIRS_HISTORY = 10          # 保留最近 N 轮对话
EXIT_WORD = "quit"

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
DEEPSEEK_BASE_URL = os.getenv("DEEPSEEK_BASE_URL")

# 环境变量校验
if not DEEPSEEK_API_KEY:
    raise ValueError("请在 .env 文件中设置 DEEPSEEK_API_KEY")
if not DEEPSEEK_BASE_URL:
    raise ValueError("请在 .env 文件中设置 DEEPSEEK_BASE_URL")

# ---------- 初始化模型 ----------
model = ChatDeepSeek(
    model=MODEL_NAME,
    api_key=DEEPSEEK_API_KEY,
    api_base=DEEPSEEK_BASE_URL,
)

# ---------- 历史记忆管理函数 ----------
def keep_recent_messages(messages, max_pairs):
    """
    保留 system 消息 + 最近 max_pairs 轮对话（每轮包含 user 和 assistant 各一条）
    """
    # 至少要有 system 消息
    if len(messages) <= 1:
        return messages
    # 计算需要保留的消息数量：1 条 system + 2 * max_pairs 条问答
    keep_count = 1 + 2 * max_pairs
    if len(messages) <= keep_count:
        return messages
    # 保留第一条 system 消息 + 最后 2*max_pairs 条消息
    return [messages[0]] + messages[-2 * max_pairs:]

# ---------- 初始化消息列表 ----------
messages = [
    {
        "role": "system",
        "content": "你是小谷姐姐，是尚硅谷教育的数字员工，也是一名耐心、友好的AI助手，可以回答学员的问题"
    }
]

print(f"请输入具体的问题，当输入 {EXIT_WORD} 时结束对话。")

i = 1  # 对话轮数
while True:
    print("\n", "=" * 10, f"第 {i} 轮对话开始", "=" * 10, "\n")

    user_input = input("请输入：")

    # 退出判断
    if user_input.strip().lower() == EXIT_WORD:
        print("会话已结束，欢迎下次再来！")
        break

    # 将用户消息加入列表
    messages.append({"role": "user", "content": user_input})

    print("小谷姐姐：", end="", flush=True)

    # 获取优化后的历史消息（限制记忆长度）
    memory_messages = keep_recent_messages(messages, max_pairs=MAX_PAIRS_HISTORY)

    # 流式输出模型回复
    reply_content = ""
    try:
        for chunk in model.stream(memory_messages):
            if chunk.content:
                print(chunk.content, end="", flush=True)
                reply_content += chunk.content
    except Exception as e:
        print(f"\n[错误] 模型调用失败：{e}")
        # 将错误信息作为 AI 回复写入历史，避免对话中断
        reply_content = f"[系统提示：AI 回复出错] {e}"

    print("\n", "=" * 10, f"第 {i} 轮对话结束", "=" * 10, "\n")

    # 将 AI 回复加入消息列表
    messages.append({"role": "assistant", "content": reply_content})

    i += 1