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

操作:
- add: 追加内容到指定 section
- replace: 替换 section 内的文本
- remove: 删除整个 section
- read: 读取文件内容

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

from __future__ import annotations

from typing import Any

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


VALID_FILES = {"soul", "user", "memory", "daily"}
VALID_ACTIONS = {"add", "replace", "remove", "read"}


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",
                    },
                    "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",
                    },
                },
                "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}"}

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

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