"""Beaver 的精炼长期记忆存储层。

这个文件实现的是 Beaver curated memory 模型，目标不是
“把所有历史都存下来”，而是只保存跨会话仍然值得保留的稳定事实。

核心设计：
1. 只保留两个持久化记忆桶：
   - ``memory``: agent 自己对环境、项目、工具 quirks 的长期备注
   - ``user``: 对用户偏好、习惯、身份信息的长期理解
2. ``replace`` / ``remove`` 不使用 UUID，而是使用短语义片段做子串匹配。
   这是为了适配 LLM 更擅长“记住一句话片段”而不是“追踪一个随机 ID”的现实。
3. 写入前先做安全扫描，避免把 prompt injection / secrets exfiltration
   一类危险内容写入长期记忆，再在未来会话中反向污染 system prompt。
4. 写入协议严格遵守：
   - scan
   - lock
   - reload
   - validate
   - atomic write
5. 本文件维护两份状态：
   - live state: 当前内存中的真实条目，tool 写入后立刻变化
   - frozen snapshot: 会话开始时冻结的一份 prompt 注入快照

其中最重要的一点是：本会话中新增的记忆会立刻写盘，但不会反向修改本会话
已经冻结的 system prompt。这样可以保住 prefix cache，也避免“会话中途 prompt
变了导致行为抖动”的问题。
"""

from __future__ import annotations

import os
import re
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Any

try:
    import fcntl
except ImportError:  # pragma: no cover - Windows fallback
    fcntl = None

try:
    import msvcrt
except ImportError:  # pragma: no cover - Unix platforms
    msvcrt = None

ENTRY_DELIMITER = "\n§\n"
DEFAULT_MEMORY_FILENAME = "MEMORY.md"
DEFAULT_USER_FILENAME = "USER.md"

_MEMORY_THREAT_PATTERNS: list[tuple[str, str]] = [
    (r"ignore\s+(previous|all|above|prior)\s+instructions", "prompt_injection"),
    (r"you\s+are\s+now\s+", "role_hijack"),
    (r"do\s+not\s+tell\s+the\s+user", "deception_hide"),
    (r"system\s+prompt\s+override", "sys_prompt_override"),
    (r"disregard\s+(your|all|any)\s+(instructions|rules|guidelines)", "disregard_rules"),
    (r"act\s+as\s+(if|though)\s+you\s+(have\s+no|don't\s+have)\s+(restrictions|limits|rules)", "bypass_restrictions"),
    (r"curl\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", "exfil_curl"),
    (r"wget\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", "exfil_wget"),
    (r"cat\s+[^\n]*(\.env|credentials|\.netrc|\.pgpass|\.npmrc|\.pypirc)", "read_secrets"),
    (r"authorized_keys", "ssh_backdoor"),
    (r"\$HOME/\.ssh|\~/\.ssh", "ssh_access"),
    (r"\$HOME/\.beaver/\.env|\~/\.beaver/\.env", "beaver_env"),
]

_INVISIBLE_CHARS = {
    "\u200b",
    "\u200c",
    "\u200d",
    "\u2060",
    "\ufeff",
    "\u202a",
    "\u202b",
    "\u202c",
    "\u202d",
    "\u202e",
}


def scan_memory_content(content: str) -> str | None:
    """扫描待写入内容，拦截明显危险的记忆条目。

    这里不是在做完备的安全审计，而是在做“进入长期记忆之前的最低限度闸门”。
    因为长期记忆会在未来会话中重新注入 system prompt，所以一旦把恶意文本写进去，
    风险远高于普通临时上下文。
    """

    for char in _INVISIBLE_CHARS:
        if char in content:
            return (
                f"Blocked: content contains invisible unicode character "
                f"U+{ord(char):04X}."
            )

    for pattern, pattern_id in _MEMORY_THREAT_PATTERNS:
        if re.search(pattern, content, re.IGNORECASE):
            return (
                f"Blocked: content matches threat pattern '{pattern_id}'. "
                "Memory entries are injected into future system prompts."
            )

    return None


