import json
import os
import time
from typing import Any, Dict, List, Optional

from fastmcp import FastMCP, mcp
from pydantic import BaseModel, ValidationError, Field

# Path to store memory data
DATA_FILE = os.path.join(os.path.dirname(__file__), "memory_data.json")


class MemoryItem(BaseModel):
    key: str
    value: Any
    timestamp: float


class MemoryServer:
    def __init__(self, data_file: str = DATA_FILE):
        self.data_file = data_file
        self.memory: Dict[str, MemoryItem] = {}
        self._load_memory()
        self.mcp = FastMCP()
        self._register_tools()

    def _load_memory(self):
        if not os.path.exists(self.data_file):
            self.memory = {}
            return
        try:
            with open(self.data_file, "r", encoding="utf-8") as f:
                raw = json.load(f)
            for k, v in raw.items():
                try:
                    self.memory[k] = MemoryItem(**v)
                except ValidationError:
                    continue
        except Exception:
            self.memory = {}

    def _save_memory(self):
        try:
            raw = {k: v.dict() for k, v in self.memory.items()}
            with open(self.data_file, "w", encoding="utf-8") as f:
                json.dump(raw, f, indent=2)
        except Exception:
            pass

    def _register_tools(self):
        @self.mcp.tool("save", "Saves a key-value pair.")
        def save(key: str, value: Any) -> bool:
            if not self._validate_key(key):
                return False
            self.memory[key] = MemoryItem(key=key, value=value, timestamp=time.time())
            self._save_memory()
            return True

        @self.mcp.tool("get", "Retrieves a key-value pair.")
        def get(key: str) -> Optional[Dict[str, Any]]:
            item = self.memory.get(key)
            if item:
                return item.dict()
            return None

        @self.mcp.tool("delete", "Deletes a key-value pair.")
        def delete(key: str) -> bool:
            if key in self.memory:
                del self.memory[key]
                self._save_memory()
                return True
            return False

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

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

        @self.mcp.tool("get_by_namespace", "Gets all items in a namespace.")
        def get_by_namespace(namespace: str = "default") -> List[Dict[str, Any]]:
            prefix = f"{namespace}:"
            return [item.dict() for k, item in self.memory.items() if k.startswith(prefix)]

    def _validate_key(self, key: str) -> bool:
        # disallow colon to avoid confusion with namespace separator
        return ":" not in key

    def run(self):
        self.mcp.run()


if __name__ == "__main__":
    server = MemoryServer()
    server.run()