Tasks¶

class TaskManager:
    """Asynchronous task manager with state tracking, pause/resume, and retries.

    This manager orchestrates the interaction between task contexts and their
    state machines, maintaining clean separation of concerns.

    """

    def __init__(
        self,
        logger: logging.Logger | None = None,
        on_error: Callable[[str, Exception], None] | None = None,
        on_state_change: Callable[[str, str, str], None] | None = None,
    ) -> None:
        """Initialize the task manager.

        Args:
            logger: Logger instance for task events. If None, creates a default logger.
            on_error: Optional callback function called when a task encounters an error. Receives task name
                and exception.
            on_state_change: Optional callback function called when a task changes state. Receives task
                name, source state, and target state.

        """
        self._logger: logging.Logger = logger or logging.getLogger(__name__)
        self._contexts: Dict[str, TaskContext] = {}
        self._state_machines: Dict[str, TaskStateMachine] = {}
        self._on_error = on_error
        self._on_state_change = on_state_change
        self._groups: Dict[str, Set[str]] = {}

    # --- Properties ---

    @property
    def task_status(self) -> Dict[str, Dict[str, Any]]:
        """Get the status of all tasks including state and error information.

        Returns:
            A dictionary mapping task names to their status information. Each status contains: state, error,
            retries, and auto_restart.

        """
        return {
            name: {
                "state": ctx.state,
                "error": ctx.error,
                "retries": ctx.retry_count,
                "auto_restart": ctx.auto_restart,
            }
            for name, ctx in self._contexts.items()
        }

    @property
    def tasks(self) -> List[str]:
        """Get a list of all task names.

        Returns:
            A list of all registered task names.

        """
        return list(self._contexts.keys())

    @property
    def task_groups(self) -> List[str]:
        """Get a list of all task group names.

        Returns:
            A list of all registered task group names.

        """
        return list(self._groups.keys())

    # --- Concrete Methods ---

    def _get_context(self, name: str) -> TaskContext:
        """Retrieve a task context by name.

        Args:
            name: The name of the task to retrieve.

        Returns:
            The task context for the specified task.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        if name not in self._contexts:
            raise ValueError(f"Task '{name}' not found")
        return self._contexts[name]

    def _get_state_machine(self, name: str) -> TaskStateMachine:
        """Retrieve a state machine by task name.

        Args:
            name: The name of the task.

        Returns:
            The state machine for the specified task.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        if name not in self._state_machines:
            raise ValueError(f"State machine for task '{name}' not found")
        return self._state_machines[name]

    def _transition(self, name: str, transition_name: str) -> bool:
        """Execute a state transition and update the task context.

        This method orchestrates the state transition by:
        1. Getting the current state from the state machine
        2. Executing the transition
        3. Updating the context's state string
        4. Firing the state change callback

        Args:
            name: The name of the task to transition.
            transition_name: The name of the transition to execute (e.g., 'start', 'pause', 'complete').

        Returns:
            True if the transition was successful, False if not allowed.

        """
        sm = self._get_state_machine(name)
        ctx = self._get_context(name)

        old_state = sm.current_state.name

        try:
            # Execute the transition on the state machine
            transition = getattr(sm, transition_name)
            transition()

            # Update the context with the new state
            new_state = sm.current_state.name
            ctx.state = new_state

            # Fire the state change callback
            if self._on_state_change:
                self._on_state_change(name, old_state, new_state)

            self._logger.debug("Task '%s' transitioned from %s to %s", name, old_state, new_state)
            return True

        except TransitionNotAllowed:
            self._logger.warning(
                "Transition '%s' not allowed for task '%s' in state %s", transition_name, name, old_state
            )
            return False

    async def add_task(  # noqa
        self,
        name: str,
        coro_func: Callable[..., Coroutine[Any, Any, Any]],
        *args: Any,
        auto_restart: bool = False,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        group: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Add a new task to the manager and start it immediately.

        Args:
            name: Unique name for the task.
            coro_func: The coroutine function to execute.
            *args: Positional arguments to pass to the coroutine function.
            auto_restart: Whether to automatically restart the task on failure. Defaults to False.
            max_retries: Maximum number of retry attempts. Defaults to 3.
            retry_delay: Delay in seconds between retry attempts. Defaults to 1.0.
            group: Optional group name to organize tasks. Defaults to None.
            **kwargs: Keyword arguments to pass to the coroutine function.

        Raises:
            ValueError: If a task with the given name already exists.

        """
        if name in self._contexts:
            raise ValueError(f"Task '{name}' already exists")

        # Create the context (pure data)
        context = TaskContext(
            name=name,
            coro_func=coro_func,
            args=args,
            kwargs=kwargs,
            auto_restart=auto_restart,
            max_retries=max_retries,
            retry_delay=retry_delay,
        )

        # Create the state machine (pure state transitions)
        state_machine = TaskStateMachine()

        # Store both
        self._contexts[name] = context
        self._state_machines[name] = state_machine

        # Add to group if specified
        if group:
            self._groups.setdefault(group, set()).add(name)

        # Start the task
        await self._start_task(context)
        self._logger.info("Task '%s' added and started", name)

    async def remove_task(self, name: str) -> None:
        """Remove a task from the manager.

        Stops the task if it's running and removes it from all groups.

        Args:
            name: The name of the task to remove.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        context = self._get_context(name)

        # Stop the task if it's running
        if context.task and not context.task.done():
            try:
                await self.stop_task(name)
            except Exception as e:
                self._logger.warning("Error while stopping task '%s' during removal: %s", name, e)

        # Remove from groups
        for group, group_tasks in list(self._groups.items()):
            group_tasks.discard(name)
            if not group_tasks:
                del self._groups[group]

        # Remove from storage
        del self._contexts[name]
        del self._state_machines[name]

        self._logger.info("Task '%s' removed", name)

    async def _start_task(self, context: TaskContext) -> None:
        """Start a task by transitioning to running state and creating the asyncio task.

        Args:
            context: The task context to start.

        """
        if self._transition(context.name, "start"):
            context.task = asyncio.create_task(self._run_task(context))

    async def _run_task(self, context: TaskContext) -> None:
        """Execute the task with retry logic and error handling.

        This method handles the main execution loop for a task, including
        retries, error handling, and state transitions. It supports both
        standard coroutines and pausable tasks that implement the
        PausableTask protocol.

        Args:
            context: The task context containing execution parameters.

        """
        while context.retry_count <= context.max_retries:
            try:
                # Check if this is a pausable task
                if isinstance(context.coro_func, types.MethodType) and isinstance(
                    context.coro_func.__self__, PausableTask
                ):
                    pausable = cast(PausableTask, context.coro_func.__self__)
                    await self._run_pausable_task(context, pausable)
                else:
                    await self._run_standard_task(context)

                # If we're still running, mark as completed
                if context.state == "Running":
                    context.error = None
                    self._transition(context.name, "complete")
                break

            except asyncio.CancelledError:
                self._transition(context.name, "cancel")
                break

            except Exception as e:
                context.error = e
                self._logger.error("Task '%s' failed: %s", context.name, e)

                if self._on_error:
                    self._on_error(context.name, e)

                # Handle retry logic
                if context.auto_restart and context.retry_count < context.max_retries:
                    context.retry_count += 1
                    self._logger.info(
                        "Retrying task '%s' (attempt %s/%s)",
                        context.name,
                        context.retry_count,
                        context.max_retries,
                    )
                    await asyncio.sleep(context.retry_delay)

                    # Return to running for another attempt.
                    # From running the only path back is: fail -> reset -> start.
                    self._transition(context.name, "fail")
                    if self._transition(context.name, "reset") and self._transition(context.name, "start"):
                        continue
                    break
                else:
                    self._transition(context.name, "fail")
                    break

    async def _run_standard_task(self, context: TaskContext) -> Any:
        """Execute a standard coroutine task with pause support.

        This method wraps a standard coroutine to provide pause/resume
        functionality by monitoring the pause_requested event and
        recreating the task when resumed.

        Args:
            context: The task context containing the coroutine and parameters.

        Returns:
            The result returned by the coroutine function.

        """

        async def pausable_wrapper() -> Any:
            coro = context.coro_func(*context.args, **context.kwargs)
            task = context.inner_task = asyncio.create_task(coro)
            try:
                while not task.done():
                    if context.pause_requested.is_set():
                        task.cancel()
                        try:
                            await task
                        except asyncio.CancelledError:
                            pass

                        # Wait for resume
                        while context.pause_requested.is_set():
                            await asyncio.sleep(0.1)

                        # Recreate the task
                        coro = context.coro_func(*context.args, **context.kwargs)
                        task = context.inner_task = asyncio.create_task(coro)

                    await asyncio.sleep(0.1)
            except asyncio.CancelledError:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
                raise

            return await task

        return await pausable_wrapper()

    async def _run_pausable_task(self, context: TaskContext, pausable: PausableTask) -> Any:
        """Execute a pausable task with checkpoint support.

        This method handles tasks that implement the PausableTask protocol,
        providing checkpoint save/restore functionality for proper pause/resume
        behavior.

        Args:
            context: The task context containing execution parameters.
            pausable: The pausable task instance to execute.

        Returns:
            The result returned by the pausable task's run method.

        """
        # Restore from checkpoint if available
        if context.checkpoint_data:
            await pausable.load_checkpoint(context, context.checkpoint_data)

        while True:
            # Check for pause request
            if context.pause_requested.is_set():
                context.checkpoint_data = await pausable.save_checkpoint(context)

                # Wait for resume
                while context.pause_requested.is_set():
                    await asyncio.sleep(0.1)

                # Restore from checkpoint
                await pausable.load_checkpoint(context, context.checkpoint_data)

            # Run one iteration
            result = await pausable.run(context)

            # If result is not None, task is complete
            if result is not None:
                return result

    async def start_task(self, name: str) -> None:
        """Start or resume a task.

        This method handles different scenarios:
        - Resume a paused task
        - Restart a completed or failed task
        - Start a pending task

        Args:
            name: The name of the task to start.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        context = self._get_context(name)
        sm = self._get_state_machine(name)

        if sm.current_state == sm.paused:
            # Resume from pause
            context.pause_requested.clear()
            self._transition(name, "start")

        elif sm.current_state in [sm.completed, sm.failed, sm.cancelled]:
            # Restart from terminal state
            if self._transition(name, "reset"):
                await self._start_task(context)

        elif sm.current_state == sm.pending:
            # Start for the first time
            await self._start_task(context)

    async def pause_task(self, name: str) -> None:
        """Pause a running task.

        Sets the pause_requested event and transitions the task to paused state.
        The task will pause at the next opportunity.

        Args:
            name: The name of the task to pause.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        context = self._get_context(name)
        sm = self._get_state_machine(name)

        if sm.current_state == sm.running:
            context.pause_requested.set()
            self._transition(name, "pause")

    async def stop_task(self, name: str) -> None:
        """Stop a task by cancelling it and transitioning to cancelled state.

        Called from inside the task itself, the task is cancelled but not awaited; it unwinds at its next await.

        Args:
            name: The name of the task to stop.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        context = self._get_context(name)

        if context.task and not context.task.done():
            context.task.cancel()
            if not self._is_current(context):
                try:
                    await context.task
                except asyncio.CancelledError:
                    pass

        self._transition(name, "cancel")

    def _is_current(self, context: TaskContext) -> bool:
        """Return True when the running asyncio task is the one executing ``context``."""
        current = asyncio.current_task()
        return current is not None and current in (context.task, context.inner_task)

    def task_state(self, name: str) -> str:
        """Get the current state of a task.

        Args:
            name: The name of the task.

        Returns:
            The current state name (e.g., 'Pending', 'Running', 'Completed').

        Raises:
            ValueError: If the task with the given name is not found.

        """
        return self._get_context(name).state

    def task_error(self, name: str) -> Exception | None:
        """Get the last error that occurred for a task.

        Args:
            name: The name of the task.

        Returns:
            The last exception that occurred, or None if no error.

        Raises:
            ValueError: If the task with the given name is not found.

        """
        return self._get_context(name).error

    def get_tasks_in_group(self, group: str) -> Set[str]:
        """Get all task names in a specific group.

        Args:
            group: The name of the group.

        Returns:
            A set of task names in the group, or empty set if group doesn't exist.

        """
        return self._groups.get(group, set()).copy()

    async def wait_all(self, return_exceptions: bool = True) -> List[Any]:
        """Wait for all tasks to complete.

        Args:
            return_exceptions: If True, exceptions are returned as results instead of raising. Defaults to
                True.

        Returns:
            A list of results from all tasks. If return_exceptions is True, exceptions are included in the
            results.

        """
        tasks = [ctx.task for ctx in self._contexts.values() if ctx.task]
        if not tasks:
            return []
        return await asyncio.gather(*tasks, return_exceptions=return_exceptions)

    async def stop_all(self) -> None:
        """Stop all tasks in the manager.

        This method attempts to stop each task gracefully. If a task fails to
        stop, a warning is logged but the operation continues for other tasks.
        The calling task, if it is managed here, is stopped last.

        """
        for name in sorted(self._contexts, key=lambda n: self._is_current(self._contexts[n])):
            try:
                await self.stop_task(name)
            except Exception as e:
                self._logger.warning("Error while stopping task '%s': %s", name, e)

class DownloadTask:
    """Pausable, resumable download task with persistent checkpointing.

    Coordinates between HTTP downloading, file writing, and checkpoint storage.

    """

    def __init__(  # noqa
        self,
        url: str,
        dest: Path | str,
        chunk_size: int = 64 * 1024,
        max_retries: int = 5,
        retry_delay: float = 2,
        checkpoint_interval: int = 1024 * 1024,  # 1MB
        checkpoint_dir: Path | str | None = None,
        logger: logging.Logger | None = None,
    ):
        self._url: str = url
        self._dest = Path(dest)
        self._max_retries: int = max_retries
        self._retry_delay: float = retry_delay
        self._checkpoint_interval: int = checkpoint_interval
        self._logger = logger or logging.getLogger(__name__)

        self._downloader: HttpDownloader = HttpDownloader(chunk_size, logger=self._logger)
        self._checkpoint_storage: CheckpointStorage = self._create_checkpoint_storage(checkpoint_dir)
        self._file_writer: FileWriter = FileWriter(self._dest, self._logger)

        self._state: DownloadState = DownloadState(
            url=self._url,
            dest=self._dest,
            bytes_downloaded=0,
            total_size=None,
        )
        self._last_checkpoint_bytes: int = 0

    @property
    def progress(self) -> float | None:
        """Progress as a percentage, or None when the total size is unknown."""
        return self._state.progress_percent

    @property
    def bytes_downloaded(self) -> int:
        """Bytes written so far, including those restored from a checkpoint."""
        return self._state.bytes_downloaded

    @property
    def total_size(self) -> int | None:
        """The remote file's size, or None when the server did not report it."""
        return self._state.total_size

    async def run(self, context: TaskContext) -> int | None:
        """Download one chunk, resuming from wherever the last run stopped.

        Called repeatedly by the task manager: returning None asks for
        another iteration, and returning the byte count completes the task.
        Disk space is checked against the remaining bytes before writing, so
        a download that cannot fit fails before it starts rather than filling
        the volume.

        Args:
            context: The task context, carrying any restored checkpoint.

        Returns:
            The total bytes downloaded once the file is complete, or None
            while it is still in progress.

        """
        await self._initialize_state()

        # Check available disk space before attempting download ---
        total_size = await self._downloader.get_remote_file_size(self._url)
        if total_size is not None:
            # Update state so progress reporting works from the start
            self._state = replace(self._state, total_size=total_size)

            needed_bytes = total_size - self._state.bytes_downloaded  # account for resume
            free_bytes = shutil.disk_usage(self._dest.parent).free
            if needed_bytes > free_bytes:
                error: Exception = DiskSpaceError(needed_bytes, free_bytes)
                self._logger.error(str(error))
                context.error = error
                raise error
            self._logger.info(
                "Disk space OK: need %s, have %s free",
                format_bytes(needed_bytes),
                format_bytes(free_bytes),
            )
        else:
            self._logger.warning("Could not determine remote file size, skipping disk space check")

        last_exception: Exception | None = None

        for attempt in range(self._max_retries):
            try:
                result = await self._download_with_resume(context)

                if result is not None:
                    await self._checkpoint_storage.delete()

                return result

            except ResumeNotSupportedError:
                self._logger.warning("Resume not supported, restarting download")
                self._file_writer.delete_file()
                await self._checkpoint_storage.delete()
                self._state = DownloadState(self._url, self._dest, 0, None)
                self._last_checkpoint_bytes = 0
                continue

            except (
                httpx.HTTPError,
                httpx.ReadError,
                httpx.RemoteProtocolError,
                httpx.TimeoutException,
                asyncio.TimeoutError,
                OSError,
            ) as e:
                last_exception = e
                self._file_writer.close()

                try:
                    self._file_writer.flush()
                    await self._save_checkpoint(context)
                    self._logger.info("Saved progress at %s before retry", format_bytes(self._state.bytes_downloaded))
                except Exception as checkpoint_error:
                    self._logger.warning("Failed to save checkpoint before retry: %s", checkpoint_error)

                self._logger.warning("Download attempt %d/%d failed: %s", attempt + 1, self._max_retries, e)

                if attempt < self._max_retries - 1:
                    delay = self._retry_delay * (2**attempt)
                    self._logger.info("Retrying in %.2f seconds...", delay)
                    await asyncio.sleep(delay)

            except Exception as e:
                self._logger.exception("Unexpected error")
                context.error = e
                raise

        error = last_exception or RuntimeError(f"Download failed after {self._max_retries} attempts")
        self._logger.error("Download failed: %s", error)
        context.error = error
        raise error

    async def save_checkpoint(self, context: TaskContext) -> Dict:
        """Public API for checkpoint saving."""
        await self._save_checkpoint(context)
        return self._state.to_dict()

    async def load_checkpoint(self, context: TaskContext, checkpoint: Dict) -> None:
        """Public API for checkpoint loading."""
        self._state = DownloadState.from_dict(checkpoint)
        self._last_checkpoint_bytes = self._state.bytes_downloaded
        self._logger.info("Checkpoint loaded: resuming from %s", format_bytes(self._state.bytes_downloaded))

    def _create_checkpoint_storage(self, checkpoint_dir: Path | str | None) -> CheckpointStorage:
        """Factory for checkpoint storage."""
        if checkpoint_dir:
            cp_dir = Path(checkpoint_dir)
        else:
            cp_dir = self._dest.parent / ".checkpoints"

        return FileCheckpointStorage(cp_dir, self._dest, self._logger)

    async def _initialize_state(self) -> None:
        """Initialize or restore download state."""
        # Try to load existing checkpoint
        loaded_state: DownloadState | None = await self._checkpoint_storage.load(self._url, self._dest)

        if loaded_state:
            self._state = loaded_state
            self._last_checkpoint_bytes = loaded_state.bytes_downloaded
            self._logger.info("Resuming download from %s", format_bytes(self._state.bytes_downloaded))
        else:
            # Clean start
            self._file_writer.delete_file()
            await self._checkpoint_storage.delete()

            # Reset to initial state
            self._state = DownloadState(url=self._url, dest=self._dest, bytes_downloaded=0, total_size=None)
            self._last_checkpoint_bytes = 0
            self._logger.info("Starting fresh download")

    async def _download_with_resume(self, context: TaskContext) -> int | None:
        """Perform download with resume support."""
        # Open file appropriately
        if self._state.bytes_downloaded > 0:
            self._file_writer.open_for_append()
            self._logger.debug("Opened file for append at byte %d", self._state.bytes_downloaded)
        else:
            self._file_writer.open_for_write()
            self._logger.debug("Opened file for write")

        try:

            async def chunk_handler(chunk: bytes) -> bool:
                """Process each chunk. Returns True to stop."""
                self._file_writer.write(chunk)

                # Update state immutably using replace()
                self._state = replace(self._state, bytes_downloaded=self._state.bytes_downloaded + len(chunk))

                # Logging
                self._log_progress()

                # Checkpoint if needed
                if self._state.bytes_downloaded - self._last_checkpoint_bytes >= self._checkpoint_interval:
                    self._file_writer.flush()
                    await self._save_checkpoint(context)

                # Check for pause
                if context.pause_requested.is_set():
                    self._file_writer.flush()
                    await self._save_checkpoint(context)
                    self._logger.info("Paused at %s", format_bytes(self._state.bytes_downloaded))
                    self._file_writer.close()
                    return True  # Stop

                return False  # Continue

            # Start download
            total_size = await self._downloader.download_stream(self._url, self._state.bytes_downloaded, chunk_handler)

            # Update total size if we learned it
            if total_size and not self._state.total_size:
                self._state = replace(self._state, total_size=total_size)

            # If we got here without pause, download completed
            if not context.pause_requested.is_set():
                self._logger.info("Download completed: %s (%s)", self._dest, format_bytes(self._state.bytes_downloaded))
                return self._state.bytes_downloaded

            return None  # Paused

        finally:
            self._file_writer.close()

    async def _save_checkpoint(self, context: TaskContext) -> None:
        """Save checkpoint to storage and context."""
        await self._checkpoint_storage.save(self._state)
        context.checkpoint_data = self._state.to_dict()
        self._last_checkpoint_bytes = self._state.bytes_downloaded

    def _log_progress(self) -> None:
        """Log download progress."""
        if self._state.total_size:
            pct = self._state.progress_percent
            self._logger.debug(
                "Progress: %s/%s (%.1f%%)",
                format_bytes(self._state.bytes_downloaded),
                format_bytes(self._state.total_size),
                pct,
            )
        else:
            self._logger.debug("Downloaded: %s", format_bytes(self._state.bytes_downloaded))