#!/usr/bin/env python3
# /root/.openclaw/workspace/skills/mem0-integration/test_integration.py

import asyncio
import logging
import sys
sys.path.insert(0, '/root/.openclaw/workspace/skills/mem0-integration')

from mem0_client import mem0_client
from openclaw_interceptor import intercept_before_llm, intercept_after_response

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

async def mock_llm(system_prompt, user_message):
    """模拟 LLM 调用"""
    return f"这是模拟回复：{user_message[:20]}..."

async def test_full_flow():
    """测试完整对话流程"""
    print("=" * 60)
    print("🧪 测试 mem0 集成架构")
    print("=" * 60)
    
    context = {
        'user_id': '5237946060',
        'agent_id': 'general'
    }
    
    # ========== 测试 1: Pre-Hook 检索 ==========
    print("\n1️⃣ 测试 Pre-Hook 检索...")
    query = "我平时喜欢用什么时区？"
    memory_prompt = await intercept_before_llm(query, context)
    
    if memory_prompt:
        print(f"✅ 检索到记忆:\n{memory_prompt}")
    else:
        print("⚠️  未检索到记忆（正常，首次对话）")
    
    # ========== 测试 2: 完整对话流程 ==========
    print("\n2️⃣ 测试完整对话流程...")
    user_message = "我平时喜欢使用 UTC 时区，请用简体中文和我交流"
    
    print(f"用户：{user_message}")
    response = await mock_llm("system", user_message)
    print(f"助手：{response}")
    
    # ========== 测试 3: Post-Hook 异步写入 ==========
    print("\n3️⃣ 测试 Post-Hook 异步写入...")
    await intercept_after_response(user_message, response, context)
    print(f"✅ 对话已提交到异步队列")
    print(f"   队列大小：{len(mem0_client.async_queue.queue) if mem0_client.async_queue else 0}")
    
    # ========== 等待异步写入完成 ==========
    print("\n4️⃣ 等待异步写入 (5 秒)...")
    await asyncio.sleep(5)
    
    # ========== 测试 4: 验证记忆已存储 ==========
    print("\n5️⃣ 验证记忆已存储...")
    memories = await mem0_client.pre_hook_search("时区", **context)
    print(f"✅ 检索到 {len(memories)} 条记忆")
    for i, mem in enumerate(memories, 1):
        print(f"   {i}. {mem.get('memory', 'N/A')[:100]}")
    
    # ========== 状态报告 ==========
    print("\n" + "=" * 60)
    print("📊 系统状态:")
    status = mem0_client.get_status()
    for key, value in status.items():
        print(f"  {key}: {value}")
    print("=" * 60)
    print("✅ 测试完成")

if __name__ == '__main__':
    asyncio.run(test_full_flow())