chongming-plugins/aibox_sophon_adapter/mqtt_forward.py

"""
AI Box Sophon Adapter — 外部 MQTT 转发客户端

建立到外部 MQTT Broker 的长连接，将设备心跳数据转发到指定主题。
"""
import asyncio
import json
import logging
from urllib.parse import urlparse, unquote

from amqtt.client import MQTTClient
from amqtt.mqtt.constants import QOS_1

logger = logging.getLogger(__name__)


class MQTTForwardClient:
    """
    外部 MQTT 转发客户端

    在插件启动时建立到外部 MQTT Broker 的持久连接，
    提供 publish() 方法将数据转发到指定主题。

    支持 URI 中的 percent-encoded 凭据（如密码中的 @ 写为 %40）。

    使用方式:
        client = MQTTForwardClient(uri="mqtt://user:pass@host:port")
        await client.start()
        await client.publish("/topic", {"key": "value"})
        await client.stop()
    """

    def __init__(self, uri: str, topic: str = "/board_ping"):
        self._uri = uri
        self._topic = topic
        self._client: MQTTClient | None = None
        self._running = False
        self._lock = asyncio.Lock()

        # 预解析 URI 并解码凭据（amqtt 的 urlparse 不会自动 unquote）
        parsed = urlparse(uri)
        self._clean_uri = f"{parsed.scheme}://{parsed.hostname}:{parsed.port or 1883}"
        self._username: str | None = unquote(parsed.username) if parsed.username else None
        self._password: str | None = unquote(parsed.password) if parsed.password else None

    async def start(self) -> None:
        """建立到外部 MQTT Broker 的连接"""
        if self._running:
            return

        try:
            self._client = MQTTClient(
                config={
                    "auto_reconnect": True,
                    "reconnect_retries": -1,       # 无限重连
                    "reconnect_max_interval": 30,   # 最大重连间隔（秒）
                    "keep_alive": 60,
                    "connection_timeout": 10,        # 连接超时（秒），避免永久挂起
                }
            )
            # 手动初始化 session 并设置解码后的凭据（绕过 amqtt urlparse 不解码 %40 的问题）
            self._client.session = self._client._init_session(self._clean_uri)
            if self._username:
                self._client.session.username = self._username
            if self._password:
                self._client.session.password = self._password
            await self._client._do_connect()
            self._running = True
            logger.info(f"[MQTTForward] 已连接到外部 Broker: {self._clean_uri}")
            logger.info(f"[MQTTForward] 转发目标主题: {self._topic}")
        except Exception as e:
            logger.error(f"[MQTTForward] 连接外部 Broker 失败: {e}")
            self._client = None
            raise

    async def stop(self) -> None:
        """断开与外部 MQTT Broker 的连接"""
        if not self._running:
            return

        self._running = False
        if self._client:
            try:
                await self._client.disconnect()
            except Exception as e:
                logger.warning(f"[MQTTForward] 断开连接时异常: {e}")
            finally:
                self._client = None
        logger.info("[MQTTForward] 已断开外部 Broker 连接")

    async def publish(self, payload: dict, topic: str | None = None) -> bool:
        """
        发布消息到外部 MQTT Broker

        Args:
            payload: 要发送的 JSON 数据
            topic: 目标主题，默认使用初始化时配置的主题

        Returns:
            是否发送成功
        """
        if not self._running or self._client is None:
            logger.warning("[MQTTForward] 客户端未连接，跳过转发")
            return False

        target = topic or self._topic

        async with self._lock:
            try:
                data = json.dumps(payload, ensure_ascii=False, default=str)
                await self._client.publish(target, data.encode("utf-8"), QOS_1)
                logger.debug(f"[MQTTForward] 已转发到 {target}: {len(data)} 字节")
                return True
            except Exception as e:
                logger.error(f"[MQTTForward] 发布失败 [{target}]: {e}")
                return False

    @property
    def is_running(self) -> bool:
        return self._running