"""
JSON文件存储实现（向后兼容）
"""
import json
import uuid
from typing import List, Dict, Optional
from pathlib import Path
from datetime import datetime
import numpy as np
from .base import VectorStore, VectorMemory
from src.utils.logger import setup_logger

logger = setup_logger('JSONStore')


class JSONVectorStore(VectorStore):
    """JSON文件存储实现（向后兼容旧版本）"""
    
    def __init__(self, storage_path: Path):
        """初始化JSON存储"""
        self.storage_path = storage_path
        self.memories: Dict[str, List[VectorMemory]] = {}  # user_id -> List[VectorMemory]
        self._load()
        logger.info(f"✅ JSON存储初始化: {storage_path}")
    
    def _load(self):
        """加载记忆"""
        if self.storage_path.exists():
            try:
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                for user_id, items in data.items():
                    self.memories[user_id] = []
                    for item in items:
                        # 兼容旧格式
                        if 'id' not in item:
                            item['id'] = str(uuid.uuid4())
                        
                        memory = VectorMemory.from_dict(item)
                        self.memories[user_id].append(memory)
                
                logger.info(f"加载了 {sum(len(v) for v in self.memories.values())} 条记忆")
            except Exception as e:
                logger.error(f"加载记忆失败: {e}")
                self.memories = {}
    
    def _save(self):
        """保存记忆"""
        try:
            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
            data = {
                user_id: [memory.to_dict() for memory in memories]
                for user_id, memories in self.memories.items()
            }
            
            with open(self.storage_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存记忆失败: {e}")
    
    async def add(
        self,
        id: str,
        user_id: str,
        content: str,
        embedding: List[float],
        importance: float,
        metadata: Optional[Dict] = None
    ) -> bool:
        """添加记忆"""
        try:
            memory = VectorMemory(
                id=id,
                user_id=user_id,
                content=content,
                embedding=embedding,
                importance=importance,
                timestamp=datetime.now(),
                metadata=metadata or {},
                access_count=0,
                last_access=None
            )
            
            if user_id not in self.memories:
                self.memories[user_id] = []
            
            self.memories[user_id].append(memory)
            self._save()
            
            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]:
        """搜索相似记忆"""
        if user_id not in self.memories:
            return []
        
        memories = self.memories[user_id]
        
        # 过滤重要性
        memories = [m for m in memories if m.importance >= min_importance]
        
        if not memories:
            return []
        
        # 使用向量相似度排序
        scored_memories = []
        for memory in memories:
            if memory.embedding:
                similarity = self._cosine_similarity(query_embedding, memory.embedding)
                scored_memories.append((similarity, memory))
        
        if not scored_memories:
            # 如果没有嵌入向量，按重要性排序
            return await self.get_by_importance(user_id, limit, min_importance)
        
        scored_memories.sort(reverse=True, key=lambda x: x[0])
        return [m for _, m in scored_memories[:limit]]
    
    async def get_by_importance(
        self,
        user_id: str,
        limit: int = 5,
        min_importance: float = 0.3
    ) -> List[VectorMemory]:
        """按重要性获取记忆"""
        if user_id not in self.memories:
            return []
        
        memories = [m for m in self.memories[user_id] if m.importance >= min_importance]
        memories.sort(key=lambda m: (m.importance, m.timestamp), reverse=True)
        return memories[:limit]
    
    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
        """计算余弦相似度"""
        vec1 = np.array(vec1)
        vec2 = np.array(vec2)
        return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    
    async def update_access(self, memory_id: str) -> bool:
        """更新访问记录"""
        try:
            for memories in self.memories.values():
                for memory in memories:
                    if memory.id == memory_id:
                        memory.access_count += 1
                        memory.last_access = datetime.now()
                        self._save()
                        return True
            return False
        except Exception as e:
            logger.error(f"更新访问记录失败: {e}")
            return False
    
    async def delete(self, memory_id: str) -> bool:
        """删除记忆"""
        try:
            for user_id, memories in self.memories.items():
                for i, memory in enumerate(memories):
                    if memory.id == memory_id:
                        del self.memories[user_id][i]
                        self._save()
                        return True
            return False
        except Exception as e:
            logger.error(f"删除记忆失败: {e}")
            return False
    
    async def get_all(self, user_id: str) -> List[VectorMemory]:
        """获取用户所有记忆"""
        return self.memories.get(user_id, [])
    
    async def clear_user(self, user_id: str) -> bool:
        """清除用户所有记忆"""
        try:
            if user_id in self.memories:
                del self.memories[user_id]
                self._save()
                logger.info(f"清除用户记忆: {user_id}")
            return True
        except Exception as e:
            logger.error(f"清除用户记忆失败: {e}")
            return False
    
    async def close(self):
        """关闭连接"""
        self._save()
        logger.info("JSON存储已关闭")
