# memory_server.py
import json
import os
import sys
import asyncio
from datetime import datetime
from pathlib import Path
from typing import Any, Optional, Dict, List

from fastmcp import FastMCP, Client
from fnmatch import fnmatch

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

    def _load_memory(self) -> Dict[str, Dict[str, Any]]:
        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[str, Dict[str, Any]]):
        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 _sanitize_key(self, key: str) -> str:
        if "../" in key or key.startswith("/"):
            raise ValueError("Invalid key: contains path traversal")
        return key

    @property
    def _mcp(self):
        return self.mcp

    # Basic tools
    @self.mcp.tool()
    def save(self, key: str, value: Any) -> bool:
        """Сохраняет значение по ключу в память сервера."""
        key = self._sanitize_key(key)
        data = self._load_memory()
        data[key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
        self._save_memory(data)
        return True

    @self.mcp.tool()
    def get(self, key: str) -> Optional[Dict[str, Any]]:
        """Возвращает значение по ключу с метаданными."""
        key = self._sanitize_key(key)
        data = self._load_memory()
        if key in data:
            return {"key": key, **data[key]}
        return None

    @self.mcp.tool()
    def delete(self, key: str) -> bool:
        """Удаляет ключ из памяти сервера."""
        key = self._sanitize_key(key)
        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(self, pattern: str = "*") -> List[str]:
        """Возвращает список всех ключей с поддержкой wildcard-паттерна."""
        data = self._load_memory()
        return [k for k in data.keys() if fnmatch(k, pattern)]

    # Namespace tools
    @self.mcp.tool()
    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
        full_key = f"{namespace}:{key}"
        return self.save(full_key, value)

    @self.mcp.tool()
    def get_by_namespace(self, namespace: str = "default") -> List[Dict[str, Any]]:
        prefix = f"{namespace}:"
        data = self._load_memory()
        result = []
        for k, v in data.items():
            if k.startswith(prefix):
                result.append({"key": k, **v})
        return result

# ---------------------------------------------------------------------------
# Client for testing
# ---------------------------------------------------------------------------
async def test_client():
    client = Client("python memory_server.py")
    await client.connect()
    try:
        # Сохранение
        res = await client.call_tool("save_with_namespace", {"key": "username", "value": "Алексей", "namespace": "default"})
        print("Сохранено:", res)
        # Чтение
        res = await client.call_tool("get_by_namespace", {"namespace": "default"})
        print("Данные namespace 'default':")
        for item in res:
            print(f"  {item['key']}: {item['value']}")
        # Поиск по паттерну
        keys = await client.call_tool("list_keys", {"pattern": "*name"})
        print("Ключи с 'name':", keys)
    finally:
        await client.close()

# ---------------------------------------------------------------------------
# Entry points
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "test":
        asyncio.run(test_client())
    else:
        server = MemoryServer()
        server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")