mcbluetooth-esp32/src/mcbluetooth_esp32/tools/connection.py

"""Connection management tools for the ESP32 test harness."""

from __future__ import annotations

import asyncio
from typing import Any

from fastmcp import FastMCP

from ..protocol import CMD_GET_INFO, CMD_GET_STATUS, CMD_PING, CMD_RESET, Status
from ..serial_client import CommandTimeout, NotConnected, get_client, init_client


def register_tools(mcp: FastMCP) -> None:
    """Register connection management tools with the MCP server."""

    @mcp.tool()
    async def esp32_connect(
        port: str = "/dev/ttyUSB0",
        baudrate: int = 115200,
    ) -> dict[str, Any]:
        """Connect to an ESP32 over serial UART.

        Initialises the serial client (if needed) and opens the connection.
        Waits briefly for the ESP32 boot event before returning status.

        Args:
            port: Serial device path (e.g. "/dev/ttyUSB0", "/dev/ttyACM0")
            baudrate: UART baud rate (default 115200)

        Returns:
            Connection status including port, baudrate, and whether the
            device responded with a boot event.
        """
        from ..serial_client import get_client_or_none

        client = get_client_or_none()
        if client is None:
            client = init_client(port=port, baudrate=baudrate)

        try:
            await client.connect()
        except Exception as e:
            return {"connected": False, "error": str(e), "port": port}

        # Give the ESP32 a moment to send its boot event
        boot_received = False
        try:
            event = await client.event_queue.wait_for(event_name="boot", timeout=2.0)
            boot_received = event is not None
        except TimeoutError:
            pass

        # Ready probe: retry ping until the firmware's command handler is up.
        # The boot event fires early in app_main, before the UART command
        # task is fully initialised, so the first command can get lost.
        ready = False
        for attempt in range(5):
            try:
                resp = await client.send_command(CMD_PING, timeout=1.0)
                if resp.status == Status.OK:
                    ready = True
                    break
            except (CommandTimeout, NotConnected):
                await asyncio.sleep(0.3)

        return {
            "connected": True,
            "port": port,
            "baudrate": baudrate,
            "boot_event": boot_received,
            "ready": ready,
        }

    @mcp.tool()
    async def esp32_disconnect() -> dict[str, Any]:
        """Disconnect from the ESP32 serial device.

        Closes the UART connection. The client instance is retained so
        a subsequent esp32_connect can reuse it.

        Returns:
            Disconnection status.
        """
        try:
            client = get_client()
        except NotConnected:
            return {"connected": False, "message": "Already disconnected"}

        try:
            await client.disconnect()
            return {"connected": False, "message": "Disconnected"}
        except Exception as e:
            return {"connected": False, "error": str(e)}

    @mcp.tool()
    async def esp32_status() -> dict[str, Any]:
        """Query the ESP32 for its current status.

        Sends a get_status command and returns the device's reported
        state (Bluetooth mode, discoverable, connected peers, etc.).

        Returns:
            Status dict from the ESP32, or {"connected": false} if the
            serial link is down.
        """
        try:
            client = get_client()
        except NotConnected:
            return {"connected": False}

        try:
            resp = await client.send_command(CMD_GET_STATUS)
            if resp.status == Status.OK:
                return {"connected": True, **resp.data}
            return {"connected": True, "status": resp.status, "error": resp.data.get("error", "")}
        except NotConnected:
            return {"connected": False}
        except Exception as e:
            return {"connected": True, "error": str(e)}

    @mcp.tool()
    async def esp32_ping() -> dict[str, Any]:
        """Ping the ESP32 to verify the serial link is alive.

        Sends a lightweight ping command and expects a pong response.

        Returns:
            Response dict (typically {"pong": true}) or an error dict
            if the device is unreachable.
        """
        try:
            client = get_client()
            resp = await client.send_command(CMD_PING)
            if resp.status == Status.OK:
                return {"pong": True, **resp.data}
            return {"pong": False, "error": resp.data.get("error", "")}
        except NotConnected:
            return {"pong": False, "error": "Not connected to ESP32"}
        except Exception as e:
            return {"pong": False, "error": str(e)}

    @mcp.tool()
    async def esp32_reset() -> dict[str, Any]:
        """Reset (reboot) the ESP32.

        Sends a reset command over UART. The ESP32 will reboot, which
        drops the serial link. The client is disconnected automatically;
        call esp32_connect again after the device restarts.

        Returns:
            Reset status. The serial connection will be closed.
        """
        try:
            client = get_client()
        except NotConnected:
            return {"reset": False, "error": "Not connected to ESP32"}

        try:
            resp = await client.send_command(CMD_RESET)
            status = resp.status == Status.OK
        except Exception:
            # The reset itself may kill the link before the response arrives
            status = True

        # Connection will be lost after reboot, so disconnect cleanly
        try:
            await client.disconnect()
        except Exception:
            pass

        return {"reset": status, "message": "ESP32 is rebooting — reconnect after restart"}

    @mcp.tool()
    async def esp32_get_info() -> dict[str, Any]:
        """Retrieve device information from the ESP32.

        Returns firmware version, chip model, MAC addresses, and other
        hardware details reported by the device.

        Returns:
            Device information dict, or an error dict on failure.
        """
        try:
            client = get_client()
            resp = await client.send_command(CMD_GET_INFO)
            if resp.status == Status.OK:
                return resp.data
            return {"error": resp.data.get("error", "Failed to get device info")}
        except NotConnected:
            return {"error": "Not connected to ESP32"}
        except Exception as e:
            return {"error": str(e)}