pi-pm3/app/backend/services/system_service.py

"""System Service Layer.

This service provides system operations (shutdown, restart, info) that can be
consumed by multiple interfaces (REST API, BLE GATT, etc.).
"""
import asyncio
import os
import platform
import psutil
import shutil
from dataclasses import dataclass
from typing import Optional, Dict, Any
from pathlib import Path

from .pm3_service import PM3ServiceError, PM3ServiceResult


@dataclass
class CPUCoreInfo:
    """Per-core CPU information."""
    core_id: int
    percent: float
    online: bool = True  # Whether this core is currently online


@dataclass
class PiModelInfo:
    """Raspberry Pi model information."""
    model: str  # Full model string (e.g., "Raspberry Pi Zero 2 W Rev 1.0")
    model_short: str  # Short name (e.g., "Pi Zero 2 W")
    total_cores: int  # Total physical cores
    default_active_cores: int  # Recommended default active cores
    min_cores: int  # Minimum cores (always 1)
    max_cores: int  # Maximum configurable cores


@dataclass
class CPUCoresConfig:
    """CPU cores configuration."""
    total_cores: int
    online_cores: int
    configurable_cores: list  # List of core IDs that can be toggled (usually all except 0)
    pi_model: Optional[PiModelInfo] = None


@dataclass
class SystemInfo:
    """System information data."""
    hostname: str
    platform: str
    platform_version: str
    cpu_count: int
    cpu_percent: float
    cpu_per_core: list  # List of CPUCoreInfo
    memory_total: int
    memory_used: int
    memory_percent: float
    disk_total: int
    disk_used: int
    disk_percent: float
    uptime: float
    temperature: Optional[float] = None
    load_average: Optional[tuple] = None  # 1, 5, 15 min load averages


