OTA¶

class OTAPlugin(PluginBase, OTAMixin):
    """OTA plugin implementation.

    This plugin integrates the OTA update workflow into the Tolomeo
    plugin architecture, combining the `PluginBase` interface with
    `OTAMixin` for update state management. It provides real download,
    installation, and lifecycle coordination behavior suitable for
    deployment environments.

    The class binds OTA state transitions to task manager events and
    ensures consistent synchronization between OTA status and task
    execution state.

    """

    class Meta:
        """Plugin metadata container defining available OTA commands and options."""

        commands = OTA_COMMANDS
        config_model = OTAOptions

    def __init__(self, logger: logging.Logger | None = None) -> None:
        """Initialize the OTA plugin instance.

        Combines `PluginBase` and `OTAMixin` initializers, establishing
        logging and OTA state tracking behavior.

        Args:
            logger: Optional logger instance for OTA operations. If None, a default logger is created.

        """
        PluginBase.__init__(self, logger)
        OTAMixin.__init__(self, logger)

    async def before_setup(self) -> None:
        """Configure OTA-related callbacks prior to plugin setup.

        This hook prepares the OTA state machine and task manager by
        registering event hooks for download errors and task state
        transitions. It chains existing task manager callbacks with
        OTA-specific handlers to ensure OTA state synchronization. It also
        resolves the checkpoint file location and `OTAOptions.allow_same_version`
        from `self.options` (readable now that `connect` has assigned `self.id`)
        and assigns them to the state machine, which arms bootstrapping without
        performing it — `after_setup` performs the actual bootstrap.

        Raises:
            Exception: If callback setup fails or task manager references
                are invalid.

        """
        self._logger.debug("Setting up push callback for state machine")
        self._ota_fsm.push_callback = self._push_data_cb

        self._ota_fsm.checkpoint_file = self.options.checkpoint_file
        self._logger.info("OTA checkpoint file: %s", self._ota_fsm.checkpoint_file)
        self._ota_fsm.allow_same_version = self.options.allow_same_version

        def download_error_hook(task_name: str, exception: Exception) -> None:
            """Internal hook to handle download task errors."""
            if task_name == self._download_task_name:
                self._on_download_error(task_name, exception)

        def download_state_change_hook(task_name: str, old_state: str, new_state: str) -> None:
            """Internal hook to handle task state transitions for downloads."""
            if task_name == self._download_task_name:
                self._on_download_state_change(task_name, old_state, new_state)

        tm = self._task_manager
        tm._on_error = chain_callback(tm._on_error, download_error_hook)
        tm._on_state_change = chain_callback(tm._on_state_change, download_state_change_hook)

    async def after_setup(self) -> None:
        """Bootstrap OTA state after plugin setup.

        This hook is the single authoritative bootstrap for the OTA state
        machine: it calls `reset`, which resolves to whatever the checkpoint
        says — `updated`, `paused`, `update_available`, or any other
        checkpointable state — rather than assuming a fixed starting state.

        """
        await self._bootstrap()

