"""MemoryTool — Agent 可在对话中读写记忆的工具.

操作:
- add: 追加内容到指定 section
- replace: 替换 section 内的文本
- remove: 删除整个 section
- read: 读取文件内容
- update_soul: 动态更新 SOUL 文件（带版本追踪）

file 参数: soul | user | memory | daily
"""

from __future__ import annotations

import re
from datetime import datetime, timezone
from typing import Any

from agentkit.memory.profile import MemoryFile, MemoryStore
from agentkit.tools.base import Tool


VALID_FILES = {"soul", "user", "memory", "daily"}
VALID_ACTIONS = {"add", "replace", "remove", "read", "update_soul"}
# 受保护的 section：不允许通过 update_soul 修改，防止 Agent 覆盖版本追踪元数据
PROTECTED_SOUL_SECTIONS = {"版本", "更新历史"}


class MemoryTool(Tool):
    """Agent 可调用的记忆操作工具.

    让 Agent 在对话中读写 SOUL/USER/MEMORY/DAILY 记忆文件。
    """

    def __init__(self, memory_store: MemoryStore):
        super().__init__(
            name="memory",
            description="Read and write persistent memory files. Use to remember user preferences, project info, and notes across sessions.",
            input_schema={
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": list(VALID_ACTIONS),
                        "description": "Operation: add, replace, remove, read, update_soul",
                    },
                    "file": {
                        "type": "string",
                        "enum": list(VALID_FILES),
                        "description": "Memory file: soul (agent identity), user (user profile), memory (work notes), daily (today's log)",
                    },
                    "section": {
                        "type": "string",
                        "description": "Section name within the file (e.g. '项目信息', '偏好')",
                    },
                    "content": {
                        "type": "string",
                        "description": "Content to add or new text for replace",
                    },
                    "old_text": {
                        "type": "string",
                        "description": "Text to find for replace action",
                    },
                    "new_text": {
                        "type": "string",
                        "description": "Replacement text for replace action",
                    },
                    "reason": {
                        "type": "string",
                        "description": "Reason for update_soul action (stored in version history)",
                    },
                },
                "required": ["action", "file"],
            },
        )
        self._store = memory_store

    async def execute(self, **kwargs) -> dict[str, Any]:
        action = kwargs.get("action", "")
        file_key = kwargs.get("file", "")

        # Validate
        if file_key not in VALID_FILES:
            return {"success": False, "error": f"Invalid file: {file_key}. Must be one of {VALID_FILES}"}
        if action not in VALID_ACTIONS:
            return {"success": False, "error": f"Unknown action: {action}. Must be one of {VALID_ACTIONS}"}

        try:
            mf = self._store.get_file(file_key)

            if action == "read":
                content = mf.read()
                return {"success": True, "content": content}

            elif action == "add":
                section = kwargs.get("section", "")
                content = kwargs.get("content", "")
                if not section:
                    return {"success": False, "error": "section is required for add action"}
                mf.add_section(section, content)
                return {"success": True, "message": f"Added to {file_key}/{section}"}

            elif action == "replace":
                section = kwargs.get("section", "")
                old_text = kwargs.get("old_text", "")
                new_text = kwargs.get("new_text", "")
                if not section:
                    return {"success": False, "error": "section is required for replace action"}
                if not old_text:
                    return {"success": False, "error": "old_text is required for replace action"}
                success = mf.replace_section(section, old_text, new_text)
                if not success:
                    return {"success": False, "error": f"old_text not found in {file_key}/{section}"}
                return {"success": True, "message": f"Replaced in {file_key}/{section}"}

            elif action == "remove":
                section = kwargs.get("section", "")
                if not section:
                    return {"success": False, "error": "section is required for remove action"}
                mf.remove_section(section)
                return {"success": True, "message": f"Removed {file_key}/{section}"}

            elif action == "update_soul":
                section = kwargs.get("section", "")
                content = kwargs.get("content", "")
                reason = kwargs.get("reason", "")
                if not section:
                    return {"success": False, "error": "section is required for update_soul action"}
                if not content:
                    return {"success": False, "error": "content is required for update_soul action"}
                if section in PROTECTED_SOUL_SECTIONS:
                    return {"success": False, "error": f"Cannot modify protected section '{section}' via update_soul"}
                return await self._update_soul(mf, section, content, reason)

            return {"success": False, "error": f"Unhandled action: {action}"}

        except Exception as e:
            return {"success": False, "error": str(e)}

    async def _update_soul(
        self, mf: MemoryFile, section: str, content: str, reason: str
    ) -> dict[str, Any]:
        """执行 SOUL 动态更新，带版本追踪和更新历史."""
        # 解析当前版本号
        version = 1
        version_content = mf.read_section("版本")
        if version_content:
            match = re.search(r"版本:\s*(\d+)", version_content)
            if match:
                version = int(match.group(1))

        new_version = version + 1
        now = datetime.now(timezone.utc)
        timestamp = now.strftime("%Y-%m-%dT%H:%M:%S")
        date_str = now.strftime("%Y-%m-%d")

        # 更新目标 section
        if section in mf.list_sections():
            mf.remove_section(section)
        mf.add_section(section, content)

        # 更新版本 section
        version_text = f"版本: {new_version}\n更新时间: {timestamp}"
        if "版本" in mf.list_sections():
            mf.remove_section("版本")
        mf.add_section("版本", version_text)

        # 更新更新历史 section
        history_entry = f"- v{new_version} ({date_str}): 更新了{section}" + (f" - {reason}" if reason else "")

        history_lines: list[str] = []
        history_content = mf.read_section("更新历史")
        if history_content:
            history_lines = [line for line in history_content.strip().split("\n") if line.strip()]

        history_lines.append(history_entry)
        # 最多保留 10 条
        if len(history_lines) > 10:
            history_lines = history_lines[-10:]

        if "更新历史" in mf.list_sections():
            mf.remove_section("更新历史")
        mf.add_section("更新历史", "\n".join(history_lines))

        return {
            "success": True,
            "message": f"Updated soul/{section} to v{new_version}",
            "version": new_version,
        }
