#!/usr/bin/env python3
"""
Mem0 Python 集成脚本
被 Node.js 插件调用，执行实际的记忆操作
"""

import sys
import json
import os
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['OPENAI_API_KEY'] = os.getenv('MEM0_DASHSCOPE_API_KEY', 'sk-c1715ee0479841399fd359c574647648')

sys.path.insert(0, os.path.dirname(__file__))

from mem0_client import mem0_client


async def main():
    if len(sys.argv) < 2:
        print(json.dumps({"error": "No action specified"}))
        return

    action = sys.argv[1]
    data = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}

    try:
        if action == 'init':
            # 初始化 mem0
            await mem0_client.start()
            print(json.dumps({
                "status": "initialized",
                "qdrant": f"{mem0_client.config['qdrant']['host']}:{mem0_client.config['qdrant']['port']}"
            }))

        elif action == 'search':
            # 检索记忆
            memories = await mem0_client.pre_hook_search(
                query=data.get('query', ''),
                user_id=data.get('user_id', 'default'),
                agent_id=data.get('agent_id', 'general')
            )
            print(json.dumps({
                "memories": memories,
                "count": len(memories)
            }))

        elif action == 'add':
            # 添加记忆（异步，不等待）
            mem0_client.post_hook_add(
                user_message=data.get('user_message', ''),
                assistant_message=data.get('assistant_message', ''),
                user_id=data.get('user_id', 'default'),
                agent_id=data.get('agent_id', 'general')
            )
            print(json.dumps({"status": "queued"}))

        else:
            print(json.dumps({"error": f"Unknown action: {action}"}))

    except Exception as e:
        print(json.dumps({"error": str(e)}))


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