#!/usr/bin/env python3
"""
mem0 Client for OpenClaw (v1.0 兼容)
"""

import os
import logging
from typing import List, Dict, Optional

# 设置环境变量（在导入 mem0 之前）
os.environ['OPENAI_API_BASE'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
os.environ['OPENAI_API_KEY'] = 'sk-c1715ee0479841399fd359c574647648'

try:
    from mem0 import Memory
    from mem0.configs.base import MemoryConfig, VectorStoreConfig, LlmConfig
except ImportError as e:
    print(f"⚠️  mem0ai 导入失败：{e}")
    Memory = None

logger = logging.getLogger(__name__)

class Mem0Client:
    def __init__(self):
        self.local_memory = None
        self.init_memory()
    
    def init_memory(self):
        """初始化 mem0"""
        if Memory is None:
            logger.warning("mem0ai 未安装")
            return
        
        try:
            config = MemoryConfig(
                vector_store=VectorStoreConfig(
                    provider="qdrant",
                    config={
                        "host": os.getenv('MEM0_QDRANT_HOST', 'localhost'),
                        "port": int(os.getenv('MEM0_QDRANT_PORT', '6333')),
                        "collection_name": "mem0_local",
                        "on_disk": True
                    }
                ),
                llm=LlmConfig(
                    provider="openai",
                    config={"model": "qwen-plus"}
                )
            )
            
            self.local_memory = Memory(config=config)
            logger.info("✅ 本地记忆初始化成功")
        except Exception as e:
            logger.error(f"❌ 初始化失败：{e}")
            self.local_memory = None
    
    def add(self, messages: List[Dict], user_id: str) -> Optional[Dict]:
        """添加记忆"""
        if self.local_memory is None:
            return {"error": "mem0 not initialized"}
        
        try:
            result = self.local_memory.add(messages, user_id=user_id)
            return {"success": True}
        except Exception as e:
            return {"error": str(e)}
    
    def search(self, query: str, user_id: str, limit: int = 5) -> List[Dict]:
        """搜索记忆"""
        if self.local_memory is None:
            return []
        
        try:
            return self.local_memory.search(query, user_id=user_id, limit=limit)
        except Exception:
            return []
    
    def get_all(self, user_id: str) -> List[Dict]:
        """获取所有记忆"""
        if self.local_memory is None:
            return []
        
        try:
            return self.local_memory.get_all(user_id=user_id)
        except Exception:
            return []
    
    def delete(self, memory_id: str, user_id: str) -> bool:
        """删除记忆"""
        if self.local_memory is None:
            return False
        
        try:
            self.local_memory.delete(memory_id, user_id=user_id)
            return True
        except Exception:
            return False
    
    def get_status(self) -> Dict:
        """获取状态"""
        return {
            "initialized": self.local_memory is not None,
            "qdrant": f"{os.getenv('MEM0_QDRANT_HOST', 'localhost')}:{os.getenv('MEM0_QDRANT_PORT', '6333')}"
        }