class OTAMixin(metaclass=ABCMeta):
    """Mixin providing OTA (Over-The-Air) update lifecycle management.

    Uses `ABCMeta` so ``install_update`` and ``monitor_update`` are enforced. Both
    bodies are ``...``, so an unenforced marker means a plugin that forgot one
    instantiates cleanly and silently no-ops the firmware install.

    """

    _ota_fsm: OTAStateMachine
    _logger: logging.Logger
    _task_manager: TaskManager
    _push_data_cb: Callable | None
    _download_task_name: str = DEFAULT_DOWNLOAD_TASK_NAME

    def __init__(self, logger: logging.Logger | None = None) -> None:
        """Create the OTA state machine; the logger and callbacks belong to `PluginBase`.

        Args:
            logger: Optional logger for the state machine. If None, it uses its
                own module logger.

        """
        self._ota_fsm = OTAStateMachine(logger=logger, defer_bootstrap=True)

    # --- Properties ---

    @property
    def state(self) -> str:
        """Get the current OTA state.

        Returns:
            The string representation of the current OTA state
            (e.g., "downloading", "paused", "update_ready").

        """
        return str(self._ota_fsm.current_state.value)

    @property
    def fsm(self) -> OTAStateMachine:
        """Get the OTA state machine instance."""
        return self._ota_fsm

    # --- Supplied by the Host Class ---

    @property
    def options(self) -> OTAOptions:
        """Declared, not implemented: `PluginBase` provides this.

        Narrowed to `OTAOptions` here so the mixin's own reads are typed. The
        host class resolves it against whatever model it declares — `OTAOptions`
        or an extension of it.

        """
        raise NotImplementedError("OTAMixin must be combined with PluginBase, which supplies options")

    # --- Abstract Methods ---

    @abstractmethod
    async def install_update(self) -> None:
        """Install the downloaded OTA update.

        This method must be implemented by subclasses to perform the
        actual update process (e.g., applying firmware, rebooting,
        or verifying installation integrity).

        """
        ...

    @abstractmethod
    async def monitor_update(self) -> None:
        """Monitor the OTA update progress.

        This method must be implemented by subclasses to periodically
        check update progress, collect metrics, or detect installation
        anomalies.

        """
        ...

    # --- Protected State Transition Methods ---

    async def _bootstrap(self) -> None:
        """Bootstrap the OTA state machine to the initial state."""
        self._ota_fsm.reset()
        self._logger.info("State machine bootstrapped")

    # --- Task Manager Callbacks ---

    def _on_download_error(self, task_name: str, exception: Exception) -> None:
        """Handle errors raised during download task execution.

        Args:
            task_name: The name of the download task that failed
            exception: The exception raised during execution

        """
        self._logger.error("Download task %s failed: %s", task_name, exception)
        self._ota_fsm.download_fail()

    def _on_download_state_change(self, task_name: str, old_state: str, new_state: str) -> None:
        """Handle state changes from the task manager for the download task.

        Synchronizes the OTA state machine with the task manager's internal
        task state transitions.

        Args:
            task_name: The name of the download task
            old_state: The previous task state
            new_state: The new task state after transition

        """
        self._logger.info("Download task '%s' state changed: %s -> %s", task_name, old_state, new_state)

        if new_state == "Completed":
            self._ota_fsm.download_success()
        elif new_state == "Failed":
            self._ota_fsm.download_fail()
        elif new_state == "Paused":
            if self.state == "downloading":
                self._ota_fsm.pause()
        # Handle resume after pause or power interruption
        elif new_state == "Running" and old_state in {"Pending", "Paused"}:
            if self.state == "paused":
                self._ota_fsm.resume()

    # --- Public Methods ---

    # --- Notification Phase ---
    async def notify_update(
        self, fw_name: str, fw_version: str, url: str | None = None, dest: Path | str | None = None
    ) -> bool:
        """Notify the system that an update is available.

        Sets the pending update information and transitions the OTA state
        machine to `update_available`, unless ``fw_version`` would roll the
        installed firmware back: an older version is ignored, and so is the same
        version unless `OTAOptions.allow_same_version` is set. An ignored update
        leaves both the state and the pending update untouched.

        Args:
            fw_name: Name of the firmware update
            fw_version: Version of the firmware update, in Semantic Versioning 2.0.0
            url: Optional URL for the update package
            dest: Optional destination path for the download

        Returns:
            True if the update was accepted, False if it was ignored.

        Raises:
            ValueError: If ``fw_version`` is not a valid Semantic Versioning 2.0.0 string

        """
        if not self._ota_fsm.accepts_version(fw_version):
            self._logger.warning(
                "Ignoring update %s %s: installed firmware is %s",
                fw_name,
                fw_version,
                self._ota_fsm.current_firmware.version,
            )
            return False

        # Set pending update information
        self._ota_fsm.set_pending_update(
            name=fw_name, version=fw_version, url=url, download_path=dest if dest else self.options.download_dest
        )

        self._logger.info("Update available: %s %s from %s", fw_name, fw_version, url or "unspecified URL")

        self._ota_fsm.notify()
        self._logger.info("Update notification processed")
        return True

    # --- Download Phase ---

    async def start_download(self, url: str | None = None, dest: Path | str | None = None) -> None:
        """Start downloading the OTA update package.

        This method coordinates with the `TaskManager` to create and start
        a download task for the specified URL. It automatically manages
        existing tasks, retries, and OTA state synchronization.

        Args:
            url: Optional download URL. If None, uses the pending update URL
            dest: Optional destination path for saving the downloaded file

        Raises:
            ValueError: If the download URL is missing
            RuntimeError: If the task manager is not initialized
            Exception: If task creation or execution fails

        """
        # Update pending update info if provided
        if url or dest:
            self._ota_fsm.set_pending_update(url=url, download_path=dest)

        # Validate download URL
        download_url = self._ota_fsm.pending_update.url
        if not download_url:
            self._logger.error("Cannot start download: URL not provided")
            raise ValueError("Download URL is required")

        # Validate/set download destination
        download_path = self._ota_fsm.pending_update.download_path
        if not download_path:
            self._ota_fsm.set_pending_update(download_path=self.options.download_dest)
            download_path = str(self.options.download_dest)
            self._logger.info("Using default download destination: %s", download_path)

        # Clean up existing download task if present
        await self._ensure_clean_download_task()

        # Transition to downloading state
        self._ota_fsm.download()
        self._logger.info("Download state transition initiated")

        # Create and start download task
        try:
            await self._create_download_task(download_url, Path(download_path))
        except Exception as e:
            self._logger.error("Failed to start download: %s", e)
            self._ota_fsm.download_fail()
            raise

    async def pause_download(self) -> None:
        """Pause the ongoing OTA download.

        Requests the `TaskManager` to pause the active download task.
        The OTA state machine will automatically transition to `paused`
        through the state change callback.

        Raises:
            RuntimeError: If the task manager is not initialized
            Exception: If the pause operation fails

        """
        try:
            if self._download_task_name not in self._task_manager.tasks:
                self._logger.warning("No active download task to pause")
                return

            task_state = self._task_manager.task_state(self._download_task_name)
            if task_state != "Running":
                self._logger.warning("Download task is not running (state: %s)", task_state)
                return

            await self._task_manager.pause_task(self._download_task_name)
            self._logger.info("Download pause requested")

        except Exception as e:
            self._logger.error("Failed to pause download: %s", e)
            raise

    async def resume_download(self) -> None:
        """Resume a previously paused OTA download.

        Requests the `TaskManager` to resume the paused download task.
        If no task exists, creates a new one using the pending update info.
        The OTA state machine will transition to `downloading` through
        the state change callback.

        Raises:
            RuntimeError: If the task manager is not initialized
            ValueError: If download URL is missing
            Exception: If the resume operation fails

        """
        try:
            # If no task exists, create a new one
            if self._download_task_name not in self._task_manager.tasks:
                self._logger.warning("No download task to resume, creating new task")

                download_url = self._ota_fsm.pending_update.url
                if not download_url:
                    self._logger.error("Cannot resume download: URL not available")
                    raise ValueError("Download URL is required")

                download_path = self._ota_fsm.pending_update.download_path
                if not download_path:
                    self._ota_fsm.set_pending_update(download_path=self.options.download_dest)
                    download_path = str(self.options.download_dest)
                    self._logger.info("Using default download destination: %s", download_path)

                await self._create_download_task(download_url, Path(download_path))
                return

            # Resume existing paused task
            task_state = self._task_manager.task_state(self._download_task_name)
            if task_state != "Paused":
                self._logger.warning("Download task is not paused (state: %s)", task_state)
                return

            await self._task_manager.start_task(self._download_task_name)
            self._logger.info("Download resume requested")

        except Exception as e:
            self._logger.error("Failed to resume download: %s", e)
            self._ota_fsm.download_fail()
            raise

    async def abort_update(self) -> None:
        """Abort a failed OTA update and restart from the beginning.

        Cleans up any existing download task or partial file, and
        transitions back to `update_available`.

        Raises:
            Exception: If cleanup or state reset fails

        """
        try:
            await self._cleanup_download()
            self._ota_fsm.abort()
            self._logger.info("Update aborted")
        except Exception as e:
            self._logger.error("Failed to abort update: %s", e)
            raise

    async def cancel_update(self) -> None:
        """Cancel an in-flight update and clean up its partial download.

        Raises:
            Exception: Whatever the cleanup or the transition raised, after
                logging it.

        """
        try:
            await self._cleanup_download()
            self._ota_fsm.cancel()
            self._logger.info("Update cancelled")
        except Exception as e:
            self._logger.error("Failed to cancel update: %s", e)
            raise

    async def clear_update(self) -> None:
        """Discard the pending update and clean up its partial download.

        Raises:
            Exception: Whatever the cleanup or the transition raised, after
                logging it.

        """
        try:
            await self._cleanup_download()
            self._ota_fsm.clear_pending_update()
            self._ota_fsm.clear()
            self._logger.info("Pending update cleared")
        except Exception as e:
            self._logger.error("Failed to clear pending update: %s", e)
            raise

    async def reset_state(self) -> None:
        """Reset the OTA state machine to a clean baseline.

        Cleans up any in-flight download, clears the pending update, removes the
        persisted checkpoint, and drives the FSM through ``reset()`` back to
        ``bootstrapping``. With no checkpoint present, re-detection resolves to
        ``updated``.

        Raises:
            Exception: If cleanup or the reset transition fails.

        """
        try:
            await self._cleanup_download()
            self._ota_fsm.clear_pending_update()
            self._ota_fsm.clear_checkpoint()
            self._ota_fsm.reset()
            self._logger.info("OTA state reset")
        except Exception as e:
            self._logger.error("Failed to reset state: %s", e)
            raise

    # --- System Control ---

    async def reboot_system(self) -> None:
        """Request a system reboot using systemd via D-Bus.

        Attempts to start the `systemd-reboot.service` unit using the
        system D-Bus. This method is typically called after a successful
        update to reboot into the new firmware.

        Raises:
            Exception: If the D-Bus call fails or the service cannot be started

        """
        try:
            bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
            introspection = await bus.introspect("org.freedesktop.systemd1", "/org/freedesktop/systemd1")
            obj = bus.get_proxy_object("org.freedesktop.systemd1", "/org/freedesktop/systemd1", introspection)
            manager = obj.get_interface("org.freedesktop.systemd1.Manager")

            await self._safe_call_method(manager, "call_start_unit", SYSTEMD_REBOOT_SERVICE, "replace")
            self._logger.info("Service '%s' start requested", SYSTEMD_REBOOT_SERVICE)
        except Exception as e:
            self._logger.error("Failed to start service '%s': %s", SYSTEMD_REBOOT_SERVICE, e)
            raise

    # --- Private Helper Methods ---

    async def _ensure_clean_download_task(self) -> None:
        """Ensure no existing download task is present.

        If a download task exists, stops and removes it before proceeding.

        """
        if self._download_task_name in self._task_manager.tasks:
            task_state = self._task_manager.task_state(self._download_task_name)
            self._logger.warning("Download task already exists in state: %s. Stopping and recreating.", task_state)
            await self._task_manager.stop_task(self._download_task_name)
            self._logger.debug("Stopped existing download task")
            await self._task_manager.remove_task(self._download_task_name)
            self._logger.debug("Removed existing download task")

    async def _create_download_task(self, url: str, dest: Path) -> None:
        """Create and start a new download task.

        Args:
            url: The URL to download from
            dest: The destination path for the downloaded file

        Raises:
            Exception: If task creation or start fails

        """
        download_task = DownloadTask(url=url, dest=dest)
        self._logger.debug("Created DownloadTask instance")

        await self._task_manager.add_task(
            name=self._download_task_name,
            coro_func=download_task.run,
            auto_restart=False,
            max_retries=3,
            retry_delay=2.0,
        )
        self._logger.info("Download task started for %s", url)

    async def _cleanup_download(self) -> None:
        """Clean up download task and partial file.

        Stops and removes the download task if it exists, and deletes
        any partial download file.

        Raises:
            Exception: If cleanup operations fail

        """
        if self._download_task_name in self._task_manager.tasks:
            await self._task_manager.stop_task(self._download_task_name)
            await self._task_manager.remove_task(self._download_task_name)
            self._logger.debug("Removed download task")

        download_path = self._ota_fsm.pending_update.download_path
        if download_path:
            dest_path = Path(download_path)
            if dest_path.exists():
                dest_path.unlink()
                self._logger.debug("Removed existing file at: %s", dest_path)

    async def _safe_call_method(self, interface: Any, method_name: str, *args: Any, **kwargs: Any) -> Any:
        """Safely call a D-Bus method with validation and error handling.

        Validates the target method's existence, callability, and coroutine
        nature before invoking it asynchronously.

        Args:
            interface: The D-Bus interface containing the target method
            method_name: The name of the method to call
            *args: Positional arguments to pass to the method
            **kwargs: Keyword arguments to pass to the method

        Returns:
            The return value from the D-Bus method call, or None if validation fails

        Raises:
            Exception: If the method call fails during execution

        """
        method = getattr(interface, method_name, None)
        if method is None:
            self._logger.warning("Method '%s' not found on interface", method_name)
            return
        if not callable(method):
            self._logger.warning("Attribute '%s' is not callable", method_name)
            return
        if not inspect.iscoroutinefunction(method):
            self._logger.warning("Method '%s' is not an async coroutine", method_name)
            return

        return await method(*args, **kwargs)