# -*- coding: utf-8 -*-
"""
记忆持久化（JSONL）

每条记忆结构：
{
    "id": "mem_xxx",
    "title": "短标题",
    "content": "完整正文",
    "category": "session|preference|fact",
    "source_session": "session-name",
    "created_at": "2026-07-25T10:00:00",
    "tags": ["tag1", "tag2"]
}
"""
import json
import os
import time
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from ..security import get_default_filter


@dataclass
class MemoryEntry:
    """单条记忆"""
    id: str
    title: str
    content: str
    category: str = "session"
    source_session: str = ""
    created_at: str = ""
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "MemoryEntry":
        return cls(
            id=d.get("id") or f"mem_{uuid.uuid4().hex[:8]}",
            title=d.get("title") or "",
            content=d.get("content") or "",
            category=d.get("category") or "session",
            source_session=d.get("source_session") or "",
            created_at=d.get("created_at") or "",
            tags=list(d.get("tags") or []),
        )

    def to_index_entry(self) -> Dict[str, Any]:
        """转成 BM25Engine 可识别的索引项"""
        return {
            "id": self.id,
            "title": self.title,
            "content": self.content,
        }


class MemoryStore:
    """记忆持久化存储（JSONL）。

    使用方式：
        store = MemoryStore(Path.home() / ".abyss" / "memory.jsonl")
        store.add(MemoryEntry(id="...", title="...", content="..."))
        all_entries = store.list_all()
        store.save()  # 写入磁盘
    """

    def __init__(self, path: Path, max_entries: int = 1000):
        self._path = Path(path)
        self._max_entries = max_entries
        self._entries: List[MemoryEntry] = []
        self._loaded = False

    def _ensure_loaded(self) -> None:
        if self._loaded:
            return
        self._loaded = True
        self._entries = self._load_from_disk()

    def _load_from_disk(self) -> List[MemoryEntry]:
        if not self._path.exists():
            return []
        entries: List[MemoryEntry] = []
        try:
            with self._path.open("r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        d = json.loads(line)
                        if isinstance(d, dict):
                            entries.append(MemoryEntry.from_dict(d))
                    except (json.JSONDecodeError, ValueError):
                        continue
        except OSError:
            return []
        return entries

    def _atomic_save(self) -> None:
        """原子写入：写 .tmp 再 rename，避免写入中断损坏。"""
        self._path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self._path.with_suffix(self._path.suffix + ".tmp")
        with tmp.open("w", encoding="utf-8") as f:
            for entry in self._entries:
                f.write(json.dumps(entry.to_dict(), ensure_ascii=False) + "\n")
        # Windows 上 replace 也能原子覆盖
        os.replace(str(tmp), str(self._path))

    def add(self, entry: MemoryEntry, *, persist: bool = True) -> None:
        """添加一条记忆。超过 max_entries 时按 FIFO 淘汰最旧的。

        写入前会对 title/content 做敏感信息脱敏（借鉴 zhikuncode-main
        SensitiveDataFilter），避免密钥被持久化到 ~/.abyss/memory.jsonl。
        """
        self._ensure_loaded()
        if not entry.id:
            entry.id = f"mem_{uuid.uuid4().hex[:8]}"
        if not entry.created_at:
            entry.created_at = datetime.now().isoformat(timespec="seconds")
        # 脱敏：避免密钥污染跨会话记忆库
        redactor = get_default_filter()
        entry.title = redactor.redact(entry.title)
        entry.content = redactor.redact(entry.content)
        # tags 也脱敏（用户可能在 tag 中放密钥）
        entry.tags = [redactor.redact(t) for t in entry.tags]
        self._entries.append(entry)
        # FIFO 淘汰
        while len(self._entries) > self._max_entries:
            self._entries.pop(0)
        if persist:
            self._atomic_save()

    def list_all(self) -> List[MemoryEntry]:
        """返回所有记忆条目（按创建时间升序）"""
        self._ensure_loaded()
        return list(self._entries)

    def get(self, mem_id: str) -> Optional[MemoryEntry]:
        """按 id 获取"""
        self._ensure_loaded()
        for e in self._entries:
            if e.id == mem_id:
                return e
        return None

    def remove(self, mem_id: str, *, persist: bool = True) -> bool:
        """按 id 删除。返回是否删除成功。"""
        self._ensure_loaded()
        for i, e in enumerate(self._entries):
            if e.id == mem_id:
                self._entries.pop(i)
                if persist:
                    self._atomic_save()
                return True
        return False

    def clear(self, *, persist: bool = True) -> int:
        """清空所有记忆，返回被清除的数量。"""
        self._ensure_loaded()
        n = len(self._entries)
        self._entries.clear()
        if persist:
            self._atomic_save()
        return n

    def save(self) -> None:
        """显式保存到磁盘"""
        self._ensure_loaded()
        self._atomic_save()

    def search(self, query: str, top_k: int = 5) -> List[MemoryEntry]:
        """简单关键词搜索（向后兼容用，BM25 检索请用 BM25Engine）"""
        self._ensure_loaded()
        if not query:
            return []
        q_lower = query.lower()
        scored: List[tuple] = []
        for i, e in enumerate(self._entries):
            text = (e.title + " " + e.content).lower()
            score = text.count(q_lower)
            if score > 0:
                scored.append((score, i, e))
        scored.sort(key=lambda x: (-x[0], x[1]))
        return [e for _, _, e in scored[:top_k]]

    @property
    def count(self) -> int:
        self._ensure_loaded()
        return len(self._entries)

    @property
    def path(self) -> Path:
        return self._path
