chongming-plugins/aibox_sophon_adapter/__init__.py

"""
AI Box Sophon Adapter 插件

基于 MQTT 的 AI Box 设备心跳适配器，实现：
  1. 通过 MQTT 主题 /board_ping 接收设备心跳上报
  2. 解析设备系统参数（温度、硬盘、内存、通道、任务、TPU 等）
  3. 持久化心跳数据到 SQLite 数据库
  4. 提供 HTTP API 查询设备在线状态和详情
  5. 后台定时检查离线设备，触发异常告警
  6. 视频上传接口 — 盒子告警时录制视频上传（FORM-DATA）
  7. 告警上报接口 — 违规事件 JSON 报文字动解析入库
  8. 告警查询接口 — 告警列表/详情查询

心跳超时机制:
  - 设备按配置间隔（默认 5 秒）上报心跳
  - 服务端每 check_interval 秒检查一次
  - 超过 heartbeat_timeout 秒未收到心跳的设备标记为离线
"""
import asyncio
import logging
from datetime import datetime
from pathlib import Path

from chongming_worker import auto_bind_models
from chongming_worker.runtime import create_workers as _create_workers

from .router import *
from .routes import router  # 自定义 FastAPI 路由（网关自动发现）

logger = logging.getLogger(__name__)

# 自动绑定 Pydantic 模型（用于 gen_models 生成的模型）
auto_bind_models("aibox_sophon_adapter")

# 离线检查任务引用
_offline_check_task: asyncio.Task | None = None


async def _check_offline_devices(interval: float):
    """
    后台任务：定期检查离线设备

    每隔 interval 秒扫描数据库，将超过 heartbeat_timeout 秒
    未收到心跳的设备标记为离线，并输出告警日志。
    """
    from .global_var import db, heartbeat_timeout

    while True:
        await asyncio.sleep(interval)
        try:
            if db is None:
                continue

            async with db.get_session() as session:
                from .models.orm import BoardHeartbeat
                from sqlalchemy import select

                stmt = select(BoardHeartbeat).where(BoardHeartbeat.is_online == True) # type: ignore
                result = await session.execute(stmt)
                records = result.scalars().all()

                now = datetime.now()
                offline_count = 0
                for record in records:
                    elapsed = (now - record.last_ping_time).total_seconds()
                    if elapsed > heartbeat_timeout:
                        record.is_online = False
                        record.updated_at = now
                        offline_count += 1
                        logger.warning(
                            f"[Offline] 设备 {record.board_id} (IP: {record.board_ip}) "
                            f"已离线 {elapsed:.0f} 秒（超时阈值: {heartbeat_timeout}s）"
                        )

                if offline_count > 0:
                    await session.commit()
                    logger.info(f"[Offline] 本轮标记 {offline_count} 个设备为离线")

        except Exception as e:
            logger.error(f"[OfflineChecker] 检查异常: {e}")


async def create_workers(session, config, count=1, worker_manager=None):
    """
    插件入口 — 初始化数据库和全局资源，创建 Worker 副本

    由 FastAPI lifespan 在启动时调用。
    """
    from .global_var import set_global_var, heartbeat_timeout as _default_timeout
    from chongming_database import Database
    from chongming_stream_manager import StreamManager

    # 读取心跳配置
    worker_config = config.get("worker", {}).get("heartbeat", {}) if isinstance(config, dict) else {}
    hb_interval = worker_config.get("interval", 5)
    hb_timeout = worker_config.get("timeout", _default_timeout)
    check_interval = worker_config.get("check_interval", 10)
    mediamtx_addr = worker_config.get("mediamtx_addr", "rtsp://127.0.0.1:8554")
    link_path = worker_config.get("link_path", "/data/bm-app")

    # 设定全局变量
    global heartbeat_timeout
    heartbeat_timeout = hb_timeout

    # 初始化数据库
    data_dir = Path(config.get("data_dir", "./data")) if isinstance(config, dict) else Path("./data")
    db_path = data_dir / "aibox_sophon_adapter.db"
    db_path.parent.mkdir(parents=True, exist_ok=True)

    db = Database(db_path)
    await db.initialize()

    # 初始化视频流控制器
    stream_manager = StreamManager(mediamtx_addr)

    # 读取告警保留配置
    alarm_retention_config = worker_config.get("alarm_retention", {}) if isinstance(config, dict) else {}
    alarm_max = alarm_retention_config.get("max_records", 0)

    # upload 根目录（用于清理告警关联视频）
    upload_root_path = config.get("data_dir", "./data/uploads") if isinstance(config, dict) else "./data/uploads"

    # 初始化外部 MQTT 转发客户端
    mqtt_fwd_config = worker_config.get("mqtt_forward", {}) if isinstance(config, dict) else {}
    mqtt_fwd_enabled = mqtt_fwd_config.get("enabled", False)
    mqtt_fwd_client = None
    if mqtt_fwd_enabled:
        from .mqtt_forward import MQTTForwardClient
        mqtt_fwd_uri = mqtt_fwd_config.get("uri", "")
        mqtt_fwd_topic = mqtt_fwd_config.get("topic", "/board_ping")
        if mqtt_fwd_uri:
            mqtt_fwd_client = MQTTForwardClient(uri=mqtt_fwd_uri, topic=mqtt_fwd_topic)
            try:
                await mqtt_fwd_client.start()
                logger.info(f"[MQTTForward] 外部转发已启用: {mqtt_fwd_uri} → {mqtt_fwd_topic}")
            except Exception as e:
                logger.warning(f"[MQTTForward] 启动失败，将跳过转发: {e}")
                mqtt_fwd_client = None

    set_global_var("db", db)
    set_global_var("heartbeat_timeout", hb_timeout)
    set_global_var("heartbeat_interval", hb_interval)
    set_global_var("stream_manager", stream_manager)
    set_global_var("link_path", link_path)
    set_global_var("mqtt_forward", mqtt_fwd_client)
    set_global_var("upload_root", upload_root_path)
    set_global_var("alarm_max_records", alarm_max)

    logger.info(
        f"[aibox_sophon_adapter] 已初始化: "
        f"心跳间隔={hb_interval}s, 超时={hb_timeout}s, 检查间隔={check_interval}s"
    )

    # 启动离线设备检查后台任务
    global _offline_check_task
    _offline_check_task = asyncio.create_task(_check_offline_devices(check_interval))

    # 创建 Worker 副本
    return _create_workers(session, config, count, worker_manager)


async def cleanup(worker_item: str):
    """
    清理资源

    在 Worker 关闭时调用，关闭数据库连接和后台任务。
    """
    from .global_var import db, stream_manager, mqtt_forward

    global _offline_check_task
    if _offline_check_task is not None and not _offline_check_task.done():
        _offline_check_task.cancel()
        try:
            await _offline_check_task
        except asyncio.CancelledError:
            pass
        logger.info("[aibox_sophon_adapter] 离线检查任务已停止")

    if db is not None:
        await db.dispose()

    if stream_manager is not None:
        await stream_manager.stop_all()

    if mqtt_forward is not None:
        await mqtt_forward.stop()

    logger.info(f"[aibox_sophon_adapter] {worker_item} 清理完毕")