class MemoryStore:
    """带容量上限的长期记忆存储。

    这个类负责：
    1. 从磁盘加载 `MEMORY.md` / `USER.md`
    2. 在 session 启动时冻结 prompt snapshot
    3. 为 `add / replace / remove` 提供安全写接口
    4. 维护 live state 与 frozen snapshot 的边界

    它不负责：
    1. 自动从对话里抽取要记住的内容
    2. session transcript 检索
    3. skills 的学习和发布
    """

    def __init__(
        self,
        root: str | Path,
        *,
        memory_char_limit: int = 2200,
        user_char_limit: int = 1375,
    ) -> None:
        self.root = Path(root)
        self.memory_char_limit = memory_char_limit
        self.user_char_limit = user_char_limit
        self.memory_entries: list[str] = []
        self.user_entries: list[str] = []
        self._system_prompt_snapshot: dict[str, str] = {"memory": "", "user": ""}

    def load_from_disk(self) -> None:
        """从磁盘加载 live state，并冻结当前 session 的 prompt snapshot。

        调用时机应该是“会话启动时”，而不是每次工具写入后。
        如果在每次写入后都重新 load 并更新 system prompt，就会破坏 frozen snapshot
        这个设计，导致本轮会话 prompt 前缀发生变化。
        """

        self.root.mkdir(parents=True, exist_ok=True)
        self.memory_entries = list(dict.fromkeys(self._read_file(self._path_for("memory"))))
        self.user_entries = list(dict.fromkeys(self._read_file(self._path_for("user"))))
        self._system_prompt_snapshot = {
            "memory": self._render_block("memory", self.memory_entries),
            "user": self._render_block("user", self.user_entries),
        }

    @contextmanager
    def _file_lock(self, path: Path):
        """对目标记忆文件加排他锁。

        锁文件使用 sibling `.lock` 文件，而不是直接锁业务文件本身。
        原因是业务文件使用的是“临时文件写入 + os.replace 原子替换”，如果直接锁目标
        文件，替换时会让锁语义和文件句柄关系变得更脆弱。
        """

        lock_path = path.with_suffix(path.suffix + ".lock")
        lock_path.parent.mkdir(parents=True, exist_ok=True)

        if fcntl is None and msvcrt is None:
            yield
            return

        if msvcrt and (not lock_path.exists() or lock_path.stat().st_size == 0):
            lock_path.write_text(" ", encoding="utf-8")

        fd = open(lock_path, "r+" if msvcrt else "a+", encoding="utf-8")
        try:
            if fcntl is not None:
                fcntl.flock(fd, fcntl.LOCK_EX)
            elif msvcrt is not None:  # pragma: no cover - Windows fallback
                fd.seek(0)
                msvcrt.locking(fd.fileno(), msvcrt.LK_LOCK, 1)
            yield
        finally:
            if fcntl is not None:
                fcntl.flock(fd, fcntl.LOCK_UN)
            elif msvcrt is not None:  # pragma: no cover - Windows fallback
                try:
                    fd.seek(0)
                    msvcrt.locking(fd.fileno(), msvcrt.LK_UNLCK, 1)
                except OSError:
                    pass
            fd.close()

    def _path_for(self, target: str) -> Path:
        """根据目标桶返回实际文件路径。"""
        if target == "user":
            return self.root / DEFAULT_USER_FILENAME
        return self.root / DEFAULT_MEMORY_FILENAME

    def _entries_for(self, target: str) -> list[str]:
        """读取某个目标桶当前的 live entries。"""
        if target == "user":
            return self.user_entries
        return self.memory_entries

    def _set_entries(self, target: str, entries: list[str]) -> None:
        """更新某个目标桶在内存中的 live entries。"""
        if target == "user":
            self.user_entries = entries
        else:
            self.memory_entries = entries

    def _char_limit(self, target: str) -> int:
        """返回目标桶的字符预算。

        这里使用字符数而不是 token 数，是因为字符预算更稳定，也不依赖具体模型。
        """
        return self.user_char_limit if target == "user" else self.memory_char_limit

    def _char_count(self, target: str) -> int:
        """返回目标桶当前 live state 的字符占用。"""
        entries = self._entries_for(target)
        return len(ENTRY_DELIMITER.join(entries)) if entries else 0

    def _reload_target(self, target: str) -> None:
        """在持锁状态下重新从磁盘读取目标桶。

        这是并发安全协议里最关键的一步之一。
        必须在拿到锁之后 reload，才能确保当前进程不会覆盖掉其他并发会话刚刚写入
        的最新内容。
        """
        fresh = list(dict.fromkeys(self._read_file(self._path_for(target))))
        self._set_entries(target, fresh)

    def save_to_disk(self, target: str) -> None:
        """把当前 live entries 持久化到磁盘。"""
        self.root.mkdir(parents=True, exist_ok=True)
        self._write_file(self._path_for(target), self._entries_for(target))

    def add(self, target: str, content: str) -> dict[str, Any]:
        """追加一条新的长期记忆。

        规则：
        1. 空内容拒绝
        2. 安全扫描不通过拒绝
        3. 精确重复拒绝
        4. 超出字符预算拒绝
        5. 否则追加并立即写盘
        """

        content = content.strip()
        if not content:
            return {"success": False, "error": "Content cannot be empty."}

        scan_error = scan_memory_content(content)
        if scan_error:
            return {"success": False, "error": scan_error}

        with self._file_lock(self._path_for(target)):
            self._reload_target(target)
            entries = self._entries_for(target)
            if content in entries:
                return self._success_response(target, "Entry already exists (skipped duplicate).")

            new_entries = entries + [content]
            new_total = len(ENTRY_DELIMITER.join(new_entries))
            limit = self._char_limit(target)
            if new_total > limit:
                current = self._char_count(target)
                return {
                    "success": False,
                    "error": (
                        f"Memory at {current:,}/{limit:,} chars. "
                        f"Adding this entry ({len(content)} chars) would exceed the limit."
                    ),
                    "current_entries": list(entries),
                    "usage": f"{current:,}/{limit:,}",
                }

            entries.append(content)
            self._set_entries(target, entries)
            self.save_to_disk(target)

        return self._success_response(target, "Entry added.")

    def replace(self, target: str, old_text: str, new_content: str) -> dict[str, Any]:
        """用新的内容替换一条已有记忆。

        这里按 `old_text in entry` 做子串匹配，而不是要求调用方提供完整条目或 UUID。
        如果命中多条且它们内容不同，会要求调用方给出更精确的片段，避免误替换。
        """

        old_text = old_text.strip()
        new_content = new_content.strip()
        if not old_text:
            return {"success": False, "error": "old_text cannot be empty."}
        if not new_content:
            return {
                "success": False,
                "error": "new_content cannot be empty. Use remove to delete entries.",
            }

        scan_error = scan_memory_content(new_content)
        if scan_error:
            return {"success": False, "error": scan_error}

        with self._file_lock(self._path_for(target)):
            self._reload_target(target)
            entries = self._entries_for(target)
            matches = [(index, entry) for index, entry in enumerate(entries) if old_text in entry]
            if not matches:
                return {"success": False, "error": f"No entry matched '{old_text}'."}

            if len(matches) > 1:
                unique_texts = {entry for _, entry in matches}
                if len(unique_texts) > 1:
                    return {
                        "success": False,
                        "error": f"Multiple entries matched '{old_text}'. Be more specific.",
                        "matches": [
                            entry[:80] + ("..." if len(entry) > 80 else "")
                            for _, entry in matches
                        ],
                    }

            index = matches[0][0]
            candidate_entries = list(entries)
            candidate_entries[index] = new_content
            new_total = len(ENTRY_DELIMITER.join(candidate_entries))
            limit = self._char_limit(target)
            if new_total > limit:
                return {
                    "success": False,
                    "error": (
                        f"Replacement would put memory at {new_total:,}/{limit:,} chars. "
                        "Shorten the new content or remove other entries first."
                    ),
                }

            entries[index] = new_content
            self._set_entries(target, entries)
            self.save_to_disk(target)

        return self._success_response(target, "Entry replaced.")

    def remove(self, target: str, old_text: str) -> dict[str, Any]:
        """删除一条已有记忆。

        删除和替换共享同样的匹配策略：优先服务于 LLM 可操作性，而不是数据库式的强 ID。
        """

        old_text = old_text.strip()
        if not old_text:
            return {"success": False, "error": "old_text cannot be empty."}

        with self._file_lock(self._path_for(target)):
            self._reload_target(target)
            entries = self._entries_for(target)
            matches = [(index, entry) for index, entry in enumerate(entries) if old_text in entry]
            if not matches:
                return {"success": False, "error": f"No entry matched '{old_text}'."}

            if len(matches) > 1:
                unique_texts = {entry for _, entry in matches}
                if len(unique_texts) > 1:
                    return {
                        "success": False,
                        "error": f"Multiple entries matched '{old_text}'. Be more specific.",
                        "matches": [
                            entry[:80] + ("..." if len(entry) > 80 else "")
                            for _, entry in matches
                        ],
                    }

            entries.pop(matches[0][0])
            self._set_entries(target, entries)
            self.save_to_disk(target)

        return self._success_response(target, "Entry removed.")

    def format_for_system_prompt(self, target: str) -> str | None:
        """返回 session 启动时冻结下来的 prompt block。

        这里明确返回的是 frozen snapshot，而不是 live state。
        所以如果 session 中途调用 `add()` 写入了新记忆，这里不会立刻变化。
        """

        block = self._system_prompt_snapshot.get(target, "")
        return block or None

    def _success_response(self, target: str, message: str | None = None) -> dict[str, Any]:
        """统一生成 memory tool 的成功响应。

        响应里返回 live entries 和占用信息，目的是让模型能“看到自己刚写进去什么”，
        即使 system prompt 仍然保持冻结不变。
        """
        current = self._char_count(target)
        limit = self._char_limit(target)
        percent = min(100, int((current / limit) * 100)) if limit > 0 else 0
        payload: dict[str, Any] = {
            "success": True,
            "target": target,
            "entries": list(self._entries_for(target)),
            "entry_count": len(self._entries_for(target)),
            "usage": f"{percent}% — {current:,}/{limit:,} chars",
        }
        if message:
            payload["message"] = message
        return payload

    def _render_block(self, target: str, entries: list[str]) -> str:
        """把条目渲染成适合注入 system prompt 的块。"""
        if not entries:
            return ""

        current = len(ENTRY_DELIMITER.join(entries))
        limit = self._char_limit(target)
        percent = min(100, int((current / limit) * 100)) if limit > 0 else 0
        if target == "user":
            header = f"USER PROFILE (who the user is) [{percent}% — {current:,}/{limit:,} chars]"
        else:
            header = f"MEMORY (your personal notes) [{percent}% — {current:,}/{limit:,} chars]"
        separator = "═" * 46
        return f"{separator}\n{header}\n{separator}\n{ENTRY_DELIMITER.join(entries)}"

    @staticmethod
    def _read_file(path: Path) -> list[str]:
        """读取记忆文件并按 entry delimiter 拆分。

        这里不额外加读锁，因为写入采用的是原子替换：读者只会看到旧完整文件或新完整文件，
        不会看到半写入状态。
        """
        if not path.exists():
            return []
        try:
            raw = path.read_text(encoding="utf-8")
        except OSError:
            return []
        if not raw.strip():
            return []
        return [entry for entry in (item.strip() for item in raw.split(ENTRY_DELIMITER)) if entry]

    @staticmethod
    def _write_file(path: Path, entries: list[str]) -> None:
        """以原子方式写入记忆文件。

        这里不能直接 `open(path, "w")`，因为那会先截断原文件，再写新内容。
        如果恰好此时别的进程正在读，就可能读到空文件或半成品。

        正确方式是：
        1. 在同目录创建临时文件
        2. 写入并 fsync
        3. 使用 `os.replace()` 原子替换
        """
        content = ENTRY_DELIMITER.join(entries) if entries else ""
        fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp", prefix=".mem_")
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                handle.write(content)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(tmp_path, path)
        except BaseException:
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
            raise
