# -*- coding: utf-8 -*-
# MCP-сервер для управления памятью агента
# Используем FastMCP, pydantic и python-dotenv

import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional

from dotenv import load_dotenv
from fastmcp import MCPServer, MCPRequest, MCPResponse
from pydantic import BaseModel, Field

# Загружаем переменные окружения из .env
load_dotenv()

# Путь к файлу хранения данных
DATA_FILE = Path(os.getenv("MEMORY_DATA_FILE", "memory_data.json"))

# ---------------------
# Модели данных
# ---------------------
class MemoryItem(BaseModel):
    key: str = Field(..., description="Ключ памяти")
    value: Any = Field(..., description="Значение памяти")
    namespace: Optional[str] = Field(None, description="Опциональное пространство имён")

# ---------------------
# Хранилище памяти
# ---------------------
class MemoryStore:
    def __init__(self, file_path: Path):
        self.file_path = file_path
        self._data: Dict[str, Any] = {}
        self._load()

    def _load(self) -> None:
        if self.file_path.exists():
            try:
                with self.file_path.open("r", encoding="utf-8") as f:
                    self._data = json.load(f)
            except Exception:
                self._data = {}
        else:
            self._data = {}

    def _save(self) -> None:
        with self.file_path.open("w", encoding="utf-8") as f:
            json.dump(self._data, f, ensure_ascii=False, indent=2)

    def save(self, key: str, value: Any) -> None:
        self._data[key] = value
        self._save()

    def get(self, key: str) -> Any:
        return self._data.get(key)

    def delete(self, key: str) -> bool:
        if key in self._data:
            del self._data[key]
            self._save()
            return True
        return False

    def list_keys(self) -> List[str]:
        return list(self._data.keys())

    def save_with_namespace(self, key: str, value: Any, namespace: str) -> None:
        ns_key = f"{namespace}:{key}"
        self.save(ns_key, value)

    def get_by_namespace(self, key: str, namespace: str) -> Any:
        ns_key = f"{namespace}:{key}"
        return self.get(ns_key)

# ---------------------
# MCP-обработчики
# ---------------------
class MemoryMCPHandler:
    def __init__(self, store: MemoryStore):
        self.store = store

    async def handle(self, request: MCPRequest) -> MCPResponse:
        # Ожидаем, что payload содержит "action" и параметры
        payload = request.payload
        action = payload.get("action")
        if not action:
            return MCPResponse(error="Missing action")

        try:
            if action == "save":
                key = payload["key"]
                value = payload["value"]
                self.store.save(key, value)
                return MCPResponse(result="ok")
            elif action == "get":
                key = payload["key"]
                value = self.store.get(key)
                return MCPResponse(result=value)
            elif action == "delete":
                key = payload["key"]
                success = self.store.delete(key)
                return MCPResponse(result=success)
            elif action == "list_keys":
                keys = self.store.list_keys()
                return MCPResponse(result=keys)
            elif action == "save_with_namespace":
                key = payload["key"]
                value = payload["value"]
                namespace = payload["namespace"]
                self.store.save_with_namespace(key, value, namespace)
                return MCPResponse(result="ok")
            elif action == "get_by_namespace":
                key = payload["key"]
                namespace = payload["namespace"]
                value = self.store.get_by_namespace(key, namespace)
                return MCPResponse(result=value)
            else:
                return MCPResponse(error="Unknown action")
        except KeyError as e:
            return MCPResponse(error=f"Missing parameter: {e}")
        except Exception as e:
            return MCPResponse(error=str(e))

# ---------------------
# Запуск сервера
# ---------------------
async def main():
    store = MemoryStore(DATA_FILE)
    handler = MemoryMCPHandler(store)
    server = MCPServer(handler=handler, host=os.getenv("HOST", "0.0.0.0"), port=int(os.getenv("PORT", 8000)))
    await server.serve()

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())