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._register_tools()

    def _load_memory(self) -> Dict[str, dict]:
        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]) -> 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 _validate_key(self, key: str) -> bool:
        # Disallow path traversal characters
        return ".." not in key and "/" not in key and "\\" not in key

    def _register_tools(self):
        @self.mcp.tool()
        def save(key: str, value: Any) -> bool:
            """Save a value under a key."""
            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

        @self.mcp.tool()
        def get(key: str) -> Optional[dict]:
            """Retrieve a value with metadata."""
            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"]}

        @self.mcp.tool()
        def delete(key: str) -> bool:
            """Delete a key."""
            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

        @self.mcp.tool()
        def list_keys(pattern: str = "*") -> List[str]:
            """List keys matching a wildcard pattern."""
            data = self._load_memory()
            return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]

        @self.mcp.tool()
        def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool:
            """Save a value with a namespace."""
            if not self._validate_key(key) or not self._validate_key(namespace):
                return False
            full_key = f"{namespace}:{key}"
            return save(full_key, value)

        @self.mcp.tool()
        def get_by_namespace(namespace: str = "default") -> List[dict]:
            """Get all entries from a namespace."""
            if not self._validate_key(namespace):
                return []
            data = self._load_memory()
            prefix = f"{namespace}:"
            result = []
            for k, v in data.items():
                if k.startswith(prefix):
                    result.append(
                        {
                            "key": k[len(prefix) :],
                            "value": v["value"],
                            "timestamp": v["timestamp"],
                        }
                    )
            return result


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