"""
Memory Server implementing MCP protocol using FastMCP.

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

from fastmcp import FastMCP
import fnmatch

class MemoryServer:
    """A simple MCP server that stores key/value pairs in a JSON file.

    The server exposes six tools:
        * save – store a value under a key.
        * get – retrieve a stored value.
        * delete – remove a key.
        * list_keys – list all keys with optional glob pattern.
        * save_with_namespace – same as ``save`` but prefixes the key with a namespace.
        * get_by_namespace – return all items belonging to a namespace.
    """

    def __init__(self, storage_path: str | Path = "./memory_data.json"):
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path(storage_path)
        # Ensure directory exists for future writes
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)

    def _load(self) -> Dict[str, Any]:
        """Load the entire memory store from disk.

        Returns an empty dict if file does not exist or is invalid.
        """
        try:
            if self.storage_path.exists():
                with self.storage_path.open("r", encoding="utf-8") as f:
                    return json.load(f)
        except Exception:
            # Corrupted file – start fresh
            pass
        return {}

    def _save(self, data: Dict[str, Any]) -> None:
        """Persist the memory store to disk.

        The file is written atomically by writing to a temp file first.
        """
        tmp = self.storage_path.with_suffix(".tmp")
        with tmp.open("w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        tmp.replace(self.storage_path)

    def _current_timestamp(self) -> str:
        return datetime.now(timezone.utc).isoformat()

    @property
    def mcp_tool(self):
        # Helper to expose the FastMCP instance for decorators
        return self.mcp

    # ---------- Basic tools ----------
    @FastMCP.tool()
    def save(key: str, value: Any) -> bool:
        """Store *value* under *key*.

        The value is JSON‑serialisable.  Existing keys are overwritten.
        Returns ``True`` on success.
        """
        data = self._load()
        data[key] = {"value": value, "timestamp": self._current_timestamp()}
        self._save(data)
        return True

    @FastMCP.tool()
    def get(key: str) -> Optional[Dict[str, Any]]:
        """Retrieve the stored item for *key*.

        Returns a dict with ``value`` and ``timestamp`` or ``None`` if not found.
        """
        data = self._load()
        return data.get(key)

    @FastMCP.tool()
    def delete(key: str) -> bool:
        """Remove *key* from the store.

        Returns ``True`` if key existed and was removed, otherwise ``False``.
        """
        data = self._load()
        if key in data:
            del data[key]
            self._save(data)
            return True
        return False

    @FastMCP.tool()
    def list_keys(pattern: str = "*") -> List[str]:
        """Return all keys matching *pattern*.

        The pattern supports Unix shell wildcards (``*``, ``?``).
        """
        data = self._load()
        return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]

    # ---------- Namespace tools ----------
    @FastMCP.tool()
    def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool:
        """Store *value* under ``namespace:key``.

        The key is prefixed with the namespace and a colon.  Existing entries are overwritten.
        Returns ``True`` on success.
        """
        namespaced_key = f"{namespace}:{key}"
        return self.save(namespaced_key, value)

    @FastMCP.tool()
    def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]:
        """Return all items belonging to *namespace*.

        Each returned dict contains ``key`` (without namespace prefix), ``value`` and ``timestamp``.
        """
        data = self._load()
        prefix = f"{namespace}:"
        result: List[Dict[str, Any]] = []
        for full_key, meta in data.items():
            if full_key.startswith(prefix):
                key_without_ns = full_key[len(prefix) :]
                item = {"key": key_without_ns, "value": meta["value"], "timestamp": meta["timestamp"]}
                result.append(item)
        return result

# ---------- Server entry point ----------
if __name__ == "__main__":
    server = MemoryServer()
    # Run with stdio transport – suitable for local testing and client usage.
    server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")