"""NUT (UPS)-specific bot command handler.""" from __future__ import annotations import logging from collections.abc import Callable, Coroutine from typing import Any from ..database.models import CommandConfig, CommandTracker, ServiceProvider, TelegramBot from ..services import make_nut_provider from .base import CommandResponse, ProviderCommandHandler from .handler import _render_cmd_template _LOGGER = logging.getLogger(__name__) _NUT_COMMANDS = {"status", "devices", "battery"} # --------------------------------------------------------------------------- # Command dispatch table # --------------------------------------------------------------------------- _TEXT_COMMANDS: dict[str, Callable[..., Coroutine[Any, Any, dict[str, Any]]]] = {} def _text_cmd(fn: Callable[..., Coroutine[Any, Any, dict[str, Any]]]) -> Callable[..., Coroutine[Any, Any, dict[str, Any]]]: """Register a function in the text command dispatch table.""" name = fn.__name__.removeprefix("_cmd_") _TEXT_COMMANDS[name] = fn return fn class NutCommandHandler(ProviderCommandHandler): """Handles NUT-specific bot commands.""" provider_type = "nut" def get_provider_commands(self) -> set[str]: return _NUT_COMMANDS def get_rate_categories(self) -> dict[str, str]: return {"devices": "api", "battery": "api", "status": "api"} async def handle( self, cmd: str, args: str, count: int, locale: str, response_mode: str, provider: ServiceProvider, cmd_templates: dict[str, dict[str, str]], bot: TelegramBot, tracker: CommandTracker, config: CommandConfig, *, listener: Any = None, page: int = 1, ) -> CommandResponse | None: fn = _TEXT_COMMANDS.get(cmd) if fn is None: return None ctx = await fn(provider, count) return CommandResponse(text=_render_cmd_template(cmd_templates, cmd, locale, ctx)) async def _query_ups( provider: ServiceProvider, ) -> list[dict[str, Any]]: """Connect to a NUT provider and query UPS data.""" from notify_bridge_core.providers.nut.models import NutUpsData results: list[dict[str, Any]] = [] nut = make_nut_provider(provider) try: client = nut._make_client() await client.connect() try: devices = await client.list_ups() for dev in devices: variables = await client.list_var(dev.name) data = NutUpsData.from_variables(dev.name, variables) results.append({ "name": data.name, "description": data.description, "model": data.model, "manufacturer": data.manufacturer, "status": data.status, "battery_charge": int(data.battery_charge) if data.battery_charge is not None else None, "battery_runtime": data.battery_runtime_formatted, "ups_load": int(data.ups_load) if data.ups_load is not None else None, "input_voltage": str(data.input_voltage) if data.input_voltage is not None else None, "output_voltage": str(data.output_voltage) if data.output_voltage is not None else None, }) finally: await client.disconnect() except Exception as exc: _LOGGER.warning("Failed to query NUT provider %s: %s", provider.name, exc) return results @_text_cmd async def _cmd_status(provider: ServiceProvider, count: int) -> dict[str, Any]: devices = await _query_ups(provider) return {"devices": devices} @_text_cmd async def _cmd_devices(provider: ServiceProvider, count: int) -> dict[str, Any]: devices: list[dict[str, Any]] = [] nut = make_nut_provider(provider) try: device_list = await nut.list_collections() devices.extend(device_list) except Exception as exc: _LOGGER.warning("Failed to list devices from %s: %s", provider.name, exc) return {"devices": devices} @_text_cmd async def _cmd_battery(provider: ServiceProvider, count: int) -> dict[str, Any]: devices = await _query_ups(provider) return {"devices": devices}