"""
记忆系统：短期记忆、长期记忆与 RAG 检索（向量数据库）。
"""
import asyncio
import hashlib
import shutil
import time
import uuid
from typing import List, Dict, Optional, Tuple, Callable, Awaitable
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from collections import deque
from .vector_store import VectorStore, VectorMemory, ChromaVectorStore, JSONVectorStore
from src.utils.logger import setup_logger

logger = setup_logger('MemorySystem')


@dataclass
class MemoryItem:
    """记忆项（用于短期记忆）。"""
    content: str
    timestamp: datetime
    user_id: str
    importance: float = 0.5
    metadata: Dict = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        """转换为字典。"""
        return {
            'content': self.content,
            'timestamp': self.timestamp.isoformat(),
            'user_id': self.user_id,
            'importance': self.importance,
            'metadata': self.metadata
        }


class ShortTermMemory:
    """短期记忆（滑动窗口）。"""
    
    def __init__(self, max_size: int = 20, max_age_minutes: int = 30):
        self.max_size = max_size
        self.max_age = timedelta(minutes=max_age_minutes)
        self.memories: Dict[str, deque] = {}  # user_id -> deque of MemoryItem
    
    def add(self, user_id: str, content: str, metadata: Optional[Dict] = None):
        """添加短期记忆。"""
        if user_id not in self.memories:
            self.memories[user_id] = deque(maxlen=self.max_size)
        
        memory = MemoryItem(
            content=content,
            timestamp=datetime.now(),
            user_id=user_id,
            metadata=metadata or {}
        )
        
        self.memories[user_id].append(memory)
    
    def get(self, user_id: str, limit: Optional[int] = None) -> List[MemoryItem]:
        """获取短期记忆。"""
        if user_id not in self.memories:
            return []
        
        # 过滤过期记忆
        now = datetime.now()
        valid_memories = [
            m for m in self.memories[user_id]
            if now - m.timestamp <= self.max_age
        ]
        
        if limit:
            valid_memories = valid_memories[-limit:]
        
        return valid_memories
    
    def clear(self, user_id: str):
        """清除用户短期记忆。"""
        if user_id in self.memories:
            self.memories.pop(user_id, None)


