# vector_memory.py - 向量存储引擎
# BGE-M3 + Chroma + SQLite 架构

import chromadb
from chromadb.config import Settings
from openai import OpenAI
import sqlite3
import json
from datetime import datetime


class VectorMemorySystem:
    def __init__(self, persist_dir="./data", api_key: str = None):
        """初始化向量记忆系统"""
        
        # 1. 初始化硅基流动客户端
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.siliconflow.cn/v1"
        )
        
        # 2. 初始化 Chroma 向量库
        self.chroma = chromadb.Client(Settings(
            persist_directory=persist_dir,
            anonymized_telemetry=False
        ))
        self.collection = self.chroma.get_or_create_collection(
            name="openclaw_memory",
            metadata={"description": "OpenClaw long-term memory"}
        )
        
        # 3. 初始化 SQLite（用于持久化）
        self.db_path = f"{persist_dir}/memory.db"
        self._init_sqlite()
    
    def _init_sqlite(self):
        """初始化 SQLite 数据库"""
        self.conn = sqlite3.connect(self.db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS memories (
                id TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                metadata TEXT,
                importance INTEGER DEFAULT 3,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        self.conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_importance ON memories(importance)
        """)
        self.conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_created_at ON memories(created_at)
        """)
        self.conn.commit()
    
    def _get_embedding(self, text: str) -> list:
        """调用 BGE-M3 获取向量"""
        response = self.client.embeddings.create(
            model="BAAI/bge-m3",
            input=text
        )
        return response.data[0].embedding
    
    def add_memory(self, content: str, metadata: dict = None, importance: int = 3):
        """添加记忆（同时写入向量库 + SQLite）"""
        import uuid
        memory_id = str(uuid.uuid4())
        
        # 1. 生成向量并存储
        embedding = self._get_embedding(content)
        self.collection.add(
            ids=[memory_id],
            embeddings=[embedding],
            documents=[content],
            metadatas=[metadata or {}]
        )
        
        # 2. 写入 SQLite 持久化
        self.conn.execute(
            """INSERT INTO memories (id, content, metadata, importance) 
               VALUES (?, ?, ?, ?)""",
            (memory_id, content, json.dumps(metadata), importance)
        )
        self.conn.commit()
        
        return memory_id
    
    def search(self, query: str, top_k: int = 5) -> list:
        """语义搜索"""
        # 1. 查询向量
        query_embedding = self._get_embedding(query)
        
        # 2. 向量相似度搜索
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k
        )
        
        # 3. 格式化返回
        memories = []
        for i, doc in enumerate(results['documents'][0]):
            memories.append({
                'id': results['ids'][0][i],
                'content': doc,
                'distance': results['distances'][0][i],
                'metadata': results['metadatas'][0][i]
            })
        
        return memories
    
    def hybrid_search(self, query: str, keyword: str = None, top_k: int = 5):
        """混合搜索：语义 + 关键词"""
        # 1. 向量搜索
        vector_results = self.search(query, top_k * 2)
        
        # 2. 关键词过滤（可选）
        if keyword:
            vector_results = [
                r for r in vector_results 
                if keyword in r['content']
            ]
        
        return vector_results[:top_k]


if __name__ == "__main__":
    import os
    
    api_key = os.getenv("SILICONFLOW_API_KEY")
    if not api_key:
        print("请设置 SILICONFLOW_API_KEY 环境变量")
        exit(1)
    
    vm = VectorMemorySystem(persist_dir="./data/memory", api_key=api_key)
    
    # 测试添加
    memory_id = vm.add_memory(
        content="2026-03-21: 部署了向量记忆系统",
        metadata={"tags": ["系统部署"]},
        importance=4
    )
    print(f"添加记忆成功: {memory_id}")
    
    # 测试搜索
    results = vm.search("记忆系统")
    for r in results:
        print(f"- {r['content'][:50]}... (相似度: {1-r['distance']:.2%})")
