import os

import dotenv
from langchain.memory import ConversationBufferMemory, ConversationTokenBufferMemory, ConversationSummaryMemory
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.messages import HumanMessage, AIMessage, trim_messages

memory = ConversationBufferMemory()
memory.chat_memory.add_user_message("我叫小明")
memory.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory)
print(memory.memory_key)
print(memory.load_memory_variables({}))
print('---------')

memory2 = ConversationBufferMemory(memory_key="memory2")
memory2.chat_memory.add_user_message("我叫小明")
memory2.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory2)
print(memory2.memory_key)
print(memory2.load_memory_variables({}))
print('---------')

memory3 = ConversationBufferMemory(return_messages=True)
memory3.chat_memory.add_user_message("我叫小明")
memory3.chat_memory.add_ai_message("噢，好的，你叫小明")
print(memory3)
print(memory3.memory_key)
print(memory3.load_memory_variables({}))
print('---------')
dotenv.load_dotenv()

## 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("OLLAMA_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("OLLAMA_BASE_URL")

# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
llm = ChatOpenAI(model="gemma4:e2b")

memory_test = ConversationSummaryMemory(llm=llm)
memory_test.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
memory_test.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手，可以帮你解答任何问题。"})
print(memory_test.load_memory_variables({}))
# {'history': '\n\nThe human introduces themselves as xiaoming. The AI confirms the name and responds. The human then asks who the AI is. The AI introduces itself as an all-powerful AI chat assistant who can answer any questions.'}
print('---------')

# 这是openai专用的OpenAI套件
# memory_token = ConversationTokenBufferMemory(llm=llm,max_token_limit=7)
# memory_token.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
# memory_token.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手，可以帮你解答任何问题。"})
# print(memory_token.load_memory_variables({}))

# k=1 代表只記住「最新的一問一答」（1輪）
memory_window = ConversationBufferWindowMemory(k=1)
memory_window.save_context({"input": "我叫小明"}, {"output": "噢，好的，你叫小明"})
memory_window.save_context({"input": "那么你是谁呢"}, {"output": "我是一个无所不能的AI聊天助手。"})
print(memory_window.load_memory_variables({}))
# 輸出結果只會保留最後一輪，因為 k=1
print('---------')

messages = [
    HumanMessage("我叫小明"),
    AIMessage("噢，好的，你叫小明"),
    HumanMessage("那么你是谁呢"),
    AIMessage("我是一个无所不能的AI聊天助手。")
]

# 現代化做法：直接裁剪訊息列表
# 這裡我們用最簡單的 "只保留最後 2 條訊息" 策略
trimmed_messages = trim_messages(
    messages,
    max_tokens=2, 
    token_counter=len, # 關鍵：我們直接告訴它「算訊息條數」當作 token，避開底層 tiktoken 崩潰
    strategy="last",
)

print(trimmed_messages)