#!/usr/bin/env python3
"""Main entry point for MCP memory server and client demonstration.

This file contains both the server implementation (MemoryServer) and a simple
client example that demonstrates how to use the server via FastMCP.

The server exposes six tools:

* save
* get
* delete
* list_keys
* save_with_namespace
* get_by_namespace

The client shows how to call these tools using the FastMCP Client.
"""

from __future__ import annotations

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

from fastmcp import FastMCP, Client

# ---------------------------------------------------------------------------
# Server implementation
# ---------------------------------------------------------------------------

class MemoryServer:
    """Simple key‑value store backed by a JSON file.

    Keys are stored as ``namespace:key``.  Each entry contains the value and a
    timestamp of the last write.
    """

    def __init__(self, storage_path: str | Path = "./memory_data.json"):
        self.mcp = FastMCP("Memory-Server")
        self.storage_path = Path(storage_path)
        self._ensure_storage()

    # ---------------------------------------------------------------------
    # Persistence helpers
    # ---------------------------------------------------------------------

    def _ensure_storage(self) -> None:
        if not self.storage_path.parent.exists():
            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        if not self.storage_path.exists():
            self.storage_path.write_text("{}", encoding="utf-8")

    def _load(self) -> Dict[str, Dict[str, Any]]:
        try:
            return json.loads(self.storage_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return {}

    def _save(self, data: Dict[str, Dict[str, Any]]) -> None:
        self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")

    # ---------------------------------------------------------------------
    # Tools
    # ---------------------------------------------------------------------

    @self.mcp.tool()
    def save(self, key: str, value: Any) -> bool:
        """Save a value under ``key`` in the default namespace.

        Parameters
        ----------
        key: str
            Identifier for the value.
        value: Any
            JSON‑serialisable value.

        Returns
        -------
        bool
            ``True`` if the value was stored.
        """
        data = self._load()
        full_key = f"default:{key}"
        data[full_key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
        self._save(data)
        return True

    @self.mcp.tool()
    def get(self, key: str) -> Optional[Dict[str, Any]]:
        """Retrieve a value from the default namespace.

        Parameters
        ----------
        key: str
            Identifier to look up.

        Returns
        -------
        dict | None
            Dictionary with ``key``, ``value`` and ``timestamp`` or ``None``.
        """
        data = self._load()
        full_key = f"default:{key}"
        entry = data.get(full_key)
        if entry is None:
            return None
        return {"key": full_key, "value": entry["value"], "timestamp": entry["timestamp"]}

    @self.mcp.tool()
    def delete(self, key: str) -> bool:
        """Delete a key from the default namespace.

        Parameters
        ----------
        key: str
            Identifier to delete.

        Returns
        -------
        bool
            ``True`` if the key existed and was removed.
        """
        data = self._load()
        full_key = f"default:{key}"
        if full_key in data:
            del data[full_key]
            self._save(data)
            return True
        return False

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

        Parameters
        ----------
        pattern: str, optional
            Wildcard pattern (``*`` and ``?`` supported).  Defaults to ``*``.

        Returns
        -------
        list[str]
            List of matching keys.
        """
        data = self._load()
        return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]

    @self.mcp.tool()
    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
        """Save a value under a specific namespace.

        Parameters
        ----------
        key: str
            Identifier.
        value: Any
            JSON‑serialisable value.
        namespace: str, optional
            Namespace prefix.  Defaults to ``default``.

        Returns
        -------
        bool
            ``True`` if stored.
        """
        data = self._load()
        full_key = f"{namespace}:{key}"
        data[full_key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
        self._save(data)
        return True

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

        Parameters
        ----------
        namespace: str, optional
            Namespace to query.  Defaults to ``default``.

        Returns
        -------
        list[dict]
            List of dictionaries with ``key``, ``value`` and ``timestamp``.
        """
        data = self._load()
        prefix = f"{namespace}:"
        result: List[Dict[str, Any]] = []
        for k, v in data.items():
            if k.startswith(prefix):
                result.append({"key": k, "value": v["value"], "timestamp": v["timestamp"]})
        return result

# ---------------------------------------------------------------------------
# Client demonstration
# ---------------------------------------------------------------------------

async def client_demo() -> None:
    """Simple demo that exercises the server tools via FastMCP Client."""
    client = Client("python main.py")  # launch the same script as a server
    await client.connect()
    try:
        # Save a value in default namespace
        await client.call_tool("save", {"key": "user_name", "value": "Alex"})
        # Retrieve it
        res = await client.call_tool("get", {"key": "user_name"})
        print("get default:", res)
        # Save with namespace
        await client.call_tool("save_with_namespace", {"key": "session_token", "value": "abc123", "namespace": "agent_1"})
        # List all keys
        keys = await client.call_tool("list_keys", {"pattern": "*"})
        print("all keys:", keys)
        # Get by namespace
        ns = await client.call_tool("get_by_namespace", {"namespace": "agent_1"})
        print("agent_1 namespace:", ns)
    finally:
        await client.close()

# ---------------------------------------------------------------------------
# Entry points
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # If the script is launched without arguments, run the server.
    import sys
    if len(sys.argv) == 1:
        server = MemoryServer()
        server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
    else:
        # Otherwise, run the client demo.
        asyncio.run(client_demo())