tianpu-ems/backend/app/collectors/manager.py

"""Collector Manager - orchestrates all device collectors."""
import logging
from typing import Optional

from sqlalchemy import select

from app.core.config import get_settings
from app.core.database import async_session
from app.models.device import Device
from app.collectors.base import BaseCollector
from app.collectors.modbus_tcp import ModbusTcpCollector
from app.collectors.mqtt_collector import MqttCollector
from app.collectors.http_collector import HttpCollector

logger = logging.getLogger("collector.manager")

# Full registry mapping protocol names to collector classes
COLLECTOR_REGISTRY: dict[str, type[BaseCollector]] = {
    "modbus_tcp": ModbusTcpCollector,
    "mqtt": MqttCollector,
    "http_api": HttpCollector,
}


def get_enabled_collectors() -> dict[str, type[BaseCollector]]:
    """Return collector registry filtered by customer config.

    If the customer config specifies a 'collectors' list, only those
    protocols are enabled. Otherwise fall back to the full registry.
    """
    settings = get_settings()
    customer_config = settings.load_customer_config()
    enabled_list = customer_config.get("collectors")
    if enabled_list is None:
        return COLLECTOR_REGISTRY
    enabled = {}
    for name in enabled_list:
        if name in COLLECTOR_REGISTRY:
            enabled[name] = COLLECTOR_REGISTRY[name]
        else:
            logger.warning("Customer config references unknown collector '%s', skipping", name)
    return enabled


class CollectorManager:
    """Manages lifecycle of all device collectors."""

    def __init__(self):
        self._collectors: dict[int, BaseCollector] = {}  # device_id -> collector
        self._running = False

    async def start(self):
        """Load active devices from DB and start their collectors."""
        self._running = True
        await self._load_and_start_collectors()
        logger.info("CollectorManager started with %d collectors", len(self._collectors))

    async def stop(self):
        """Stop all collectors."""
        self._running = False
        for device_id, collector in self._collectors.items():
            try:
                await collector.stop()
            except Exception as e:
                logger.error("Error stopping collector for device %d: %s", device_id, e)
        self._collectors.clear()
        logger.info("CollectorManager stopped")

    async def _load_and_start_collectors(self):
        """Load active devices with supported protocols and start collectors."""
        enabled = get_enabled_collectors()
        logger.info("Enabled collectors: %s", list(enabled.keys()))
        async with async_session() as session:
            result = await session.execute(
                select(Device).where(
                    Device.is_active == True,
                    Device.protocol.in_(list(enabled.keys())),
                )
            )
            devices = result.scalars().all()

        for device in devices:
            await self.start_collector(
                device.id,
                device.code,
                device.protocol,
                device.connection_params or {},
                device.collect_interval or 15,
            )

    async def start_collector(
        self,
        device_id: int,
        device_code: str,
        protocol: str,
        connection_params: dict,
        collect_interval: int,
    ) -> bool:
        """Start a single collector for a device."""
        if device_id in self._collectors:
            logger.warning("Collector already running for device %d", device_id)
            return False

        collector_cls = COLLECTOR_REGISTRY.get(protocol)
        if not collector_cls:
            logger.warning("No collector for protocol '%s' (device %s)", protocol, device_code)
            return False

        collector = collector_cls(device_id, device_code, connection_params, collect_interval)
        self._collectors[device_id] = collector
        await collector.start()
        logger.info("Started %s collector for %s", protocol, device_code)
        return True

    async def stop_collector(self, device_id: int) -> bool:
        """Stop collector for a specific device."""
        collector = self._collectors.pop(device_id, None)
        if not collector:
            return False
        await collector.stop()
        return True

    async def restart_collector(self, device_id: int) -> bool:
        """Restart collector for a device by reloading its config from DB."""
        await self.stop_collector(device_id)
        async with async_session() as session:
            result = await session.execute(
                select(Device).where(Device.id == device_id)
            )
            device = result.scalar_one_or_none()
        if not device or not device.is_active:
            return False
        return await self.start_collector(
            device.id,
            device.code,
            device.protocol,
            device.connection_params or {},
            device.collect_interval or 15,
        )

    def get_collector(self, device_id: int) -> Optional[BaseCollector]:
        return self._collectors.get(device_id)

    def get_all_status(self) -> list[dict]:
        """Return status of all collectors."""
        return [c.get_status() for c in self._collectors.values()]

    @property
    def collector_count(self) -> int:
        return len(self._collectors)

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