from gaia.agents.base import Agent
from gaia.mcp import MCPClientMixin
from gaia.agents.base.tools import tool

class DesktopAutomationAgent(MCPClientMixin, Agent):
    """Agent that can control the Windows desktop."""

    def __init__(self):
        super().__init__()
        # Connect to Windows MCP server
        self.connect_mcp_server("windows", "uvx windows-mcp")

    @tool
    def open_app_and_click(self, app_name: str, button_text: str) -> str:
        """Open an application and click a button."""
        # Launch the application
        self.mcp_call("windows", "app-tool", {"action": "launch", "app": app_name})
        # Get current UI state to find the button
        state = self.mcp_call("windows", "state-tool", {})
        # Find button coordinates from UI tree
        coords = self._find_element_by_text(state["ui_tree"], button_text)
        # Click the button
        self.mcp_call("windows", "click-tool", {"x": coords[0], "y": coords[1]})
        return f"Launched {app_name} and clicked '{button_text}'"

from gaia.agents.base import Agent
from gaia.mcp import MCPClientMixin
from gaia.agents.base.tools import tool

class ComputerUseAgent(MCPClientMixin, Agent):
    """Agent that can control the computer like a human."""

    def __init__(self):
        super().__init__()
        self.connect_mcp_server("windows", "uvx windows-mcp")

    @tool
    def take_screenshot_and_analyze(self) -> str:
        """Capture the screen and describe what's visible."""
        state = self.mcp_call("windows", "state-tool", {"screenshot": True})
        # Use VLM to analyze the screenshot
        return self.analyze_image(state["screenshot"])

    @tool
    def click_element(self, description: str) -> str:
        """Click on an element matching the description."""
        state = self.mcp_call("windows", "state-tool", {})
        # Find element coordinates from UI tree
        coords = self.find_element(state["ui_tree"], description)
        self.mcp_call("windows", "click-tool", {"x": coords[0], "y": coords[1]})
        return f"Clicked on {description}"

    @tool
    def type_text(self, text: str) -> str:
        """Type text into the focused element."""
        self.mcp_call("windows", "type-tool", {"text": text})
        return f"Typed: {text}"

from gaia.agents.base import Agent
from gaia.mcp import MCPClientMixin
from gaia.agents.base.tools import tool

class FullStackAgent(MCPClientMixin, Agent):
    """Agent that combines multiple MCP capabilities."""

    def __init__(self):
        super().__init__()
        # Connect to multiple MCP servers
        self.connect_mcp_server("windows", "uvx windows-mcp")
        self.connect_mcp_server("browser", "npx @anthropic/puppeteer-mcp")
        self.connect_mcp_server("db", "uvx sqlite-mcp")

    @tool
    def end_to_end_workflow(self, task: str) -> str:
        """Execute a workflow spanning desktop, browser, and database."""
        # Open browser via Windows MCP
        self.mcp_call("windows", "app-tool", {"action": "launch", "app": "chrome"})
        # Navigate via Puppeteer MCP
        self.mcp_call("browser", "navigate", {"url": "https://example.com"})
        # Store results via SQLite MCP
        self.mcp_call("db", "query", {"sql": "INSERT INTO logs ..."})
        return f"Completed workflow: {task}"

# Example: GitHub automation agent
class GitHubAgent(MCPClientMixin, Agent):
    def __init__(self):
        super().__init__()
        self.connect_mcp_server("github", url="http://localhost:8080", transport="http")

    @tool
    def create_issue(self, repo: str, title: str, body: str) -> str:
        result = self.mcp_call("github", "create-issue", {
            "repo": repo, "title": title, "body": body
        })
        return f"Created issue #{result['number']}"

from typing import Any, Dict, List, Optional

