from fastmcp import FastMCP

import json
from datetime import datetime
from pathlib import Path
from typing import Any, Optional
import fnmatch


class MemoryServer:
    def __init__(self):
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path("./memory_data.json")
        self.register_tools()

    def _load_memory(self) -> dict:
        """Load memory from JSON file."""
        if not self.storage_path.exists():
            return {}
        with open(self.storage_path, 'r', encoding='utf-8') as f:
            return json.load(f)

    def _save_memory(self, data: dict):
        """Save memory to JSON file."""
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.storage_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def register_tools(self):
        @self.mcp.tool()
        def save(key: str, value: Any) -> bool:
            """Save a value by key to server memory.

            Args:
                key: Unique identifier for storage.
                value: Any serializable value

            Returns:
                True if successfully saved, False otherwise.
            """
            if key and "/../" in key:
                return False
            try:
                data = self._load_memory()
                data[key] = {
                    "value": value,
                    "timestamp": datetime.now().isoformat()
                }
                self._save_memory(data)
                return True
            except Exception:
                return False

        @self.mcp.tool()
        def get(key: str) -> Optional[dict]:
            """Get a value by key with metadata.

            Args:
                key: Unique identifier to look up.

            Returns:
                Dict with {"key": ..., "value": ..., "timestamp": ...}
                or None if key not found.
            """
            try:
                data = self._load_memory()
                if key in data:
                    return {
                        "key": key,
                        "value": data[key]["value"],
                        "timestamp": data[key]["timestamp"]
                    }
                return None
            except Exception:
                return None

        @self.mcp.tool()
        def delete(key: str) -> bool:
            """Delete a key from server memory.

            Args:
                key: Unique identifier to delete.

            Returns:
                True if successfully deleted, False if key not found.
            """
            try:
                data = self._load_memory()
                if key in data:
                    del data[key]
                    self._save_memory(data)
                    return True
                return False
            except Exception:
                return False

        @self.mcp.tool()
        def list_keys(pattern: str = "*") -> list[str]:
            """List all keys matching a wildcard pattern.

            Args:
                pattern: Filter pattern (supports * and ?)

            Returns:
                List of matching keys.
            """
            try:
                data = self._load_memory()
                return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]
            except Exception:
                return []

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

            Args:
                key: Unique identifier.
                value: Value to store.
                namespace: Namespace (default 'default').

            Returns:
                True if successfully saved, False otherwise.
            """
            combined_key = f"{namespace}:{key}"
            return save(combined_key, value)

        @self.mcp.tool()
        def get_by_namespace(namespace: str = "default") -> list[dict]:
            """Get all keys from a specified namespace.

            Args:
                namespace: Namespace to read from.

            Returns:
                List of dicts with metadata for all keys in namespace.
            """
            prefix = f"{namespace}:"
            try:
                data = self._load_memory()
                result = []
                for key, val_data in data.items():
                    if key.startswith(prefix):
                        result.append({
                            "key": key.replace(prefix, "", 1),
                            "value": val_data["value"],
                            "timestamp": val_data["timestamp"]
                        })
                return result
            except Exception:
                return []


if __name__ == "__main__":
    server = MemoryServer()
    server.mcp.run(
        transport="stdio",
        show_banner=False,
        log_level='ERROR'
    )