"""
Chroma向量数据库实现
"""
import uuid
from typing import List, Dict, Optional
from pathlib import Path
from datetime import datetime
import chromadb
from chromadb.config import Settings
from .base import VectorStore, VectorMemory
from src.utils.logger import setup_logger

logger = setup_logger('ChromaStore')


class ChromaVectorStore(VectorStore):
    """Chroma向量数据库实现"""
    
    def __init__(self, persist_directory: Path):
        """初始化Chroma客户端"""
        self.persist_directory = persist_directory
        self.persist_directory.mkdir(parents=True, exist_ok=True)
        
        # 创建Chroma客户端
        self.client = chromadb.PersistentClient(
            path=str(persist_directory),
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
        
        # 获取或创建集合
        self.collection = self.client.get_or_create_collection(
            name="long_term_memory",
            metadata={"description": "Long-term memory storage"}
        )
        
        logger.info(f"✅ Chroma向量数据库初始化: {persist_directory}")

    @staticmethod
    def _to_list(value):
        """将可能的 numpy 数组等序列转为 Python list。"""
        if value is None:
            return []
        if hasattr(value, "tolist"):
            return value.tolist()
        if isinstance(value, list):
            return value
        return list(value)
    
    async def add(
        self,
        id: str,
        user_id: str,
        content: str,
        embedding: List[float],
        importance: float,
        metadata: Optional[Dict] = None
    ) -> bool:
        """添加记忆"""
        try:
            if not embedding:
                logger.warning(f"跳过写入长期记忆（embedding为空）: {id}")
                return False
            
            now = datetime.now()
            
            # 准备元数据
            meta = {
                'user_id': user_id,
                'importance': importance,
                'timestamp': now.isoformat(),
                'access_count': 0,
                'last_access': '',
            }
            
            if metadata:
                # 只保存可序列化的元数据
                for key, value in metadata.items():
                    if isinstance(value, (str, int, float, bool)):
                        meta[f'meta_{key}'] = value
            
            # 添加到Chroma
            self.collection.add(
                ids=[id],
                embeddings=[embedding],
                documents=[content],
                metadatas=[meta]
            )
            
            logger.debug(f"添加记忆: {id} (用户: {user_id})")
            return True
            
        except Exception as e:
            logger.error(f"添加记忆失败: {e}")
            return False
    
    async def search(
        self,
        user_id: str,
        query_embedding: List[float],
        limit: int = 5,
        min_importance: float = 0.3
    ) -> List[VectorMemory]:
        """搜索相似记忆"""
        try:
            # 查询Chroma
            results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=limit * 2,  # 多查询一些，用于过滤
                where={"user_id": user_id}
            )
            
            ids = self._to_list(results.get('ids'))
            if len(ids) == 0:
                return []
            first_row_ids = self._to_list(ids[0])
            if len(first_row_ids) == 0:
                return []

            metadatas = self._to_list(results.get('metadatas'))
            documents = self._to_list(results.get('documents'))
            embeddings = self._to_list(results.get('embeddings'))
            first_row_metas = self._to_list(metadatas[0]) if len(metadatas) > 0 else []
            first_row_docs = self._to_list(documents[0]) if len(documents) > 0 else []
            first_row_embeddings = self._to_list(embeddings[0]) if len(embeddings) > 0 else []
            
            # 转换结果
            memories = []
            for i in range(len(first_row_ids)):
                if i >= len(first_row_metas) or i >= len(first_row_docs):
                    continue

                meta = first_row_metas[i]
                
                # 过滤重要性
                if meta['importance'] < min_importance:
                    continue
                
                # 提取自定义元数据
                custom_meta = {}
                for key, value in meta.items():
                    if key.startswith('meta_'):
                        custom_meta[key[5:]] = value
                
                memory = VectorMemory(
                    id=first_row_ids[i],
                    user_id=meta['user_id'],
                    content=first_row_docs[i],
                    embedding=self._to_list(first_row_embeddings[i]) if i < len(first_row_embeddings) else [],
                    importance=meta['importance'],
                    timestamp=datetime.fromisoformat(meta['timestamp']),
                    metadata=custom_meta,
                    access_count=meta.get('access_count', 0),
                    last_access=datetime.fromisoformat(meta['last_access']) if meta.get('last_access') else None
                )
                memories.append(memory)
                
                if len(memories) >= limit:
                    break
            
            return memories
            
        except Exception as e:
            logger.error(f"搜索记忆失败: {e}")
            return []
    
    async def get_by_importance(
        self,
        user_id: str,
        limit: int = 5,
        min_importance: float = 0.3
    ) -> List[VectorMemory]:
        """按重要性获取记忆"""
        try:
            # 获取所有用户记忆
            results = self.collection.get(
                where={"user_id": user_id},
                include=['embeddings', 'documents', 'metadatas']
            )
            
            ids = self._to_list(results.get('ids'))
            if len(ids) == 0:
                return []

            metadatas = self._to_list(results.get('metadatas'))
            documents = self._to_list(results.get('documents'))
            embeddings = self._to_list(results.get('embeddings'))
            
            # 转换并排序
            memories = []
            for i in range(len(ids)):
                if i >= len(metadatas) or i >= len(documents):
                    continue
                meta = metadatas[i]
                
                if meta['importance'] < min_importance:
                    continue
                
                custom_meta = {}
                for key, value in meta.items():
                    if key.startswith('meta_'):
                        custom_meta[key[5:]] = value
                
                memory = VectorMemory(
                    id=ids[i],
                    user_id=meta['user_id'],
                    content=documents[i],
                    embedding=self._to_list(embeddings[i]) if i < len(embeddings) else [],
                    importance=meta['importance'],
                    timestamp=datetime.fromisoformat(meta['timestamp']),
                    metadata=custom_meta,
                    access_count=meta.get('access_count', 0),
                    last_access=datetime.fromisoformat(meta['last_access']) if meta.get('last_access') else None
                )
                memories.append(memory)
            
            # 按重要性和时间排序
            memories.sort(key=lambda m: (m.importance, m.timestamp), reverse=True)
            return memories[:limit]
            
        except Exception as e:
            logger.error(f"获取记忆失败: {e}")
            return []
    
    async def update_access(self, memory_id: str) -> bool:
        """更新访问记录"""
        try:
            # 获取当前记忆
            result = self.collection.get(
                ids=[memory_id],
                include=['metadatas']
            )
            
            ids = self._to_list(result.get('ids'))
            if len(ids) == 0:
                return False
            
            metadatas = self._to_list(result.get('metadatas'))
            if len(metadatas) == 0:
                return False
            meta = metadatas[0]
            meta['access_count'] = meta.get('access_count', 0) + 1
            meta['last_access'] = datetime.now().isoformat()
            
            # 更新元数据
            self.collection.update(
                ids=[memory_id],
                metadatas=[meta]
            )
            
            return True
            
        except Exception as e:
            logger.error(f"更新访问记录失败: {e}")
            return False
    
    async def delete(self, memory_id: str) -> bool:
        """删除记忆"""
        try:
            self.collection.delete(ids=[memory_id])
            return True
        except Exception as e:
            logger.error(f"删除记忆失败: {e}")
            return False
    
    async def get_all(self, user_id: str) -> List[VectorMemory]:
        """获取用户所有记忆"""
        try:
            results = self.collection.get(
                where={"user_id": user_id},
                include=['embeddings', 'documents', 'metadatas']
            )
            
            ids = self._to_list(results.get('ids'))
            if len(ids) == 0:
                return []

            metadatas = self._to_list(results.get('metadatas'))
            documents = self._to_list(results.get('documents'))
            embeddings = self._to_list(results.get('embeddings'))
            
            memories = []
            for i in range(len(ids)):
                if i >= len(metadatas) or i >= len(documents):
                    continue
                meta = metadatas[i]
                
                custom_meta = {}
                for key, value in meta.items():
                    if key.startswith('meta_'):
                        custom_meta[key[5:]] = value
                
                memory = VectorMemory(
                    id=ids[i],
                    user_id=meta['user_id'],
                    content=documents[i],
                    embedding=self._to_list(embeddings[i]) if i < len(embeddings) else [],
                    importance=meta['importance'],
                    timestamp=datetime.fromisoformat(meta['timestamp']),
                    metadata=custom_meta,
                    access_count=meta.get('access_count', 0),
                    last_access=datetime.fromisoformat(meta['last_access']) if meta.get('last_access') else None
                )
                memories.append(memory)
            
            return memories
            
        except Exception as e:
            logger.error(f"获取所有记忆失败: {e}")
            return []
    
    async def clear_user(self, user_id: str) -> bool:
        """清除用户所有记忆"""
        try:
            # 获取用户所有记忆ID
            results = self.collection.get(
                where={"user_id": user_id}
            )
            
            ids = self._to_list(results.get('ids'))
            if len(ids) > 0:
                self.collection.delete(ids=ids)
            
            logger.info(f"清除用户记忆: {user_id}")
            return True
            
        except Exception as e:
            logger.error(f"清除用户记忆失败: {e}")
            return False
    
    async def close(self):
        """关闭连接"""
        # Chroma会自动持久化，无需特殊关闭操作
        logger.info("Chroma向量数据库已关闭")