class MCPClientMixin:
    """Mixin for agents that consume external MCP servers."""

    def connect_mcp_server(
        self,
        name: str,
        command: Optional[str] = None,
        url: Optional[str] = None,
        transport: str = "stdio",
        env: Optional[Dict[str, str]] = None,
        working_dir: Optional[str] = None,
        auto_register: bool = False,
        tool_prefix: Optional[str] = None
    ) -> None:
        """
        Connect to an MCP server.

        Args:
            name: Unique identifier for this connection
            command: Command to start the MCP server (for stdio transport)
            url: URL of the MCP server (for http/websocket transport)
            transport: Transport type - "stdio", "http", or "websocket"
            env: Environment variables for the server process
            working_dir: Working directory for the server
            auto_register: If True, register MCP tools as agent tools
            tool_prefix: Prefix for auto-registered tool names (e.g., "win_")

        Raises:
            MCPConnectionError: If server fails to start or initialize
            MCPTimeoutError: If connection times out

        Example:
            # stdio transport (local subprocess)
            self.connect_mcp_server("windows", command="uvx windows-mcp")

            # HTTP transport (remote server)
            self.connect_mcp_server("api", url="http://localhost:8080", transport="http")

            # WebSocket transport (persistent connection)
            self.connect_mcp_server("realtime", url="ws://localhost:9000", transport="websocket")

            # With auto-registration
            self.connect_mcp_server("win", command="uvx windows-mcp", auto_register=True, tool_prefix="win_")
        """

    def disconnect_mcp_server(self, name: str) -> None:
        """Disconnect from an MCP server and clean up resources."""

    def list_mcp_tools(self, server_name: str) -> List[Dict]:
        """
        List available tools from an MCP server.

        Returns:
            List of tool definitions with name, description, and inputSchema

        Raises:
            MCPNotConnectedError: If server is not connected
        """

    def mcp_call(
        self,
        server_name: str,
        tool_name: str,
        arguments: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Call a tool on an MCP server.

        Args:
            server_name: Name of the connected server
            tool_name: Name of the tool to call
            arguments: Tool arguments

        Returns:
            Tool execution result

        Raises:
            MCPNotConnectedError: If server is not connected
            MCPToolError: If tool execution fails

        Example:
            result = self.mcp_call("windows", "click-tool", {"x": 100, "y": 200})
        """

    def get_mcp_resources(self, server_name: str) -> List[Dict]:
        """List resources available from an MCP server."""

    def read_mcp_resource(self, server_name: str, uri: str) -> Any:
        """Read a resource from an MCP server."""

from gaia.mcp.exceptions import MCPConnectionError, MCPToolError, MCPNotConnectedError

class MyAgent(MCPClientMixin, Agent):
    @tool
    def safe_click(self, x: int, y: int) -> str:
        """Click with error handling."""
        try:
            self.mcp_call("windows", "click-tool", {"x": x, "y": y})
            return f"Clicked at ({x}, {y})"
        except MCPNotConnectedError:
            return "Error: Windows MCP server not connected"
        except MCPToolError as e:
            return f"Click failed: {e}"

from gaia.agents.base import Agent
from gaia.mcp import MCPClientMixin

class SmartAgent(MCPClientMixin, Agent):
    def __init__(self):
        super().__init__()
        # Connect and auto-register tools with prefix
        self.connect_mcp_server(
            name="windows",
            command="uvx windows-mcp",
            auto_register=True,  # Register as agent tools
            tool_prefix="win_"   # Prefix for tool names
        )

        # Now the LLM can directly call these tools:
        # - win_click_tool(x, y)
        # - win_type_tool(text)
        # - win_state_tool()
        # - win_app_tool(action, app)
        # etc.

from typing import Dict
from gaia.agents.base import Agent
from gaia.mcp import MCPClientMixin
from gaia.agents.base.tools import tool

class FormFillerAgent(MCPClientMixin, Agent):
    """Fill out forms in any Windows application."""

    def __init__(self):
        super().__init__()
        self.connect_mcp_server("windows", "uvx windows-mcp")

    @tool
    def fill_form(self, app_name: str, form_data: Dict[str, str]) -> str:
        """Open an app and fill out a form with the provided data."""
        # Launch application
        self.mcp_call("windows", "app-tool", {
            "action": "launch",
            "app": app_name
        })

        # Wait for app to load
        self.mcp_call("windows", "wait-tool", {"seconds": 2})

        # Get current UI state
        state = self.mcp_call("windows", "state-tool", {})

        # Fill each field
        for field_name, value in form_data.items():
            # Find field (using UI tree or coordinates)
            coords = self._find_field(state, field_name)
            self.mcp_call("windows", "click-tool", {"x": coords[0], "y": coords[1]})
            self.mcp_call("windows", "type-tool", {"text": value, "clear": True})
            self.mcp_call("windows", "shortcut-tool", {"keys": ["tab"]})

        return f"Filled {len(form_data)} fields in {app_name}"

# Run a CUA session
gaia cua

# Run with specific MCP servers
gaia cua --mcp windows --mcp browser

# Interactive mode with Windows automation
gaia cua --mcp windows --task "Open Calculator and compute 2+2"

# Add an MCP server to config
gaia mcp add windows --command "uvx windows-mcp"
gaia mcp add github --url "http://localhost:8080" --transport http

# List configured MCP servers
gaia mcp list

# Test MCP server connection
gaia mcp test windows

# Show tools from an MCP server
gaia mcp tools windows

# Remove an MCP server from config
gaia mcp remove windows

# In your agent
class MyAgent(MCPClientMixin, Agent):
    MCP_SERVERS = {
        # stdio transport (default)
        "windows": {
            "command": "uvx windows-mcp",
            "transport": "stdio",
            "auto_register": True,
            "tool_prefix": "win_"
        },
        # HTTP transport
        "github": {
            "url": "http://localhost:8080",
            "transport": "http",
            "auto_register": True,
            "tool_prefix": "gh_"
        },
        # WebSocket transport
        "realtime": {
            "url": "ws://localhost:9000",
            "transport": "websocket"
        }
    }

# ~/.gaia/config.yaml
mcp:
  servers:
    # stdio transport (local subprocess)
    windows:
      command: uvx windows-mcp
      transport: stdio
      enabled: true
      auto_register: true
      tool_prefix: win_

    # HTTP transport (remote server)
    github:
      url: http://localhost:8080
      transport: http
      enabled: true

    # WebSocket transport (persistent connection)
    realtime:
      url: ws://localhost:9000
      transport: websocket
      enabled: false

class FullyConnectedAgent(MCPClientMixin, MCPAgent, Agent):
    """Agent that both exposes and consumes MCP tools."""

    # Can be called by Claude Code, Cursor, etc.
    def get_mcp_tool_definitions(self):
        return [...]

    # Can call Windows MCP, browser MCP, etc.
    def __init__(self):
        super().__init__()
        self.connect_mcp_server("windows", "uvx windows-mcp")