class SystemService:
    """System operations service.

    Provides reusable business logic for:
    - System information queries
    - Shutdown/restart operations
    - Service log retrieval
    """

    def __init__(self):
        """Initialize system service."""
        pass

    async def get_info(self) -> PM3ServiceResult:
        """Get system information.

        Returns:
            PM3ServiceResult with system info (CPU, memory, disk, etc.)
        """
        try:
            # Get CPU usage (blocking call, run in executor)
            loop = asyncio.get_event_loop()

            # Get per-core CPU percentages (requires percpu=True)
            # First call to initialize, then wait briefly for measurement
            await loop.run_in_executor(None, lambda: psutil.cpu_percent(percpu=True))
            await asyncio.sleep(0.1)  # Brief pause for accurate measurement
            cpu_per_core_raw = await loop.run_in_executor(
                None,
                lambda: psutil.cpu_percent(percpu=True)
            )

            # Check which cores are online
            total_cores = psutil.cpu_count(logical=True) or 1
            core_online_status = {}
            for i in range(total_cores):
                online_file = Path(f"/sys/devices/system/cpu/cpu{i}/online")
                if i == 0:
                    # Core 0 is always online
                    core_online_status[i] = True
                elif online_file.exists():
                    core_online_status[i] = online_file.read_text().strip() == "1"
                else:
                    # File doesn't exist, core is always on
                    core_online_status[i] = True

            # Build per-core info with online status
            cpu_per_core = [
                CPUCoreInfo(
                    core_id=i,
                    percent=pct if core_online_status.get(i, True) else 0.0,
                    online=core_online_status.get(i, True)
                )
                for i, pct in enumerate(cpu_per_core_raw)
            ]

            # Overall CPU percentage
            cpu_percent = sum(cpu_per_core_raw) / len(cpu_per_core_raw) if cpu_per_core_raw else 0.0

            # Get load average (Unix only)
            try:
                load_average = os.getloadavg()
            except (OSError, AttributeError):
                load_average = None

            # Get memory info
            memory = psutil.virtual_memory()

            # Get disk info for root partition
            disk = psutil.disk_usage('/')

            # Get system uptime
            boot_time = psutil.boot_time()
            uptime = psutil.time.time() - boot_time

            # Try to get CPU temperature (Raspberry Pi specific)
            temperature = await self._get_cpu_temperature()

            system_info = SystemInfo(
                hostname=platform.node(),
                platform=platform.system(),
                platform_version=platform.release(),
                cpu_count=psutil.cpu_count(),
                cpu_percent=cpu_percent,
                cpu_per_core=cpu_per_core,
                memory_total=memory.total,
                memory_used=memory.used,
                memory_percent=memory.percent,
                disk_total=disk.total,
                disk_used=disk.used,
                disk_percent=disk.percent,
                uptime=uptime,
                temperature=temperature,
                load_average=load_average
            )

            return PM3ServiceResult(
                success=True,
                data={
                    "hostname": system_info.hostname,
                    "platform": system_info.platform,
                    "platform_version": system_info.platform_version,
                    "cpu": {
                        "count": system_info.cpu_count,
                        "percent": system_info.cpu_percent,
                        "temperature": system_info.temperature,
                        "per_core": [
                            {"core_id": c.core_id, "percent": c.percent, "online": c.online}
                            for c in system_info.cpu_per_core
                        ],
                        "load_average": list(system_info.load_average) if system_info.load_average else None
                    },
                    "memory": {
                        "total": system_info.memory_total,
                        "used": system_info.memory_used,
                        "percent": system_info.memory_percent
                    },
                    "disk": {
                        "total": system_info.disk_total,
                        "used": system_info.disk_used,
                        "percent": system_info.disk_percent
                    },
                    "uptime": system_info.uptime
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="system_info_error",
                    message="Failed to get system information",
                    details=str(e)
                )
            )

    async def _get_cpu_temperature(self) -> Optional[float]:
        """Get CPU temperature (Raspberry Pi specific).

        Returns:
            Temperature in Celsius or None if unavailable
        """
        try:
            temp_file = Path("/sys/class/thermal/thermal_zone0/temp")
            if temp_file.exists():
                temp_str = temp_file.read_text().strip()
                # Temperature is in millidegrees
                return float(temp_str) / 1000.0
        except Exception:
            pass
        return None

    async def shutdown(self, delay: int = 0) -> PM3ServiceResult:
        """Initiate system shutdown.

        Args:
            delay: Delay in seconds before shutdown (default: 0)

        Returns:
            PM3ServiceResult indicating success/failure
        """
        try:
            if delay > 0:
                # Schedule shutdown
                await self._run_command(f"sudo shutdown -h +{delay // 60}")
                return PM3ServiceResult(
                    success=True,
                    data={
                        "message": f"Shutdown scheduled in {delay} seconds",
                        "delay": delay
                    }
                )
            else:
                # Immediate shutdown
                await self._run_command("sudo shutdown -h now")
                return PM3ServiceResult(
                    success=True,
                    data={"message": "Shutdown initiated"}
                )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="shutdown_error",
                    message="Failed to initiate shutdown",
                    details=str(e)
                )
            )

    async def restart(self, delay: int = 0) -> PM3ServiceResult:
        """Initiate system restart.

        Args:
            delay: Delay in seconds before restart (default: 0)

        Returns:
            PM3ServiceResult indicating success/failure
        """
        try:
            if delay > 0:
                # Schedule restart
                await self._run_command(f"sudo shutdown -r +{delay // 60}")
                return PM3ServiceResult(
                    success=True,
                    data={
                        "message": f"Restart scheduled in {delay} seconds",
                        "delay": delay
                    }
                )
            else:
                # Immediate restart
                await self._run_command("sudo shutdown -r now")
                return PM3ServiceResult(
                    success=True,
                    data={"message": "Restart initiated"}
                )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="restart_error",
                    message="Failed to initiate restart",
                    details=str(e)
                )
            )

    async def cancel_shutdown(self) -> PM3ServiceResult:
        """Cancel a scheduled shutdown or restart.

        Returns:
            PM3ServiceResult indicating success/failure
        """
        try:
            await self._run_command("sudo shutdown -c")
            return PM3ServiceResult(
                success=True,
                data={"message": "Shutdown/restart cancelled"}
            )
        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="cancel_shutdown_error",
                    message="Failed to cancel shutdown",
                    details=str(e)
                )
            )

    async def get_logs(
        self,
        service: str = "dangerous-pi",
        lines: int = 100
    ) -> PM3ServiceResult:
        """Get service logs using journalctl.

        Args:
            service: Service name to get logs for
            lines: Number of lines to retrieve (default: 100)

        Returns:
            PM3ServiceResult with log content
        """
        try:
            output = await self._run_command(
                f"sudo journalctl -u {service} -n {lines} --no-pager"
            )

            return PM3ServiceResult(
                success=True,
                data={
                    "service": service,
                    "lines": lines,
                    "logs": output
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="logs_error",
                    message=f"Failed to get logs for service '{service}'",
                    details=str(e)
                )
            )

    async def get_service_status(self, service: str) -> PM3ServiceResult:
        """Get systemd service status.

        Args:
            service: Service name

        Returns:
            PM3ServiceResult with service status
        """
        try:
            output = await self._run_command(
                f"sudo systemctl status {service}",
                check=False  # Don't raise on non-zero exit
            )

            # Parse status
            is_active = "active (running)" in output.lower()
            is_enabled = await self._is_service_enabled(service)

            return PM3ServiceResult(
                success=True,
                data={
                    "service": service,
                    "active": is_active,
                    "enabled": is_enabled,
                    "status": output
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="service_status_error",
                    message=f"Failed to get status for service '{service}'",
                    details=str(e)
                )
            )

    async def _is_service_enabled(self, service: str) -> bool:
        """Check if a systemd service is enabled.

        Args:
            service: Service name

        Returns:
            True if service is enabled
        """
        try:
            output = await self._run_command(
                f"sudo systemctl is-enabled {service}",
                check=False
            )
            return output.strip() == "enabled"
        except Exception:
            return False

    async def _run_command(
        self,
        command: str,
        check: bool = True
    ) -> str:
        """Run a shell command asynchronously.

        Args:
            command: Command to run
            check: Raise exception on non-zero exit code

        Returns:
            Command output

        Raises:
            RuntimeError: If command fails and check=True
        """
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        stdout, stderr = await process.communicate()

        if check and process.returncode != 0:
            raise RuntimeError(
                f"Command failed with exit code {process.returncode}: "
                f"{stderr.decode()}"
            )

        return stdout.decode()

    # Pi model defaults: maps model substring to (short_name, default_active_cores, physical_cores)
    # default_active_cores: recommended cores for thermal/power management
    # physical_cores: actual hardware cores regardless of boot config
    PI_MODEL_DEFAULTS = {
        "Pi Zero 2": ("Pi Zero 2 W", 2, 4),   # 4 cores, default to 2 for thermal
        "Pi Zero W": ("Pi Zero W", 1, 1),     # Single core
        "Pi Zero": ("Pi Zero", 1, 1),         # Single core
        "Pi 4": ("Pi 4", 4, 4),
        "Pi 3": ("Pi 3", 4, 4),
        "Pi 2": ("Pi 2", 4, 4),
        "Pi 5": ("Pi 5", 4, 4),
    }

    def _get_physical_cores(self) -> int:
        """Get the actual physical core count (not limited by maxcpus).

        Reads from /sys/devices/system/cpu/possible to get the full range
        of possible CPUs regardless of boot-time restrictions.

        Returns:
            Physical core count
        """
        try:
            # /sys/devices/system/cpu/possible contains the full range like "0-3" for 4 cores
            possible_file = Path("/sys/devices/system/cpu/possible")
            if possible_file.exists():
                content = possible_file.read_text().strip()
                # Parse "0-3" format to get count of 4
                if '-' in content:
                    start, end = content.split('-')
                    return int(end) - int(start) + 1
                elif content.isdigit():
                    return int(content) + 1
        except Exception:
            pass

        # Fallback to psutil (may be limited by maxcpus)
        return psutil.cpu_count(logical=True) or 1

    async def get_pi_model(self) -> PM3ServiceResult:
        """Detect Raspberry Pi model.

        Reads from /proc/device-tree/model or /proc/cpuinfo.

        Returns:
            PM3ServiceResult with PiModelInfo data
        """
        try:
            model_str = "Unknown"
            model_short = "Unknown"
            physical_cores = self._get_physical_cores()
            total_cores = physical_cores  # Use physical cores as total
            default_cores = total_cores

            # Try device tree first (most reliable)
            model_file = Path("/proc/device-tree/model")
            if model_file.exists():
                model_str = model_file.read_text().strip().rstrip('\x00')
            else:
                # Fallback to cpuinfo
                cpuinfo = Path("/proc/cpuinfo")
                if cpuinfo.exists():
                    for line in cpuinfo.read_text().split('\n'):
                        if line.startswith("Model"):
                            model_str = line.split(':')[1].strip()
                            break

            # Determine short name and default/physical cores based on model
            for pattern, (short_name, def_cores, phys_cores) in self.PI_MODEL_DEFAULTS.items():
                if pattern in model_str:
                    model_short = short_name
                    # Use the known physical cores from model defaults if detected
                    total_cores = max(physical_cores, phys_cores)
                    default_cores = min(def_cores, total_cores)
                    break
            else:
                # Not a known Pi model, use detected physical cores
                if "Raspberry" in model_str:
                    model_short = "Raspberry Pi"
                else:
                    model_short = model_str[:20] if len(model_str) > 20 else model_str

            pi_model = PiModelInfo(
                model=model_str,
                model_short=model_short,
                total_cores=total_cores,
                default_active_cores=default_cores,
                min_cores=1,
                max_cores=total_cores
            )

            return PM3ServiceResult(
                success=True,
                data={
                    "model": pi_model.model,
                    "model_short": pi_model.model_short,
                    "total_cores": pi_model.total_cores,
                    "default_active_cores": pi_model.default_active_cores,
                    "min_cores": pi_model.min_cores,
                    "max_cores": pi_model.max_cores
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="pi_model_error",
                    message="Failed to detect Pi model",
                    details=str(e)
                )
            )

    async def get_cpu_cores_config(self) -> PM3ServiceResult:
        """Get CPU cores configuration.

        Returns info about total cores, online cores, configured cores (from cmdline.txt),
        and Pi model with defaults.

        Returns:
            PM3ServiceResult with CPUCoresConfig data
        """
        try:
            # Get physical core count (not limited by maxcpus)
            physical_cores = self._get_physical_cores()

            # Get currently online cores (runtime state)
            # On Pi, we can use nproc or count from /proc/cpuinfo
            try:
                import subprocess
                result = subprocess.run(
                    ["nproc"],
                    capture_output=True,
                    text=True,
                    timeout=5
                )
                online_cores = int(result.stdout.strip()) if result.returncode == 0 else physical_cores
            except Exception:
                online_cores = physical_cores

            # Get configured cores from cmdline.txt (boot-time setting)
            configured_cores = self._get_configured_maxcpus()

            # Use physical cores as total (what the hardware actually has)
            total_cores = physical_cores

            # Configurable cores are all cores except 0 (which is always on)
            configurable_cores = list(range(1, total_cores))

            # Get Pi model info
            model_result = await self.get_pi_model()
            pi_model_data = model_result.data if model_result.success else None

            return PM3ServiceResult(
                success=True,
                data={
                    "total_cores": total_cores,
                    "online_cores": online_cores,
                    "configured_cores": configured_cores,  # From cmdline.txt
                    "configurable_cores": configurable_cores,
                    "pi_model": pi_model_data
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="cpu_cores_error",
                    message="Failed to get CPU cores config",
                    details=str(e)
                )
            )

    # Possible locations for cmdline.txt (varies by Pi OS version)
    CMDLINE_PATHS = [
        Path("/boot/firmware/cmdline.txt"),  # Newer Pi OS (Bookworm+)
        Path("/boot/cmdline.txt"),            # Older Pi OS
    ]

    def _find_cmdline_path(self) -> Optional[Path]:
        """Find the cmdline.txt file location.

        Returns:
            Path to cmdline.txt or None if not found
        """
        for path in self.CMDLINE_PATHS:
            if path.exists():
                return path
        return None

    def _get_configured_maxcpus(self) -> Optional[int]:
        """Read the maxcpus value from cmdline.txt.

        Returns:
            Configured maxcpus value, or None if not set
        """
        cmdline_path = self._find_cmdline_path()
        if not cmdline_path:
            return None

        try:
            content = cmdline_path.read_text().strip()
            # Parse cmdline parameters
            for param in content.split():
                if param.startswith("maxcpus="):
                    return int(param.split("=")[1])
        except Exception:
            pass
        return None

    async def set_cpu_cores(
        self,
        num_cores: int,
        persist: bool = True,
        reboot: bool = True
    ) -> PM3ServiceResult:
        """Set the number of active CPU cores via kernel parameter.

        On Pi Zero 2 W (and other ARM systems), CPU hotplug via sysfs doesn't work.
        Instead, we modify the maxcpus=N kernel parameter in /boot/cmdline.txt.
        This requires a reboot to take effect.

        Args:
            num_cores: Number of cores to use (1 to total_cores)
            persist: Must be True (always persists to cmdline.txt)
            reboot: If True, reboot the system immediately

        Returns:
            PM3ServiceResult indicating success/failure
        """
        try:
            # Use physical cores (not limited by current maxcpus setting)
            total_cores = self._get_physical_cores()

            # Validate input
            if num_cores < 1:
                return PM3ServiceResult(
                    success=False,
                    error=PM3ServiceError(
                        code="invalid_cores",
                        message="Must have at least 1 core active",
                        details=f"Requested: {num_cores}"
                    )
                )

            if num_cores > total_cores:
                return PM3ServiceResult(
                    success=False,
                    error=PM3ServiceError(
                        code="invalid_cores",
                        message=f"Cannot exceed {total_cores} cores",
                        details=f"Requested: {num_cores}, Available: {total_cores}"
                    )
                )

            # Find cmdline.txt
            cmdline_path = self._find_cmdline_path()
            if not cmdline_path:
                return PM3ServiceResult(
                    success=False,
                    error=PM3ServiceError(
                        code="cmdline_not_found",
                        message="Could not find /boot/cmdline.txt or /boot/firmware/cmdline.txt",
                        details="This feature requires a Raspberry Pi with standard boot configuration"
                    )
                )

            # Update cmdline.txt with maxcpus parameter
            update_result = await self._update_cmdline_maxcpus(cmdline_path, num_cores)
            if not update_result.success:
                return update_result

            # Reboot if requested
            if reboot:
                await self._run_command("sudo shutdown -r +0", check=False)
                return PM3ServiceResult(
                    success=True,
                    data={
                        "message": f"CPU cores set to {num_cores}. Rebooting...",
                        "requested": num_cores,
                        "rebooting": True
                    }
                )

            return PM3ServiceResult(
                success=True,
                data={
                    "message": f"CPU cores set to {num_cores}. Reboot required to apply.",
                    "requested": num_cores,
                    "rebooting": False,
                    "reboot_required": True
                }
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="set_cores_error",
                    message="Failed to set CPU cores",
                    details=str(e)
                )
            )

    async def _update_cmdline_maxcpus(
        self,
        cmdline_path: Path,
        num_cores: int
    ) -> PM3ServiceResult:
        """Update the maxcpus parameter in cmdline.txt.

        Args:
            cmdline_path: Path to cmdline.txt
            num_cores: Number of cores to set

        Returns:
            PM3ServiceResult indicating success/failure
        """
        try:
            # Read current cmdline
            content = cmdline_path.read_text().strip()

            # Parse and update parameters
            params = content.split()
            new_params = []
            maxcpus_found = False

            for param in params:
                if param.startswith("maxcpus="):
                    # Replace existing maxcpus
                    new_params.append(f"maxcpus={num_cores}")
                    maxcpus_found = True
                else:
                    new_params.append(param)

            # Add maxcpus if not present
            if not maxcpus_found:
                new_params.append(f"maxcpus={num_cores}")

            new_content = " ".join(new_params)

            # Backup original
            await self._run_command(
                f"sudo cp {cmdline_path} {cmdline_path}.bak",
                check=True
            )

            # Write new cmdline.txt
            # Use a temp file and move to avoid corruption
            await self._run_command(
                f"echo '{new_content}' | sudo tee {cmdline_path}.new > /dev/null",
                check=True
            )
            await self._run_command(
                f"sudo mv {cmdline_path}.new {cmdline_path}",
                check=True
            )

            return PM3ServiceResult(
                success=True,
                data={"message": f"Updated {cmdline_path} with maxcpus={num_cores}"}
            )

        except Exception as e:
            return PM3ServiceResult(
                success=False,
                error=PM3ServiceError(
                    code="cmdline_update_error",
                    message="Failed to update cmdline.txt",
                    details=str(e)
                )
            )

    def get_configured_cpu_cores(self) -> Optional[int]:
        """Get the configured CPU cores from cmdline.txt.

        Returns:
            Configured maxcpus value, or None if not explicitly set
        """
        return self._get_configured_maxcpus()