CalejoControl/src/core/auto_discovery.py

"""
Auto-Discovery Module for Calejo Control Adapter.

Automatically discovers pump stations and pumps from database on startup.
"""

from typing import Dict, List, Optional
import asyncio
import structlog
from datetime import datetime, timedelta

from src.database.client import DatabaseClient

logger = structlog.get_logger()


class AutoDiscovery:
    """Auto-discovery module for pump stations and pumps."""
    
    def __init__(self, db_client: DatabaseClient, refresh_interval_minutes: int = 60):
        self.db_client = db_client
        self.refresh_interval_minutes = refresh_interval_minutes
        self.pump_stations: Dict[str, Dict] = {}
        self.pumps: Dict[str, List[Dict]] = {}
        self.last_discovery: Optional[datetime] = None
        self.discovery_running = False
    
    async def discover(self):
        """Discover all pump stations and pumps from database."""
        if self.discovery_running:
            logger.warning("auto_discovery_already_running")
            return
            
        self.discovery_running = True
        try:
            logger.info("auto_discovery_started")
            
            # Clear previous discovery
            self.pump_stations.clear()
            self.pumps.clear()
            
            # Discover pump stations
            stations = self.db_client.get_pump_stations()
            for station in stations:
                self.pump_stations[station['station_id']] = station
            
            # Discover pumps
            pumps = self.db_client.get_pumps()
            for pump in pumps:
                station_id = pump['station_id']
                if station_id not in self.pumps:
                    self.pumps[station_id] = []
                self.pumps[station_id].append(pump)
            
            self.last_discovery = datetime.now()
            
            logger.info(
                "auto_discovery_completed",
                station_count=len(self.pump_stations),
                pump_count=len(pumps),
                last_discovery=self.last_discovery.isoformat()
            )
            
            # Log discovered stations and pumps
            for station_id, station in self.pump_stations.items():
                station_pumps = self.pumps.get(station_id, [])
                logger.info(
                    "station_discovered",
                    station_id=station_id,
                    station_name=station['station_name'],
                    pump_count=len(station_pumps)
                )
                
                for pump in station_pumps:
                    logger.info(
                        "pump_discovered",
                        station_id=station_id,
                        pump_id=pump['pump_id'],
                        pump_name=pump['pump_name'],
                        control_type=pump['control_type'],
                        default_setpoint=pump['default_setpoint_hz']
                    )
                    
        except Exception as e:
            logger.error("auto_discovery_failed", error=str(e))
            raise
        finally:
            self.discovery_running = False
    
    async def start_periodic_discovery(self):
        """Start periodic auto-discovery in background."""
        logger.info(
            "periodic_discovery_started",
            refresh_interval_minutes=self.refresh_interval_minutes
        )
        
        while True:
            await asyncio.sleep(self.refresh_interval_minutes * 60)  # Convert to seconds
            try:
                await self.discover()
            except Exception as e:
                logger.error("periodic_discovery_failed", error=str(e))
    
    def get_stations(self) -> Dict[str, Dict]:
        """Get all discovered pump stations."""
        return self.pump_stations.copy()
    
    def get_pumps(self, station_id: Optional[str] = None) -> List[Dict]:
        """Get all discovered pumps, optionally filtered by station."""
        if station_id:
            return self.pumps.get(station_id, []).copy()
        else:
            all_pumps = []
            for station_pumps in self.pumps.values():
                all_pumps.extend(station_pumps)
            return all_pumps
    
    def get_pump(self, station_id: str, pump_id: str) -> Optional[Dict]:
        """Get a specific pump by station and pump ID."""
        station_pumps = self.pumps.get(station_id, [])
        for pump in station_pumps:
            if pump['pump_id'] == pump_id:
                return pump.copy()
        return None
    
    def get_station(self, station_id: str) -> Optional[Dict]:
        """Get a specific station by ID."""
        return self.pump_stations.get(station_id)
    
    def get_discovery_status(self) -> Dict[str, any]:
        """Get auto-discovery status information."""
        return {
            "last_discovery": self.last_discovery.isoformat() if self.last_discovery else None,
            "station_count": len(self.pump_stations),
            "pump_count": sum(len(pumps) for pumps in self.pumps.values()),
            "refresh_interval_minutes": self.refresh_interval_minutes,
            "discovery_running": self.discovery_running
        }
    
    def is_stale(self, max_age_minutes: int = 120) -> bool:
        """Check if discovery data is stale."""
        if not self.last_discovery:
            return True
        
        age = datetime.now() - self.last_discovery
        return age > timedelta(minutes=max_age_minutes)
    
    def validate_discovery(self) -> Dict[str, any]:
        """Validate discovered data for consistency."""
        issues = []
        
        # Check if all pumps have corresponding stations
        for station_id, pumps in self.pumps.items():
            if station_id not in self.pump_stations:
                issues.append(f"Pumps found for unknown station: {station_id}")
        
        # Check if all stations have pumps
        for station_id in self.pump_stations:
            if station_id not in self.pumps or not self.pumps[station_id]:
                issues.append(f"No pumps found for station: {station_id}")
        
        # Check pump data completeness
        for station_id, pumps in self.pumps.items():
            for pump in pumps:
                if not pump.get('control_type'):
                    issues.append(f"Pump {station_id}/{pump['pump_id']} missing control_type")
                if not pump.get('default_setpoint_hz'):
                    issues.append(f"Pump {station_id}/{pump['pump_id']} missing default_setpoint_hz")
        
        return {
            "valid": len(issues) == 0,
            "issues": issues,
            "station_count": len(self.pump_stations),
            "pump_count": sum(len(pumps) for pumps in self.pumps.values())
        }