#!/usr/bin/env python3
"""
mem0 验证脚本 - 测试 DashScope 1024 维度配置
"""

import os
import sys
import asyncio

# 设置环境变量
os.environ['OPENAI_API_BASE'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
os.environ['OPENAI_BASE_URL'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
os.environ['MEM0_DASHSCOPE_API_KEY'] = 'sk-4111c9dba5334510968f9ae72728944e'
os.environ['OPENAI_API_KEY'] = 'sk-4111c9dba5334510968f9ae72728944e'

# 添加路径
sys.path.insert(0, '/root/.openclaw/workspace/skills/mem0-integration')

print("=" * 60)
print("🔍 mem0 验证测试 - DashScope 1024 维度配置")
print("=" * 60)

# 测试 1: 检查环境变量
print("\n[1/4] 检查环境变量...")
print(f"  OPENAI_BASE_URL: {os.environ.get('OPENAI_BASE_URL')}")
print(f"  OPENAI_API_KEY: {os.environ.get('OPENAI_API_KEY')[:10]}...")
print(f"  ✅ 环境变量已设置")

# 测试 2: 导入 mem0_client
print("\n[2/4] 导入 mem0_client 模块...")
try:
    from mem0_client import Mem0Client
    print("  ✅ 模块导入成功")
except Exception as e:
    print(f"  ❌ 模块导入失败：{e}")
    sys.exit(1)

# 测试 3: 初始化客户端
print("\n[3/4] 初始化 Mem0Client...")
try:
    client = Mem0Client()
    print(f"  ✅ 客户端初始化成功")
    print(f"     - local_memory: {client.local_memory is not None}")
    print(f"     - qdrant_host: {client.config['qdrant']['host']}")
    print(f"     - embedder_model: {client.config['embedder']['config'].get('model')}")
    print(f"     - embedding_dims: {client.config['embedder']['config'].get('dimensions')}")
except Exception as e:
    print(f"  ❌ 初始化失败：{e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# 测试 4: 测试向量生成（验证 1024 维度）
print("\n[4/4] 测试向量生成（验证 1024 维度）...")
async def test_embedding():
    try:
        # 使用 DashScope API 直接测试
        import httpx
        
        api_key = os.environ.get('OPENAI_API_KEY')
        api_base = os.environ.get('OPENAI_BASE_URL')
        
        payload = {
            "model": "text-embedding-v3",
            "input": "测试文本 - 验证 1024 维度配置",
            "dimensions": 1024
        }
        
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        
        async with httpx.AsyncClient(timeout=30.0) as http:
            response = await http.post(
                f"{api_base}/embeddings",
                json=payload,
                headers=headers
            )
        
        if response.status_code == 200:
            data = response.json()
            embedding = data['data'][0]['embedding']
            print(f"  ✅ 向量生成成功")
            print(f"     - 维度：{len(embedding)}")
            print(f"     - 状态码：{response.status_code}")
            
            if len(embedding) == 1024:
                print(f"  🎉 维度验证通过！(1024)")
                return True
            else:
                print(f"  ⚠️ 维度不匹配！期望 1024，实际 {len(embedding)}")
                return False
        else:
            print(f"  ❌ API 请求失败：{response.status_code}")
            print(f"     响应：{response.text[:200]}")
            return False
            
    except Exception as e:
        print(f"  ❌ 测试失败：{e}")
        import traceback
        traceback.print_exc()
        return False

# 运行异步测试
result = asyncio.run(test_embedding())

print("\n" + "=" * 60)
if result:
    print("✅ 所有测试通过！mem0 配置正确")
    sys.exit(0)
else:
    print("⚠️ 部分测试失败，请检查配置")
    sys.exit(1)