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

"""UPS Manager for Dangerous Pi.

Handles battery monitoring, safe shutdown triggers,
and battery percentage reporting for various UPS HAT devices.

Supports multiple UPS types via driver architecture:
- PiSugar (PiSugar 2/3 series via daemon)
- I2C Fuel Gauge (MAX17040/MAX17048)
"""
import asyncio
import subprocess
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Optional, Dict, Any

from .. import config
from .ups_drivers import UPSDriver, PiSugarDriver, I2CFuelGaugeDriver, auto_detect_driver


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, driver: Optional[UPSDriver] = None):
        """Initialize the UPS manager.

        Args:
            driver: UPS driver instance. If None, creates driver based on config.
        """
        self._check_interval = config.UPS_CHECK_INTERVAL
        self._driver = driver or self._create_driver()
        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._battery_capacity_mah = 1200  # Default for PiSugar 2 (1200mAh)
        self._shutdown_initiated = False
        self._config_warning_sent = False  # Only send config warning once
        self._event_callbacks = []

    def _create_driver(self) -> Optional[UPSDriver]:
        """Create UPS driver based on configuration.

        Returns:
            Appropriate UPS driver instance, or None for auto/none types
        """
        ups_type = config.UPS_TYPE.lower()

        if ups_type == "pisugar":
            return PiSugarDriver(
                host=config.UPS_PISUGAR_HOST,
                port=config.UPS_PISUGAR_PORT
            )
        elif ups_type == "i2c":
            i2c_address = int(config.UPS_I2C_ADDRESS, 16)
            return I2CFuelGaugeDriver(i2c_address=i2c_address)
        elif ups_type == "auto":
            # Auto-detection happens in initialize()
            return None
        elif ups_type == "none":
            # No UPS configured
            return None
        else:
            # Default to auto-detection for unknown types
            print(f"Unknown UPS type '{ups_type}', will auto-detect")
            return None

    async def initialize(self, startup_delay: float = 1.0) -> bool:
        """Initialize connection to UPS hardware.

        Args:
            startup_delay: Delay before detection to allow I2C bus to settle

        Returns:
            True if initialization successful, False otherwise
        """
        try:
            # If no driver set, try auto-detection (for "auto" or unknown types)
            if self._driver is None:
                ups_type = config.UPS_TYPE.lower()

                if ups_type == "none":
                    # Explicitly disabled
                    print("UPS disabled by configuration (UPS_TYPE=none)")
                    self._status.is_available = False
                    self._status.error_message = "UPS disabled"
                    return False

                # Small delay to ensure I2C bus is ready after boot
                if startup_delay > 0:
                    await asyncio.sleep(startup_delay)

                # Run auto-detection
                print("Auto-detecting UPS hardware...")
                driver, message = await auto_detect_driver(
                    pisugar_host=config.UPS_PISUGAR_HOST,
                    pisugar_port=config.UPS_PISUGAR_PORT
                )

                if driver:
                    self._driver = driver
                    print(message)
                else:
                    print(message)
                    self._status.is_available = False
                    self._status.error_message = message
                    return False

            # Initialize the driver
            success = await self._driver.initialize()

            if success:
                self._status.is_available = True
                self._status.error_message = None
                print(f"UPS initialized: {self._driver.get_model_name()}")
            else:
                self._status.is_available = False
                self._status.error_message = f"Failed to initialize {self._driver.get_model_name()}"

            return success

        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

    def get_power_restrictions(self) -> Dict[str, Any]:
        """Get current power restrictions based on hardware state.

        Returns:
            Dict with restriction info including:
            - restricted: bool - Whether operations are restricted
            - reason: Optional[str] - Why operations are restricted
            - power_source: str - Current power source
            - ups_available: bool - Whether UPS hardware is present
            - battery_percentage: Optional[float] - Current battery level
            - allow_firmware_flash: bool - Can flash firmware (fullimage)
            - allow_bootloader_flash: bool - Can flash bootloader
            - allow_intensive_operations: bool - Can run intensive ops
            - message: Optional[str] - User-friendly message
            - warning: Optional[str] - Warning message
            - shutdown_imminent: Optional[bool] - Shutdown about to happen
        """
        # UPS not available = assume AC power, no restrictions
        if not self._status.is_available:
            return {
                "restricted": False,
                "reason": None,
                "power_source": "assumed_ac",
                "ups_available": False,
                "battery_percentage": None,
                "allow_firmware_flash": True,
                "allow_bootloader_flash": True,
                "allow_intensive_operations": True,
                "message": "UPS not detected. Assuming stable AC power. Ensure power stability before firmware operations."
            }

        # UPS available - check power source and battery state
        if self._status.power_source == PowerSource.AC:
            return {
                "restricted": False,
                "reason": None,
                "power_source": "ac",
                "ups_available": True,
                "battery_percentage": self._status.battery_percentage,
                "allow_firmware_flash": True,
                "allow_bootloader_flash": True,
                "allow_intensive_operations": True,
                "message": "On AC power. All operations allowed."
            }

        # On battery power - apply restrictions based on battery level
        battery_pct = self._status.battery_percentage

        if battery_pct >= 80:
            return {
                "restricted": False,
                "reason": None,
                "power_source": "battery",
                "ups_available": True,
                "battery_percentage": battery_pct,
                "allow_firmware_flash": True,
                "allow_bootloader_flash": True,
                "allow_intensive_operations": True,
                "warning": "On battery power. Bootloader flashing not recommended below 80%."
            }
        elif battery_pct >= 50:
            return {
                "restricted": True,
                "reason": "Battery level below 80%",
                "power_source": "battery",
                "ups_available": True,
                "battery_percentage": battery_pct,
                "allow_firmware_flash": True,   # Fullimage only
                "allow_bootloader_flash": False,
                "allow_intensive_operations": True,
                "message": "Bootloader flashing disabled. Battery too low (minimum 80% required)."
            }
        elif battery_pct >= 20:
            return {
                "restricted": True,
                "reason": "Battery level below 50%",
                "power_source": "battery",
                "ups_available": True,
                "battery_percentage": battery_pct,
                "allow_firmware_flash": False,
                "allow_bootloader_flash": False,
                "allow_intensive_operations": True,
                "message": "Firmware flashing disabled. Battery too low (minimum 50% required)."
            }
        elif battery_pct >= 10:
            return {
                "restricted": True,
                "reason": "Battery critically low",
                "power_source": "battery",
                "ups_available": True,
                "battery_percentage": battery_pct,
                "allow_firmware_flash": False,
                "allow_bootloader_flash": False,
                "allow_intensive_operations": False,
                "message": "Critical battery level. Read-only mode active."
            }
        else:
            return {
                "restricted": True,
                "reason": "Battery emergency level",
                "power_source": "battery",
                "ups_available": True,
                "battery_percentage": battery_pct,
                "allow_firmware_flash": False,
                "allow_bootloader_flash": False,
                "allow_intensive_operations": False,
                "message": "Emergency battery level. System will shutdown soon.",
                "shutdown_imminent": True
            }

    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 UPS hardware."""
        if not self._status.is_available:
            return

        try:
            # Read data from driver
            data = await self._driver.read_data()

            # Check for misconfigured UPS (all values are zero) - only notify once
            if data.percentage == 0.0 and data.voltage == 0.0 and data.current == 0.0:
                if not self._config_warning_sent:
                    self._config_warning_sent = True
                    model_name = self._driver.get_model_name() if self._driver else "Unknown"
                    await self._trigger_event("ups_config_required", {
                        "model": model_name,
                        "message": f"UPS '{model_name}' detected but returning no data. May need reconfiguration.",
                        "hint": "Run 'sudo dpkg-reconfigure pisugar-server' to select the correct PiSugar model."
                    })
            else:
                # Reset flag when we get valid data
                self._config_warning_sent = False

            # Update status with new data
            self._status.battery_percentage = data.percentage
            self._status.voltage = data.voltage
            self._status.current = data.current
            self._status.temperature = data.temperature

            # Calculate time_remaining if not provided by driver
            if data.time_remaining is not None:
                self._status.time_remaining = data.time_remaining
            elif data.current < 0 and data.percentage > 0:
                # Discharging - estimate time remaining
                # current is in Amps (negative when discharging)
                draw_ma = abs(data.current * 1000)
                if draw_ma > 10:  # Only calculate if meaningful draw
                    remaining_mah = self._battery_capacity_mah * (data.percentage / 100)
                    hours_remaining = remaining_mah / draw_ma
                    self._status.time_remaining = int(hours_remaining * 60)  # Minutes
                else:
                    self._status.time_remaining = None
            else:
                self._status.time_remaining = None

            # Determine power source and battery status from driver data
            if data.is_charging:
                self._status.power_source = PowerSource.AC
                if data.percentage >= 99.0:
                    self._status.battery_status = BatteryStatus.FULL
                else:
                    self._status.battery_status = BatteryStatus.CHARGING
            else:
                self._status.power_source = PowerSource.BATTERY
                if data.percentage < self._critical_threshold:
                    self._status.battery_status = BatteryStatus.CRITICAL
                else:
                    self._status.battery_status = BatteryStatus.DISCHARGING

            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 _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 UPS driver connection."""
        if self._driver:
            self._driver.close()


# 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