Session Buddy

"""Oneiric MCP client for storage adapter discovery.

This module provides an MCP client wrapper that Session-Buddy can use to
communicate with Oneiric MCP server for storage adapter discovery and resolution.

Example:
    >>> from session_buddy.mcp_clients import OneiricMCPClient
    >>> async with OneiricMCPClient() as client:
    ...     adapters = await client.list_storage_adapters()
    ...     print(f"Found {len(adapters)} storage adapters")
"""

from __future__ import annotations

import asyncio
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast

if TYPE_CHECKING:
    from mcp import ClientSession, StdioServerParameters

logger = logging.getLogger(__name__)


class OneiricMCPClient:
    """MCP client for communicating with Oneiric MCP server.

    This client provides a simplified interface for querying Oneiric MCP
    for storage adapter discovery, resolution, and health checks.

    Attributes:
        server_path: Path to Oneiric MCP server directory
        timeout_seconds: Timeout for MCP tool calls

    Example:
        >>> async with OneiricMCPClient() as client:
        ...     adapters = await client.list_storage_adapters()
        ...     for adapter in adapters:
        ...         print(f"Provider: {adapter['provider']}")
    """

    def __init__(
        self,
        server_path: str | Path | None = None,
        timeout_seconds: int = 10,
    ) -> None:
        """Initialize Oneiric MCP client.

        Args:
            server_path: Path to Oneiric MCP server directory
            timeout_seconds: Timeout for MCP tool calls in seconds
        """
        self.server_path = (
            Path(server_path) if server_path else self._default_server_path()
        )
        self.timeout_seconds = timeout_seconds
        self._session: ClientSession | None = None
        self._initialized = False

    def _default_server_path(self) -> Path:
        """Get default Oneiric MCP server path."""
        # Try common locations
        candidates = [
            Path.home() / "Projects" / "oneiric-mcp",
            Path("/Users/les/Projects/oneiric-mcp"),
            Path.cwd().parent / "oneiric-mcp",
        ]

        for path in candidates:
            if (path / "oneiric_mcp" / "__init__.py").exists():
                return path

        # Fallback to home directory
        return Path.home() / "Projects" / "oneiric-mcp"

    def _create_server_params(self) -> StdioServerParameters:
        """Create stdio server parameters for Oneiric MCP.

        Returns:
            StdioServerParameters configured for Oneiric MCP

        Raises:
            ImportError: If mcp package is not available
        """
        try:
            from mcp import StdioServerParameters
        except ImportError as e:
            msg = "mcp package is required. Install with: pip install mcp"
            raise ImportError(msg) from e

        return StdioServerParameters(
            command="uv",
            args=[
                "--directory",
                str(self.server_path),
                "run",
                "python",
                "-m",
                "oneiric_mcp",
            ],
        )

    async def __aenter__(self) -> OneiricMCPClient:
        """Enter context manager and connect to MCP server.

        Returns:
            Self for fluent chaining

        Raises:
            RuntimeError: If connection fails
        """
        try:
            from mcp import ClientSession
        except ImportError as e:
            msg = "mcp package is required. Install with: pip install mcp"
            raise ImportError(msg) from e

        server_params = self._create_server_params()
        self._session = ClientSession(server_params)

        try:
            await self._session.__aenter__()
            await self._session.initialize()
            self._initialized = True
            logger.info(f"Connected to Oneiric MCP at {self.server_path}")
        except Exception as e:
            logger.error(f"Failed to connect to Oneiric MCP: {e}")
            self._session = None
            self._initialized = False
            raise

        return self

    async def __aexit__(self, *args: Any) -> None:
        """Exit context manager and close connection."""
        if self._session:
            try:
                await self._session.__aexit__(*args)
                logger.info("Disconnected from Oneiric MCP")
            except Exception as e:
                logger.warning(f"Error closing Oneiric MCP connection: {e}")
            finally:
                self._session = None
                self._initialized = False

    def _ensure_connected(self) -> ClientSession:
        """Ensure client is connected.

        Returns:
            The MCP client session

        Raises:
            RuntimeError: If client is not connected
        """
        if not self._initialized or not self._session:
            msg = (
                "Client not connected. Use async context manager: "
                "'async with OneiricMCPClient() as client:'"
            )
            raise RuntimeError(msg)
        return self._session

    async def list_storage_adapters(self) -> list[dict[str, Any]]:
        """List all available storage adapters from Oneiric.

        Returns:
            List of storage adapter dictionaries with provider, priority, etc.

        Raises:
            RuntimeError: If client is not connected
            asyncio.TimeoutError: If request times out

        Example:
            >>> async with OneiricMCPClient() as client:
            ...     adapters = await client.list_storage_adapters()
            ...     for adapter in adapters:
            ...         print(f"{adapter['provider']}: priority={adapter['priority']}")
        """
        session = self._ensure_connected()

        try:
            result = await asyncio.wait_for(
                session.call_tool("list_adapters", {"category": "storage"}),
                timeout=self.timeout_seconds,
            )

            # Cast to dict for dict-like access
            result_dict = cast(dict[str, Any], result)  # type: ignore[assignment]
            adapters = result_dict.get("adapters", [])  # type: ignore[attr-defined]
            logger.info(f"Listed {len(adapters)} storage adapters")
            return result_dict.get("adapters", [])  # type: ignore[return-value]

        except TimeoutError:
            logger.error(
                f"Timeout listing storage adapters after {self.timeout_seconds}s"
            )
            raise
        except Exception as e:
            logger.error(f"Failed to list storage adapters: {e}")
            return []

    async def resolve_storage_backend(
        self,
        provider: str,
    ) -> dict[str, Any]:
        """Resolve a specific storage backend via Oneiric.

        Args:
            provider: Storage provider name (e.g., "s3", "local", "azure")

        Returns:
            Resolved adapter information with selected status

        Raises:
            RuntimeError: If client is not connected
            asyncio.TimeoutError: If request times out

        Example:
            >>> async with OneiricMCPClient() as client:
            ...     result = await client.resolve_storage_backend("s3")
            ...     if result["selected"]:
            ...         print("S3 storage backend resolved successfully")
        """
        session = self._ensure_connected()

        try:
            result = await asyncio.wait_for(
                session.call_tool(
                    "resolve_adapter",
                    {
                        "domain": "adapter",
                        "key": "storage",
                        "provider": provider,
                    },
                ),
                timeout=self.timeout_seconds,
            )

            # Cast to dict for dict-like access
            result_dict = cast(dict[str, Any], result)  # type: ignore[assignment]
            logger.info(
                f"Resolved storage backend: {provider}, selected={result_dict.get('selected', False)}"  # type: ignore[attr-defined]
            )
            return result_dict  # type: ignore[return-value]

        except TimeoutError:
            logger.error(
                f"Timeout resolving {provider} backend after {self.timeout_seconds}s"
            )
            raise
        except Exception as e:
            logger.error(f"Failed to resolve storage backend {provider}: {e}")
            return {
                "error": str(e),
                "provider": provider,
                "selected": False,
                "healthy": False,
            }

    async def check_storage_health(self, provider: str) -> dict[str, Any]:
        """Check health status of a storage backend.

        Args:
            provider: Storage provider name to check

        Returns:
            Health check result with healthy status

        Raises:
            RuntimeError: If client is not connected
            asyncio.TimeoutError: If request times out

        Example:
            >>> async with OneiricMCPClient() as client:
            ...     health = await client.check_storage_health("s3")
            ...     if health["healthy"]:
            ...         print("S3 backend is healthy")
        """
        session = self._ensure_connected()

        try:
            result = await asyncio.wait_for(
                session.call_tool(
                    "get_adapter_health",
                    {
                        "domain": "adapter",
                        "key": "storage",
                        "provider": provider,
                    },
                ),
                timeout=self.timeout_seconds,
            )

            # Cast to dict for dict-like access
            result_dict = cast(dict[str, Any], result)  # type: ignore[assignment]
            logger.info(
                f"Health check for {provider}: healthy={result_dict.get('healthy', False)}"  # type: ignore[attr-defined]
            )
            return result_dict  # type: ignore[return-value]

        except TimeoutError:
            logger.error(
                f"Timeout checking {provider} health after {self.timeout_seconds}s"
            )
            raise
        except Exception as e:
            logger.error(f"Failed to check storage health for {provider}: {e}")
            return {
                "error": str(e),
                "provider": provider,
                "healthy": False,
                "has_health_check": False,
            }

    async def list_adapter_categories(self) -> dict[str, Any]:
        """List all available adapter categories.

        Returns:
            Dictionary with categories list

        Example:
            >>> async with OneiricMCPClient() as client:
            ...     categories = await client.list_adapter_categories()
            ...     print(f"Available: {categories['categories']}")
        """
        session = self._ensure_connected()

        try:
            result = await asyncio.wait_for(
                session.call_tool("list_adapter_categories", {}),
                timeout=self.timeout_seconds,
            )

            # Cast to dict for dict-like access
            result_dict = cast(dict[str, Any], result)  # type: ignore[assignment]
            logger.info(f"Listed {result_dict.get('count', 0)} adapter categories")  # type: ignore[attr-defined]
            return result_dict  # type: ignore[return-value]

        except TimeoutError:
            logger.error(f"Timeout listing categories after {self.timeout_seconds}s")
            raise
        except Exception as e:
            logger.error(f"Failed to list adapter categories: {e}")
            return {"count": 0, "categories": []}


__all__ = ["OneiricMCPClient"]