from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import ChatOpenAI
import gradio as gr;

from env_util import DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL

# 0、llm~~
llm = ChatOpenAI(
    model = "qwen-plus",
    base_url=DASHSCOPE_BASE_URL,
    api_key=DASHSCOPE_API_KEY,
    temperature=0.8,
);

# ===============================================================================================

# 1、定义专门做聊天的提示词模板
prompt = ChatPromptTemplate.from_messages([
    ('system', "{system_message}"), # 系统提示词
    MessagesPlaceholder(variable_name='chat_history', optional=True), #消息占位符
    ('human', '{input}') #用户提示词，input用户传的问题
])

chain = prompt | llm;

# ===============================================================================================
# 2、存储聊天记录： 存的谁？存的第10行的内容（存到哪里？内存、关系型数据库或者redis数据库）

store = {}

def get_session_history(session_id: str):
    """从关系型数据库的历史消息列表中 返回当前会话 的所有历史消息"""
    #  SQLChatMessageHistory是langchain提供的
    return SQLChatMessageHistory(
        session_id=session_id,
        # 这里url换为自己的数据库即可
        connection_string='sqlite:///chat_history.db',
    )



# ===============================================================================================
# 3、创建带历史记录功能的处理链，帮我自动存储历史记录

chain_with_message_history = RunnableWithMessageHistory(
    chain, # 基础执行链
    get_session_history, # 指定工厂函数，返回指定session_id的聊天记录
    input_messages_key='input', # 指定用户输入的消息的key
    history_messages_key='chat_history', # 历史消息记录的key
)

# ===============================================================================================
#4、剪辑和摘要历史上下文消息
# 比如：保留最近的前2条消息（随意指定条数），把之前的所有消息形成摘要
# 定义summarize_messages函数，把当前用户的这次输入传进来
def summarize_messages(current_input):
    """剪辑和摘要上下午，历史记录"""
    # 从current_input取出session_id
    session_id = current_input['config']["configurable"]["session_id"]

    if not session_id:
        raise ValueError("必须通过config参数提供session_id")

    # 获取当前会话ID的所有历史聊天记录
    chat_history = get_session_history(session_id)
    # 从历史聊天记录取出message这个聊天列表（是个数组，里面是一条一条消息，包括AiMessage，toolmessage，humanmessage，systemmessage。。）
    stored_messages = chat_history.messages

    # 如果stored_messages长度太短，只有两条以内聊天记录，不需要摘要
    if len(stored_messages) <= 2:  # 保留最近2条消息的阈值
        return {"original_messages": stored_messages, "summary": None}  # 不满足摘要条件时返回原始消息-保留的最后2条消息

    # 超过2条再剪辑
    # 剪辑消息列表
    last_two_messages = stored_messages[-2:]  # 保留的最后2条消息
    messages_to_summarize = stored_messages[:-2]  # 需要进行摘要的消息列表（最后2条之前的记录）

    # 构建摘要：就是让大模型去帮你构建摘要，帮你做这件事情，在这里需要再调用大模型
    # 所以先构建提示词模版
    summarization_prompt = ChatPromptTemplate.from_messages([
        ("system", "请将以下对话历史压缩为一条保留关键信息的摘要消息。"),
        ("placeholder", "{chat_history}"),
        ("human", "请生成包含上述对话核心内容的摘要，保留重要事实和决策。")
    ])
    # 组成链调用大模型
    summarization_chain = summarization_prompt | llm
    # 大模型会帮你生成摘要(AIMessage) ——》需要生成摘要的是messages_to_summarize，last_two_messages不需要生成摘要
    summary_message = summarization_chain.invoke({'chat_history': messages_to_summarize})

    # 返回结构化结果（不调用chat_history.clear()）
    return {
        "original_messages": last_two_messages,  # 保留的最后2条消息
        "summary": summary_message  # 需要进行摘要的消息列表（最后2条之前的记录）
    }


# 5、最终的链 （需要LCEL完成）
# RunnablePassthrough 默认会将输入数据原样传递到下游，通过管道传给下一个组件RunnablePassthrough
# 在第二个RunnablePassthrough收到了前面返回得两个key，分别是original_messages，summary
# 然后再第二个RunnablePassthrough分别把内容注入到chat_history，以及修改的提示词模板的system里，system也换成了占位符，这里相当于赋值
final_chain = RunnablePassthrough.assign(messages_summarized=summarize_messages) | RunnablePassthrough.assign(
    input = lambda x: x['input'],
    chat_history = lambda x: x['messages_summarized']['original_messages'],
    system_message = lambda x: f"你是一个乐于助人的助手。尽你所能回答所有问题。摘要：{x ['messages_summarized']['summary']}" if x['messages_summarized'].get('summary') else "",
) | chain_with_message_history;

# ===============================================================================================
#6、 启动Web页面，实现内容交互
#  web界面中的核心函数
# 第一个参数：聊天记录，第二个参数：用户输入的消息
# 这个功能只是把用户输入的消息添加到历史记录中，并没有调用大模型
def add_message(chat_history, user_message):
    if user_message:
        # 角色：用户角色
        chat_history.append({"role": "user", "content": user_message})
    return chat_history, ''


# 调用大模型
def execute_chain(chat_history):
    # 拿最后一条记录，最后一条就是我们刚才在文本框输入的内容
    input = chat_history[-1]
    # 调用模型
    result = final_chain.invoke({'input': input['content'], "config": {"configurable": {"session_id": "user123"}}},
                                config={"configurable": {"session_id": "user123"}})
    # 模型返回结果追加到历史记录
    chat_history.append({'role': 'assistant', 'content': result.content})
    return chat_history


# 1.开发一个聊天机器人的Web界面 gr.Blocks
#  title:标题，theme=gr.themes.Soft() 界面默认样式
with gr.Blocks(title='多模态聊天机器人', theme=gr.themes.Soft()) as block:
    # 2.加入聊天历史记录的组件 Chatbot
    chatbot = gr.Chatbot(height=500, label='聊天机器人')
    # 3. 行组件，
    with gr.Row():
        # gr.Row() 表示 “行”，gr.Column() 表示 “列”，结合 scale 参数就能实现你说的左右分布，且列的宽度比例是 4:1。我会用新手能听懂的方式讲清楚每一部分的含义和效果。
        # 行组件中左面：文字输入的区域
        with gr.Column(scale=4):
            user_input = gr.Textbox(placeholder='请给机器人发送消息...', label='文字输入', max_lines=5)

            submit_btn = gr.Button('发送', variant="primary")
        # 行组件中右面：语音输入区域，type='filepath'录制音频的文件，后缀：wav，你用mp3也可以
        with gr.Column(scale=1):
            audio_input = gr.Audio(sources=['microphone'], label='语音输入', type='filepath', format='wav')

    # 添加提交事件(要调用的函数，函数传入的组件列表，函数输出的组件列表)
    chat_msg = user_input.submit(add_message, [chatbot, user_input], [chatbot, user_input])
    # 再去调用大模型
    chat_msg.then(execute_chain, chatbot, chatbot)

if __name__ == '__main__':
    block.launch()  # 界面跑起来