"""分层记忆系统 — SOUL/USER/MEMORY/DAILY 文件管理.

参考 Hermes/OpenClaw 架构，实现 Agent 人格、用户档案、工作笔记、
日志的持久化存储与 system prompt 注入。
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable


class MemoryFile:
    """单个记忆文件的管理器，支持 section 级别 CRUD 和容量控制.

    文件格式为 Markdown，使用 `## Section` 组织内容::

        ## 身份
        我是小王，一个专业的 AI 助手。

        ## 性格
        友好、耐心、注重细节

    """

    def __init__(self, path: Path, char_budget: int | None = None,
                 protected_sections: set[str] | None = None):
        self.path = Path(path)
        self.char_budget = char_budget
        self._protected_sections = protected_sections or set()

    def read(self) -> str:
        """读取整个文件内容，文件不存在返回空字符串."""
        if not self.path.exists():
            return ""
        return self.path.read_text(encoding="utf-8")

    def write(self, content: str) -> None:
        """写入内容，自动创建父目录，超容量时自动裁剪.

        在内存中完成裁剪后一次性写入，避免中间不一致状态。
        """
        self.path.parent.mkdir(parents=True, exist_ok=True)
        if self.char_budget and len(content) > self.char_budget:
            content = self._trim_content(content, self._protected_sections or None)
        self.path.write_text(content, encoding="utf-8")

    def read_section(self, name: str) -> str:
        """读取指定 section 的内容（不含标题行）."""
        content = self.read()
        if not content:
            return ""
        # 匹配 ## name 后面的内容，直到下一个 ## 或文件末尾
        pattern = rf"^## {re.escape(name)}\s*\n(.*?)(?=^## |\Z)"
        match = re.search(pattern, content, re.MULTILINE | re.DOTALL)
        if match:
            return match.group(1).strip()
        return ""

    def add_section(self, name: str, content: str) -> None:
        """追加内容到指定 section，不存在则创建."""
        existing = self.read()
        section_content = self.read_section(name)
        if section_content:
            # 追加到已有 section
            old_text = section_content
            new_text = f"{old_text}\n{content}"
            self.replace_section(name, old_text, new_text)
        else:
            # 创建新 section
            new_section = f"\n## {name}\n{content}"
            if existing and not existing.endswith("\n"):
                new_section = "\n" + new_section
            self.write(existing + new_section)

    def replace_section(self, name: str, old_text: str, new_text: str) -> bool:
        """替换 section 内的文本，返回是否成功."""
        section_content = self.read_section(name)
        if old_text not in section_content:
            return False
        full_content = self.read()
        # 替换 section 内的文本
        pattern = rf"(^## {re.escape(name)}\s*\n)(.*?)(?=^## |\Z)"
        match = re.search(pattern, full_content, re.MULTILINE | re.DOTALL)
        if not match:
            return False
        original_section_body = match.group(2)
        new_section_body = original_section_body.replace(old_text, new_text, 1)
        updated = full_content[: match.start(2)] + new_section_body + full_content[match.end(2) :]
        self.write(updated)
        return True

    def remove_section(self, name: str) -> None:
        """删除整个 section（含标题行）."""
        content = self.read()
        if not content:
            return
        pattern = rf"^## {re.escape(name)}\s*\n.*?(?=^## |\Z)"
        new_content = re.sub(pattern, "", content, flags=re.MULTILINE | re.DOTALL).strip()
        self.write(new_content)

    def list_sections(self) -> list[str]:
        """列出所有 section 名称."""
        content = self.read()
        if not content:
            return []
        return re.findall(r"^## (.+)$", content, re.MULTILINE)

    def trim_to_budget(self, protected_sections: set[str] | None = None) -> None:
        """裁剪内容到容量上限，按 section 边界截断.

        保持原始 section 顺序，仅从后向前移除非保护 section。
        protected_sections 中的 section 始终保留，不参与裁剪。
        """
        if not self.char_budget:
            return
        content = self.read()
        if len(content) <= self.char_budget:
            return
        trimmed = self._trim_content(content, protected_sections)
        self.path.write_text(trimmed, encoding="utf-8")

    def _trim_content(self, content: str, protected_sections: set[str] | None = None) -> str:
        """在内存中裁剪内容到容量上限，返回裁剪后的字符串（不写文件）.

        保持原始 section 顺序，仅从后向前移除非保护 section。
        """
        if not self.char_budget or len(content) <= self.char_budget:
            return content

        protected = protected_sections or set()

        # 解析所有 section 及其位置
        sections: list[tuple[str, int, int]] = []  # (name, start, end)
        for match in re.finditer(r"^## (.+)$", content, re.MULTILINE):
            name = match.group(1).strip()
            start = match.start()
            next_match = re.search(r"^## ", content[match.end():], re.MULTILINE)
            if next_match:
                end = match.end() + next_match.start()
            else:
                end = len(content)
            sections.append((name, start, end))

        if not sections:
            return content[:self.char_budget]

        # 保持原始顺序，标记每个 section 是否受保护
        ordered: list[tuple[str, str, bool]] = []  # (name, text, is_protected)
        for name, start, end in sections:
            ordered.append((name, content[start:end], name in protected))

        # 从后向前移除非保护 section，直到总长度在预算内
        while ordered:
            total = len("\n\n".join(s[1] for s in ordered))
            if total <= self.char_budget:
                break
            # 从后向前找第一个非保护 section 移除
            for i in range(len(ordered) - 1, -1, -1):
                if not ordered[i][2]:
                    ordered.pop(i)
                    break
            else:
                break  # 所有剩余 section 都是受保护的

        return "\n\n".join(s[1] for s in ordered).strip()


@dataclass
class MemorySnapshot:
    """一次加载的所有记忆文件快照."""

    soul: str = ""
    user: str = ""
    memory: str = ""
    daily: str = ""
    total_chars: int = 0

    def is_empty(self) -> bool:
        return not any([self.soul, self.user, self.memory, self.daily])


# 容量上限常量（字符数）
SOUL_BUDGET = 2000
USER_BUDGET = 1400
MEMORY_BUDGET = 2200
DAILY_BUDGET = 1000  # 每天日志上限

# 默认 SOUL.md 内容
DEFAULT_SOUL = """## 身份
我是 AgentKit，一个专业的 AI 助手。

