pi-pm3/app/backend/managers/ups_manager.py

"""UPS Manager for Dangerous Pi.

Handles I2C battery monitoring, safe shutdown triggers,
and battery percentage reporting for UPS HAT devices.
"""
import asyncio
import subprocess
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Optional, Dict, Any
try:
    from smbus2 import SMBus
except ImportError:
    SMBus = None  # Mock for development environments

from .. import config


class BatteryStatus(str, Enum):
    """Battery status enum."""
    CHARGING = "charging"
    DISCHARGING = "discharging"
    FULL = "full"
    CRITICAL = "critical"
    UNKNOWN = "unknown"


class PowerSource(str, Enum):
    """Power source enum."""
    AC = "ac"
    BATTERY = "battery"
    UNKNOWN = "unknown"


@dataclass
class UPSStatus:
    """UPS status information."""
    battery_percentage: float
    voltage: float
    current: float
    power_source: PowerSource
    battery_status: BatteryStatus
    time_remaining: Optional[int] = None  # Minutes remaining on battery
    temperature: Optional[float] = None
    last_updated: Optional[str] = None
    is_available: bool = True
    error_message: Optional[str] = None


class UPSManager:
    """Manages UPS battery monitoring and safe shutdown."""

    def __init__(self):
        """Initialize the UPS manager."""
        self._i2c_address = int(config.UPS_I2C_ADDRESS, 16)
        self._check_interval = config.UPS_CHECK_INTERVAL
        self._bus: Optional[SMBus] = None
        self._status = UPSStatus(
            battery_percentage=0.0,
            voltage=0.0,
            current=0.0,
            power_source=PowerSource.UNKNOWN,
            battery_status=BatteryStatus.UNKNOWN,
            is_available=False
        )
        self._shutdown_threshold = 10.0  # Shutdown at 10% battery
        self._warning_threshold = 20.0   # Warning at 20% battery
        self._critical_threshold = 15.0  # Critical at 15% battery
        self._shutdown_initiated = False
        self._event_callbacks = []

    async def initialize(self) -> bool:
        """Initialize I2C connection to UPS.

        Returns:
            True if initialization successful, False otherwise
        """
        if SMBus is None:
            self._status.is_available = False
            self._status.error_message = "smbus2 library not available"
            return False

        try:
            # Try to open I2C bus (usually bus 1 on Raspberry Pi)
            self._bus = SMBus(1)

            # Try to read from the device to verify it exists
            await self._read_battery_data()

            self._status.is_available = True
            self._status.error_message = None
            return True

        except Exception as e:
            self._status.is_available = False
            self._status.error_message = f"Failed to initialize UPS: {e}"
            return False

    async def start_monitoring(self):
        """Start periodic battery monitoring in background."""
        if not await self.initialize():
            print(f"UPS not available: {self._status.error_message}")
            return

        while True:
            try:
                await self._update_battery_status()
                await self._check_battery_thresholds()
                await asyncio.sleep(self._check_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"Error in UPS monitoring: {e}")
                self._status.error_message = str(e)
                await asyncio.sleep(self._check_interval)

    async def get_status(self) -> UPSStatus:
        """Get current UPS status.

        Returns:
            UPSStatus object with current battery information
        """
        if self._status.is_available:
            await self._update_battery_status()
        return self._status

    async def shutdown_device(self, delay: int = 30):
        """Initiate safe shutdown of the device.

        Args:
            delay: Delay in seconds before shutdown (default: 30)
        """
        if self._shutdown_initiated:
            return

        self._shutdown_initiated = True

        # Notify all registered callbacks
        await self._trigger_event("shutdown_initiated", {
            "delay": delay,
            "battery_percentage": self._status.battery_percentage
        })

        # Wait for the delay
        await asyncio.sleep(delay)

        # Execute shutdown command
        try:
            subprocess.run(["sudo", "shutdown", "-h", "now"], check=True)
        except Exception as e:
            print(f"Failed to shutdown: {e}")
            self._shutdown_initiated = False

    def register_event_callback(self, callback):
        """Register a callback for UPS events.

        Args:
            callback: Async function to call on events
        """
        self._event_callbacks.append(callback)

    async def _update_battery_status(self):
        """Update battery status from I2C device."""
        if not self._bus or not self._status.is_available:
            return

        try:
            data = await self._read_battery_data()

            # Update status with new data
            self._status.battery_percentage = data['percentage']
            self._status.voltage = data['voltage']
            self._status.current = data['current']
            self._status.power_source = data['power_source']
            self._status.battery_status = data['battery_status']
            self._status.time_remaining = data.get('time_remaining')
            self._status.temperature = data.get('temperature')
            self._status.last_updated = datetime.utcnow().isoformat()
            self._status.error_message = None

        except Exception as e:
            print(f"Error reading battery data: {e}")
            self._status.error_message = str(e)

    async def _read_battery_data(self) -> Dict[str, Any]:
        """Read battery data from I2C device.

        This implementation is for MAX17040/MAX17048 fuel gauge.
        Adjust register addresses for different UPS HAT models.

        Returns:
            Dictionary with battery data
        """
        if not self._bus:
            raise RuntimeError("I2C bus not initialized")

        # Run I2C read in thread pool to avoid blocking
        loop = asyncio.get_event_loop()

        # Read voltage (registers 0x02-0x03)
        voltage_bytes = await loop.run_in_executor(
            None,
            self._bus.read_i2c_block_data,
            self._i2c_address,
            0x02,
            2
        )
        voltage_raw = (voltage_bytes[0] << 8) | voltage_bytes[1]
        voltage = (voltage_raw >> 4) * 1.25  # mV

        # Read state of charge (registers 0x04-0x05)
        soc_bytes = await loop.run_in_executor(
            None,
            self._bus.read_i2c_block_data,
            self._i2c_address,
            0x04,
            2
        )
        soc_raw = (soc_bytes[0] << 8) | soc_bytes[1]
        percentage = soc_raw / 256.0

        # Estimate current based on voltage change
        # This is a simple estimation; adjust based on your UPS HAT
        current = 0.0  # Some UPS HATs don't provide current readings

        # Determine power source and battery status
        # Typically voltage > 4.1V means charging
        if voltage > 4100:  # 4.1V in mV
            power_source = PowerSource.AC
            if percentage >= 99.0:
                battery_status = BatteryStatus.FULL
            else:
                battery_status = BatteryStatus.CHARGING
        else:
            power_source = PowerSource.BATTERY
            if percentage < self._critical_threshold:
                battery_status = BatteryStatus.CRITICAL
            else:
                battery_status = BatteryStatus.DISCHARGING

        # Estimate time remaining (simplified)
        time_remaining = None
        if power_source == PowerSource.BATTERY and current > 0:
            # Rough estimation: battery_mah * (percentage/100) / current_ma
            # This would need actual battery capacity from config
            pass

        return {
            'percentage': percentage,
            'voltage': voltage,
            'current': current,
            'power_source': power_source,
            'battery_status': battery_status,
            'time_remaining': time_remaining,
            'temperature': None  # Not all UPS HATs provide temperature
        }

    async def _check_battery_thresholds(self):
        """Check battery levels and trigger actions."""
        percentage = self._status.battery_percentage
        power_source = self._status.power_source

        # Only check thresholds when on battery power
        if power_source != PowerSource.BATTERY:
            self._shutdown_initiated = False
            return

        # Critical threshold - initiate shutdown
        if percentage <= self._shutdown_threshold and not self._shutdown_initiated:
            await self._trigger_event("battery_critical", {
                "percentage": percentage,
                "action": "shutdown_initiated"
            })
            await self.shutdown_device(delay=60)  # 60 second warning

        # Warning threshold
        elif percentage <= self._warning_threshold:
            await self._trigger_event("battery_warning", {
                "percentage": percentage,
                "threshold": self._warning_threshold
            })

        # Critical threshold (but not shutdown yet)
        elif percentage <= self._critical_threshold:
            await self._trigger_event("battery_low", {
                "percentage": percentage,
                "threshold": self._critical_threshold
            })

    async def _trigger_event(self, event_type: str, data: Dict[str, Any]):
        """Trigger event callbacks.

        Args:
            event_type: Type of event (e.g., "battery_warning")
            data: Event data
        """
        event = {
            "type": event_type,
            "timestamp": datetime.utcnow().isoformat(),
            "data": data
        }

        # Call all registered callbacks
        for callback in self._event_callbacks:
            try:
                await callback(event)
            except Exception as e:
                print(f"Error in event callback: {e}")

    def set_shutdown_threshold(self, threshold: float):
        """Set battery percentage threshold for automatic shutdown.

        Args:
            threshold: Battery percentage (0-100)
        """
        if 0 <= threshold <= 100:
            self._shutdown_threshold = threshold

    def set_warning_threshold(self, threshold: float):
        """Set battery percentage threshold for warnings.

        Args:
            threshold: Battery percentage (0-100)
        """
        if 0 <= threshold <= 100:
            self._warning_threshold = threshold

    def close(self):
        """Close I2C bus connection."""
        if self._bus:
            self._bus.close()
            self._bus = None


# Global UPS manager instance
_ups_manager: Optional[UPSManager] = None


def get_ups_manager() -> UPSManager:
    """Get the global UPS manager instance."""
    global _ups_manager
    if _ups_manager is None:
        _ups_manager = UPSManager()
    return _ups_manager