agentic-mobile-control/device/manager.py

from __future__ import annotations

from collections.abc import Callable
from dataclasses import replace
from threading import RLock
from time import sleep

from core.errors import (
    DeviceBusyError,
    DeviceNotFoundError,
    DeviceOfflineError,
    DriverError,
)
from core.models import Device, DeviceStatus
from driver.base import Driver

DriverFactory = Callable[[], Driver]


class DeviceManager:
    def __init__(self) -> None:
        self._devices: dict[str, Device] = {}
        self._factories: dict[str, DriverFactory] = {}
        self._drivers: dict[str, Driver] = {}
        self._lock = RLock()

    def register_device(
        self,
        device_id: str,
        driver_factory: DriverFactory,
        *,
        name: str | None = None,
        driver_type: str = "wda",
        connection_info: dict[str, object] | None = None,
        status: DeviceStatus = "idle",
        capability_tags: list[str] | None = None,
    ) -> Device:
        with self._lock:
            device = Device(
                id=device_id,
                name=name,
                status=status,
                driver_type=driver_type,
                connection_info=dict(connection_info or {}),
                capability_tags=list(capability_tags or []),
            )
            self._devices[device_id] = device
            self._factories[device_id] = driver_factory
            self._drivers.pop(device_id, None)
            return replace(device)

    def unregister_device(self, device_id: str) -> None:
        with self._lock:
            driver = self._drivers.pop(device_id, None)
            if driver:
                driver.disconnect()
            self._devices.pop(device_id, None)
            self._factories.pop(device_id, None)

    def list_devices(self) -> list[Device]:
        with self._lock:
            return [replace(device) for device in self._devices.values()]

    def status(self, device_id: str) -> DeviceStatus:
        with self._lock:
            return self._device(device_id).status

    def connect(
        self,
        device_id: str,
        *,
        max_retries: int = 2,
        retry_backoff_seconds: float = 0.25,
    ) -> Driver:
        if max_retries < 1:
            raise ValueError("max_retries must be at least 1")

        last_error: Exception | None = None
        for attempt in range(1, max_retries + 1):
            with self._lock:
                device = self._device(device_id)
                if device.status == "busy" and device_id in self._drivers:
                    return self._drivers[device_id]
                if device.status == "busy":
                    raise DeviceBusyError(f"device {device_id} is busy")
                driver = self._factories[device_id]()

            try:
                driver.connect()
            except Exception as exc:
                last_error = exc
                with self._lock:
                    status: DeviceStatus = (
                        "offline" if attempt == max_retries else "error"
                    )
                    self._set_status(device_id, status)
                if attempt < max_retries:
                    sleep(retry_backoff_seconds)
                continue

            with self._lock:
                self._drivers[device_id] = driver
                self._set_status(device_id, "busy")
            return driver

        raise DeviceOfflineError(
            f"device {device_id} is offline"
        ) from last_error

    def disconnect(self, device_id: str) -> None:
        with self._lock:
            self._device(device_id)
            driver = self._drivers.pop(device_id, None)
        if driver:
            driver.disconnect()
        with self._lock:
            self._set_status(device_id, "idle")

    def mark_error(self, device_id: str, *, offline: bool = False) -> None:
        with self._lock:
            self._device(device_id)
            self._drivers.pop(device_id, None)
            self._set_status(device_id, "offline" if offline else "error")

    def probe(self, device_id: str) -> bool:
        """Check whether an established driver is still reachable."""
        with self._lock:
            self._device(device_id)
            driver = self._drivers.get(device_id)
            if driver is None:
                return False
        try:
            health_check = getattr(driver, "health_check", None)
            if callable(health_check):
                health_check()
            else:
                # Compatibility for drivers implemented before health_check
                # existed. Built-in drivers use the non-screen health check.
                driver.screenshot()
        except Exception:
            self.mark_error(device_id, offline=True)
            return False
        return True

    def active_driver(self, device_id: str | None = None) -> Driver:
        with self._lock:
            if device_id is None:
                busy_ids = [
                    known_id
                    for known_id, device in self._devices.items()
                    if device.status == "busy" and known_id in self._drivers
                ]
                if len(busy_ids) != 1:
                    raise DriverError(
                        "exactly one connected device is required when no device_id is given"
                    )
                device_id = busy_ids[0]

            self._device(device_id)
            driver = self._drivers.get(device_id)
            if not driver:
                raise DeviceOfflineError(f"device {device_id} is not connected")
            return driver

    def _device(self, device_id: str) -> Device:
        try:
            return self._devices[device_id]
        except KeyError as exc:
            raise DeviceNotFoundError(f"unknown device {device_id}") from exc

    def _set_status(self, device_id: str, status: DeviceStatus) -> None:
        device = self._devices[device_id]
        self._devices[device_id] = replace(device, status=status)


DEFAULT_MANAGER = DeviceManager()