## 性格
专业、友好、注重细节

## 说话方式
简洁清晰，偶尔使用比喻帮助理解

## 做事准则
- 准确回答用户问题
- 主动记住用户提到的偏好和信息
- 不确定时坦诚说明
"""


class MemoryStore:
    """管理 SOUL/USER/MEMORY/DAILY 四类记忆文件.

    存储路径::

        base_dir/
        ├── SOUL.md
        ├── memories/
        │   ├── USER.md
        │   ├── MEMORY.md
        │   └── daily/
        │       ├── 2026-06-07.md
        │       └── 2026-06-08.md

    """

    def __init__(self, base_dir: Path | str | None = None,
                 on_change: Callable[[str], None] | None = None):
        if base_dir is None:
            base_dir = Path.home() / ".agentkit"
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self._on_change = on_change
        self._base_prompt: str = ""

        # 初始化四个 MemoryFile
        self._soul = MemoryFile(
            self.base_dir / "SOUL.md",
            char_budget=SOUL_BUDGET,
            protected_sections={"版本", "更新历史"},
        )
        self._user = MemoryFile(self.base_dir / "memories" / "USER.md", char_budget=USER_BUDGET)
        self._memory = MemoryFile(self.base_dir / "memories" / "MEMORY.md", char_budget=MEMORY_BUDGET)
        self._daily_dir = self.base_dir / "memories" / "daily"
        self._daily_dir.mkdir(parents=True, exist_ok=True)

    def get_file(self, file_key: str) -> MemoryFile:
        """获取指定类型的 MemoryFile.

        Args:
            file_key: "soul" | "user" | "memory" | "daily"
        """
        mapping = {
            "soul": self._soul,
            "user": self._user,
            "memory": self._memory,
        }
        if file_key in mapping:
            return mapping[file_key]
        if file_key == "daily":
            # daily 返回今天的日志文件
            today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
            return MemoryFile(self._daily_dir / f"{today}.md", char_budget=DAILY_BUDGET)
        raise ValueError(f"Invalid file_key: {file_key}. Must be soul/user/memory/daily")

    def ensure_defaults(self) -> None:
        """首次运行时创建默认 SOUL.md."""
        if not self._soul.read():
            self._soul.write(DEFAULT_SOUL.strip())

    def load_all(self) -> MemorySnapshot:
        """加载所有记忆文件."""
        soul = self._soul.read()
        user = self._user.read()
        memory = self._memory.read()
        daily = self.load_daily_logs()
        total = len(soul) + len(user) + len(memory) + len(daily)
        return MemorySnapshot(
            soul=soul,
            user=user,
            memory=memory,
            daily=daily,
            total_chars=total,
        )

    def load_daily_logs(self, days: int = 2) -> str:
        """加载最近 N 天的日志."""
        parts: list[str] = []
        for i in range(days):
            from datetime import timedelta

            date = datetime.now(timezone.utc) - timedelta(days=i)
            filename = f"{date.strftime('%Y-%m-%d')}.md"
            daily_file = MemoryFile(self._daily_dir / filename)
            content = daily_file.read()
            if content:
                parts.append(f"### {date.strftime('%Y-%m-%d')}\n{content}")
        return "\n\n".join(parts)

    def archive_old_dailies(self, keep_days: int = 2) -> int:
        """归档超过 N 天的日志（删除旧文件）."""
        from datetime import timedelta

        count = 0
        cutoff = datetime.now(timezone.utc) - timedelta(days=keep_days)
        if not self._daily_dir.exists():
            return 0
        for f in self._daily_dir.glob("*.md"):
            # 从文件名解析日期
            try:
                date_str = f.stem  # e.g. "2026-06-07"
                file_date = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
                if file_date < cutoff:
                    f.unlink()
                    count += 1
            except ValueError:
                continue
        return count

    def build_system_prompt(self, snapshot: MemorySnapshot, base_prompt: str = "") -> str:
        """将记忆注入 system prompt.

        格式::

            <agent-identity>
            [SOUL.md content]
            </agent-identity>

            <user-profile>
            [USER.md content]
            </user-profile>

            <agent-notes>
            [MEMORY.md content]
            </agent-notes>

            <recent-activity>
            [DAILY.md content]
            </recent-activity>

            [base_prompt]
        """
        # 保存 base_prompt 供后续刷新使用
        if base_prompt:
            self._base_prompt = base_prompt

        parts: list[str] = []

        if snapshot.soul:
            parts.append(f"<agent-identity>\n{snapshot.soul}\n</agent-identity>")
        if snapshot.user:
            parts.append(f"<user-profile>\n{snapshot.user}\n</user-profile>")
        if snapshot.memory:
            parts.append(f"<agent-notes>\n{snapshot.memory}\n</agent-notes>")
        if snapshot.daily:
            parts.append(f"<recent-activity>\n{snapshot.daily}\n</recent-activity>")

        if base_prompt:
            parts.append(base_prompt)

        return "\n\n".join(parts) if parts else base_prompt

    def refresh_system_prompt(self) -> str:
        """重新加载所有记忆文件并构建 system prompt.

        在 MemoryTool 写入记忆后调用，确保 agent 的 _system_prompt
        反映最新的记忆内容。
        """
        snapshot = self.load_all()
        return self.build_system_prompt(snapshot, self._base_prompt)

    def notify_change(self) -> None:
        """记忆文件变更后通知回调，刷新所有订阅者的 system prompt."""
        if self._on_change is None:
            return
        try:
            new_prompt = self.refresh_system_prompt()
            self._on_change(new_prompt)
        except Exception:
            import logging
            logging.getLogger(__name__).warning("Memory notify_change failed", exc_info=True)