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

from fastmcp import FastMCP


class MemoryServer:
    def __init__(self):
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path("./memory_data.json")
        # Register tools
        self.mcp.tool()(self.save)
        self.mcp.tool()(self.get)
        self.mcp.tool()(self.delete)
        self.mcp.tool()(self.list_keys)
        self.mcp.tool()(self.save_with_namespace)
        self.mcp.tool()(self.get_by_namespace)

    def _load_memory(self) -> Dict[str, Dict]:
        """Load memory from JSON file."""
        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, Dict]):
        """Save memory to JSON file."""
        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 _validate_key(self, key: str) -> bool:
        """Simple validation to avoid path traversal etc."""
        if ".." in key or "/" in key or "\\" in key:
            return False
        return True

    def save(self, key: str, value: Any) -> bool:
        """
        Save a value under a key.

        Args:
            key: Unique identifier.
            value: Any JSON-serializable value.

        Returns:
            True on success, False otherwise.
        """
        if not self._validate_key(key):
            return False
        data = self._load_memory()
        data[key] = {
            "value": value,
            "timestamp": datetime.utcnow().isoformat()
        }
        self._save_memory(data)
        return True

    def get(self, key: str) -> Optional[Dict]:
        """
        Retrieve a value with metadata.

        Args:
            key: Identifier to look up.

        Returns:
            Dict with keys 'key', 'value', 'timestamp' or None.
        """
        if not self._validate_key(key):
            return None
        data = self._load_memory()
        entry = data.get(key)
        if entry is None:
            return None
        return {
            "key": key,
            "value": entry["value"],
            "timestamp": entry["timestamp"]
        }

    def delete(self, key: str) -> bool:
        """
        Delete a key from memory.

        Args:
            key: Identifier to delete.

        Returns:
            True if the key existed and was removed, False otherwise.
        """
        if not self._validate_key(key):
            return False
        data = self._load_memory()
        if key in data:
            del data[key]
            self._save_memory(data)
            return True
        return False

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

        Args:
            pattern: Wildcard pattern (* and ? supported).

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

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

        Args:
            key: Identifier.
            value: JSON-serializable value.
            namespace: Namespace name.

        Returns:
            True on success, False otherwise.
        """
        full_key = f"{namespace}:{key}"
        return self.save(full_key, value)

    def get_by_namespace(self, namespace: str = "default") -> List[Dict]:
        """
        Retrieve all entries belonging to a namespace.

        Args:
            namespace: Namespace to query.

        Returns:
            List of dictionaries with metadata for each entry.
        """
        prefix = f"{namespace}:"
        data = self._load_memory()
        result = []
        for k, v in data.items():
            if k.startswith(prefix):
                result.append({
                    "key": k,
                    "value": v["value"],
                    "timestamp": v["timestamp"]
                })
        return result


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