"""
Memory Server for MCP protocol.

Provides basic memory operations with optional namespace support.
"""
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastmcp import FastMCP

class MemoryServer:
    """FastMCP based memory server.

    Stores data in a JSON file located at ``./memory_data.json``.
    Each entry is stored as ``{namespace:key: {value, timestamp}}``.
    """

    def __init__(self) -> None:
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path("./memory_data.json")

    # ---------------------------------------------------------------------
    # Internal helpers
    # ---------------------------------------------------------------------
    def _load_memory(self) -> Dict[str, Any]:
        if not self.storage_path.exists():
            return {}
        with self.storage_path.open("r", encoding="utf-8") as f:
            return json.load(f)

    def _save_memory(self, data: Dict[str, Any]) -> None:
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        with self.storage_path.open("w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _make_key(self, key: str, namespace: str = "default") -> str:
        return f"{namespace}:{key}"

    # ---------------------------------------------------------------------
    # Tools exposed via FastMCP
    # ---------------------------------------------------------------------
    @self.mcp.tool()
    def save(self, key: str, value: Any) -> bool:
        """Save a value under ``default`` namespace.

        Args:
            key: Identifier for the value.
            value: Serializable data.

        Returns:
            True if saved successfully.
        """
        return self.save_with_namespace(key, value, "default")

    @self.mcp.tool()
    def get(self, key: str) -> Optional[Dict[str, Any]]:
        """Retrieve a value from ``default`` namespace.

        Args:
            key: Identifier to fetch.

        Returns:
            Dictionary with ``value`` and ``timestamp`` or None if missing.
        """
        return self.get_by_namespace("default").get(key)

    @self.mcp.tool()
    def delete(self, key: str) -> bool:
        """Delete a key from ``default`` namespace.

        Args:
            key: Identifier to remove.

        Returns:
            True if removed, False otherwise.
        """
        return self.delete_from_namespace(key, "default")

    @self.mcp.tool()
    def list_keys(self, pattern: str = "*") -> List[str]:
        """List all keys in ``default`` namespace matching a glob pattern.

        Args:
            pattern: Glob pattern (supports * and ?).

        Returns:
            List of key names without namespace prefix.
        """
        return self.list_keys_in_namespace("default", pattern)

    @self.mcp.tool()
    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
        """Save a value under a specific namespace.

        Args:
            key: Identifier for the value.
            value: Serializable data.
            namespace: Namespace name.

        Returns:
            True if saved successfully.
        """
        full_key = self._make_key(key, namespace)
        data = self._load_memory()
        data[full_key] = {
            "value": value,
            "timestamp": datetime.utcnow().isoformat() + "Z",
        }
        self._save_memory(data)
        return True

    @self.mcp.tool()
    def get_by_namespace(self, namespace: str = "default") -> Dict[str, Any]:
        """Return all key/value pairs in a namespace.

        Args:
            namespace: Namespace name.

        Returns:
            Mapping of key to dict with ``value`` and ``timestamp``.
        """
        data = self._load_memory()
        prefix = f"{namespace}:"
        result: Dict[str, Any] = {}
        for k, v in data.items():
            if k.startswith(prefix):
                key_name = k[len(prefix) :]
                result[key_name] = v
        return result

    @self.mcp.tool()
    def delete_from_namespace(self, key: str, namespace: str = "default") -> bool:
        """Delete a key from a specific namespace.

        Args:
            key: Identifier to remove.
            namespace: Namespace name.

        Returns:
            True if removed, False otherwise.
        """
        full_key = self._make_key(key, namespace)
        data = self._load_memory()
        if full_key in data:
            del data[full_key]
            self._save_memory(data)
            return True
        return False

    @self.mcp.tool()
    def list_keys_in_namespace(self, namespace: str = "default", pattern: str = "*") -> List[str]:
        """List keys in a namespace matching a glob pattern.

        Args:
            namespace: Namespace name.
            pattern: Glob pattern.

        Returns:
            List of key names without namespace prefix.
        """
        import fnmatch

        data = self._load_memory()
        prefix = f"{namespace}:"
        keys = [k[len(prefix) :] for k in data if k.startswith(prefix)]
        return [k for k in keys if fnmatch.fnmatch(k, pattern)]

    # ---------------------------------------------------------------------
    # Entry point
    # ---------------------------------------------------------------------
    def run(self) -> None:
        self.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")

if __name__ == "__main__":
    server = MemoryServer()
    server.run()