from fastmcp.exceptions import ToolError
from fastmcp.server.middleware import Middleware, MiddlewareContext
from minitap.mobile_use.sdk import Agent


class LocalDeviceHealthMiddleware(Middleware):
    """Middleware that checks local device health before tool calls.

    Only used in local mode (when ENABLE_LOCAL_DEVICE is True).
    """

    def __init__(self, agent: Agent):
        self.agent = agent

    async def on_call_tool(self, context: MiddlewareContext, call_next):
        # Skip agent check if platform is set (Limrun mode creates its own agent)
        tool_args = getattr(context.message, "arguments", None) or {}
        if tool_args.get("cloud_platform"):
            return await call_next(context)

        if not self.agent._initialized:
            raise ToolError(
                "Agent not initialized.\nMake sure a mobile device is connected and try again."
            )
        return await call_next(context)