class MemorySystem:
    """记忆系统：整合短期记忆与长期记忆。"""
    
    def __init__(
        self,
        storage_path: Path,
        embed_func: Optional[callable] = None,
        importance_evaluator: Optional[Callable[[str, Optional[Dict]], Awaitable[float]]] = None,
        importance_threshold: float = 0.6,
        use_vector_db: bool = True,
        use_query_embedding: bool = False,
    ):
        self.short_term = ShortTermMemory()
        self.embed_func = embed_func
        self.importance_evaluator = importance_evaluator
        self.importance_threshold = importance_threshold
        # Only embed retrieval queries when explicitly enabled.
        self.use_query_embedding = use_query_embedding
        
        # 初始化向量存储
        if use_vector_db:
            chroma_path = storage_path.parent / "chroma_db"
            chroma_store = self._init_chroma_store(chroma_path)
            if chroma_store is not None:
                self.vector_store = chroma_store
                logger.info("Using Chroma vector store")
            else:
                self.vector_store = JSONVectorStore(storage_path)
        else:
            # 使用 JSON 存储（向后兼容）
            self.vector_store = JSONVectorStore(storage_path)
            logger.info("使用 JSON 存储")

    @staticmethod
    def _is_chroma_table_conflict(error: Exception) -> bool:
        msg = str(error).lower()
        return "table embeddings already exists" in msg

    @staticmethod
    def _is_chroma_trigram_error(error: Exception) -> bool:
        msg = str(error).lower()
        return "no such tokenizer: trigram" in msg

    def _init_chroma_store(self, chroma_path: Path) -> Optional[VectorStore]:
        """初始化 Chroma，遇到已知 sqlite schema 冲突时尝试修复。"""
        try:
            return ChromaVectorStore(chroma_path)
        except Exception as error:
            if self._is_chroma_trigram_error(error):
                logger.warning(
                    "Chroma 初始化失败，降级为 JSON 存储: sqlite 缺少 trigram tokenizer。"
                    "请在运行环境升级 sqlite 或安装 pysqlite3-binary。"
                )
                return None

            if not self._is_chroma_table_conflict(error):
                logger.warning(f"Chroma 初始化失败，降级为 JSON 存储: {error}")
                return None

            # 先做一次短暂重试，处理并发启动时的瞬时冲突。
            logger.warning(f"Chroma 初始化出现 schema 冲突，正在重试: {error}")
            time.sleep(0.2)
            try:
                return ChromaVectorStore(chroma_path)
            except Exception as retry_error:
                if not self._is_chroma_table_conflict(retry_error):
                    logger.warning(f"Chroma 重试失败，降级为 JSON 存储: {retry_error}")
                    return None

            backup_name = (
                f"{chroma_path.name}_backup_conflict_"
                f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
            )
            backup_path = chroma_path.parent / backup_name

            try:
                if chroma_path.exists():
                    shutil.move(str(chroma_path), str(backup_path))
                chroma_path.mkdir(parents=True, exist_ok=True)
                repaired = ChromaVectorStore(chroma_path)
                logger.warning(
                    f"检测到 Chroma 元数据库冲突，已重建目录并保留备份: {backup_path}"
                )
                return repaired
            except Exception as repair_error:
                logger.warning(f"Chroma 修复失败，降级为 JSON 存储: {repair_error}")
                return None
    
    @staticmethod
    def _normalize_embedding(values: List[float], dim: int = 1024) -> List[float]:
        if not values:
            return [0.0] * dim

        normalized = [float(v) for v in values[:dim]]
        if len(normalized) < dim:
            normalized.extend([0.0] * (dim - len(normalized)))
        return normalized

    @staticmethod
    def _local_embedding(text: str, dim: int = 1024) -> List[float]:
        if not text:
            return [0.0] * dim

        vec = [0.0] * dim
        encoded = text.encode("utf-8", errors="ignore")
        if not encoded:
            return vec

        for idx, byte in enumerate(encoded):
            bucket = idx % dim
            vec[bucket] += (byte / 255.0)

        digest = hashlib.sha256(encoded).digest()
        for idx, byte in enumerate(digest):
            bucket = idx % dim
            vec[bucket] += ((byte / 255.0) - 0.5) * 0.1

        return vec

    async def _build_embedding(self, text: str) -> List[float]:
        if self.embed_func:
            try:
                embedding = await self.embed_func(text)
                if embedding:
                    return [float(v) for v in list(embedding)]
            except Exception as e:
                logger.warning(f"embedding generation failed: {e}")

        return self._local_embedding(text)

    async def _add_vector_memory(
        self,
        memory_id: str,
        user_id: str,
        content: str,
        embedding: List[float],
        importance: float,
        metadata: Optional[Dict] = None,
    ) -> bool:
        if await self.vector_store.add(
            id=memory_id,
            user_id=user_id,
            content=content,
            embedding=embedding,
            importance=importance,
            metadata=metadata,
        ):
            return True

        # Chroma collection may have a fixed historical embedding dimension.
        candidate_dims = []
        base_len = len(embedding or [])
        for dim in [base_len, 1024, 1536, 768, 384, 3072]:
            if dim and dim > 0 and dim not in candidate_dims:
                candidate_dims.append(dim)

        for dim in candidate_dims:
            retry_embedding = self._normalize_embedding(list(embedding or []), dim=dim)
            ok = await self.vector_store.add(
                id=memory_id,
                user_id=user_id,
                content=content,
                embedding=retry_embedding,
                importance=importance,
                metadata=metadata,
            )
            if ok:
                return True

        return False

    @staticmethod
    def _normalize_importance(importance: float) -> float:
        try:
            value = float(importance)
        except (TypeError, ValueError):
            value = 0.5
        return max(0.0, min(1.0, value))

    async def add_message(
        self,
        user_id: str,
        role: str,
        content: str,
        metadata: Optional[Dict] = None
    ):
        """向短期记忆添加单条消息（不做长期记忆评分）。"""
        self.short_term.add(user_id, content, metadata)

    async def add_qa_pair(
        self,
        user_id: str,
        question: str,
        answer: str,
        metadata: Optional[Dict] = None,
    ) -> Optional[VectorMemory]:
        """
        添加最新问答对，并仅对该问答对做模型重要性评估。
        """
        user_meta = {"role": "user"}
        assistant_meta = {"role": "assistant"}
        if isinstance(metadata, dict):
            user_meta.update(metadata)
            assistant_meta.update(metadata)

        self.short_term.add(user_id, question, user_meta)
        self.short_term.add(user_id, answer, assistant_meta)

        qa_content = f"用户: {question}\n助手: {answer}"
        qa_metadata = dict(metadata or {})
        qa_metadata.update({"type": "qa_pair", "question": question, "answer": answer})

        importance = await self._evaluate_importance(qa_content, qa_metadata)
        if importance < self.importance_threshold:
            return None

        embedding = await self._build_embedding(qa_content)
        memory_id = str(uuid.uuid4())
        ok = await self._add_vector_memory(
            memory_id=memory_id,
            user_id=user_id,
            content=qa_content,
            embedding=embedding,
            importance=importance,
            metadata=qa_metadata,
        )
        if not ok:
            return None

        return await self.get_long_term(user_id, memory_id)
    
    async def _evaluate_importance(self, content: str, metadata: Optional[Dict]) -> float:
        """评估记忆重要性。"""
        if not content or not content.strip():
            return 0.0

        if self.importance_evaluator:
            try:
                score = await self.importance_evaluator(content, metadata)
                return self._normalize_importance(score)
            except Exception as e:
                logger.warning(f"importance evaluation failed, fallback to neutral score: {e}")

        # 当模型评估不可用时，使用中性分数作为兜底。
        return 0.5
    
    async def get_context(
        self,
        user_id: str,
        query: Optional[str] = None,
        max_short_term: int = 10,
        max_long_term: int = 5
    ) -> Tuple[List[MemoryItem], List[VectorMemory]]:
        """获取上下文（短期 + 长期记忆）。"""
        # 获取短期记忆
        short_term_memories = self.short_term.get(user_id, limit=max_short_term)
        
        # 获取相关长期记忆
        long_term_memories = []
        
        if query and self.use_query_embedding:
            try:
                # 使用向量检索
                query_embedding = await self._build_embedding(query)
                if query_embedding:
                    long_term_memories = await self.vector_store.search(
                        user_id=user_id,
                        query_embedding=query_embedding,
                        limit=max_long_term
                    )
            except Exception as e:
                logger.warning(f"向量检索失败，改用重要性检索: {e}")

        if query and not long_term_memories:
            query_lower = query.lower()
            try:
                candidates = await self.vector_store.get_all(user_id)
                matches = [m for m in candidates if query_lower in m.content.lower()]
                matches.sort(key=lambda m: (m.importance, m.timestamp), reverse=True)
                long_term_memories = matches[:max_long_term]
            except Exception:
                pass
        
        # 濡傛灉鍚戦噺妫€绱㈠け璐ユ垨娌℃湁缁撴灉锛屼娇鐢ㄩ噸瑕佹€ф绱?
        if not long_term_memories:
            long_term_memories = await self.vector_store.get_by_importance(
                user_id=user_id,
                limit=max_long_term
            )
        
        # 更新长期记忆访问记录
        for memory in long_term_memories:
            await self.vector_store.update_access(memory.id)
        
        return short_term_memories, long_term_memories
    
    def format_context(
        self,
        short_term: List[MemoryItem],
        long_term: List[VectorMemory]
    ) -> str:
        """格式化上下文为文本。"""
        context = ""
        
        if long_term:
            context += "## 相关历史记忆\n"
            for i, memory in enumerate(long_term, 1):
                context += f"{i}. {memory.content}\n"
            context += "\n"
        
        if short_term:
            context += "## 最近对话\n"
            for memory in short_term:
                context += f"- {memory.content}\n"
        
        return context

    async def list_long_term(
        self, user_id: str, limit: int = 20
    ) -> List[VectorMemory]:
        memories = await self.vector_store.get_all(user_id)
        memories.sort(key=lambda m: m.timestamp, reverse=True)
        if limit > 0:
            return memories[:limit]
        return memories

    async def get_long_term(self, user_id: str, memory_id: str) -> Optional[VectorMemory]:
        memories = await self.vector_store.get_all(user_id)
        for memory in memories:
            if memory.id == memory_id:
                return memory
        return None

    async def add_long_term(
        self,
        user_id: str,
        content: str,
        importance: float = 0.8,
        metadata: Optional[Dict] = None,
    ) -> Optional[VectorMemory]:
        memory_id = str(uuid.uuid4())
        importance = self._normalize_importance(importance)
        embedding = await self._build_embedding(content)

        ok = await self._add_vector_memory(
            memory_id=memory_id,
            user_id=user_id,
            content=content,
            embedding=embedding,
            importance=importance,
            metadata=metadata or {},
        )
        if not ok:
            return None

        return await self.get_long_term(user_id, memory_id)

    async def search_long_term(
        self, user_id: str, query: str, limit: int = 10
    ) -> List[VectorMemory]:
        if not query:
            return []

        if self.use_query_embedding:
            query_embedding = await self._build_embedding(query)
            results = await self.vector_store.search(
                user_id=user_id,
                query_embedding=query_embedding,
                limit=limit,
                min_importance=0.0,
            )
            if results:
                return results

        all_memories = await self.vector_store.get_all(user_id)
        query_lower = query.lower()
        matched = [m for m in all_memories if query_lower in m.content.lower()]
        matched.sort(key=lambda m: (m.importance, m.timestamp), reverse=True)
        return matched[:limit]

    async def update_long_term(
        self,
        user_id: str,
        memory_id: str,
        content: Optional[str] = None,
        importance: Optional[float] = None,
        metadata: Optional[Dict] = None,
    ) -> Optional[VectorMemory]:
        original = await self.get_long_term(user_id, memory_id)
        if not original:
            return None

        new_content = content if content is not None else original.content
        new_importance = (
            self._normalize_importance(importance)
            if importance is not None
            else original.importance
        )
        new_metadata = dict(original.metadata or {})
        if metadata is not None:
            new_metadata = metadata

        if content is not None:
            new_embedding = await self._build_embedding(new_content)
        else:
            new_embedding = self._normalize_embedding(list(original.embedding or []))

        deleted = await self.vector_store.delete(memory_id)
        if not deleted:
            return None

        added = await self._add_vector_memory(
            memory_id=memory_id,
            user_id=user_id,
            content=new_content,
            embedding=new_embedding,
            importance=new_importance,
            metadata=new_metadata,
        )
        if not added:
            return None

        return await self.get_long_term(user_id, memory_id)

    async def delete_long_term(self, user_id: str, memory_id: str) -> bool:
        memory = await self.get_long_term(user_id, memory_id)
        if not memory:
            return False
        return await self.vector_store.delete(memory_id)
    
    def clear_short_term(self, user_id: str):
        """清除短期记忆。"""
        self.short_term.clear(user_id)
    
    async def clear_long_term(self, user_id: str):
        """清除长期记忆。"""
        await self.vector_store.clear_user(user_id)
    
    async def close(self):
        """关闭记忆系统。"""
        await self.vector_store.close()
