youtube-cli/youtube_tui/services/download_manager.py

#!/usr/bin/env python3
"""
Download Manager for YouTube TUI
Handles background downloads sequentially
"""

import asyncio
from typing import Optional

from rich.console import Console

from youtube_tui.models.queue_item import QueueItem, QueueStatus
from youtube_tui.models.video import Video
from youtube_tui.services.queue import DownloadQueue
from youtube_tui.services.youtube import YouTubeService

console = Console()


class DownloadManager:
    """Manages background downloads from the queue"""

    def __init__(self, queue: DownloadQueue, youtube_service: YouTubeService):
        self._queue: DownloadQueue = queue
        self._active_task: Optional[asyncio.Task] = None
        self._current_item: Optional[QueueItem] = None
        self._is_running = False
        self._cancel_requested = False
        self._youtube_service = youtube_service

    def add_to_queue(
        self,
        video: Video,
        category: Optional[str] = None,
        network_folder: Optional[str] = None,
    ) -> QueueItem:
        """Add a video to the download queue"""
        return self._queue.add_video(video, category, network_folder)

    def remove_from_queue(self, item_id: str) -> bool:
        """Remove an item from the queue by UUID string"""
        return self._queue.remove_item(item_id)

    def cancel_active_download(self) -> None:
        """Cancel the currently active download"""
        self._cancel_requested = True
        if self._active_task:
            self._active_task.cancel()

    def get_queue_status(self) -> dict:
        """Get queue status information"""
        stats = self._queue.get_stats()
        # Use _current_item to determine if there's an active download
        has_active_download = self._current_item is not None
        return {
            "pending_count": stats["pending"],
            "downloading_count": 1 if has_active_download else stats["downloading"],
            "total_count": stats["total"],
            "has_active_download": has_active_download,
            "active_item": self._current_item.to_dict() if self._current_item else None,
        }

    def get_active_item(self) -> Optional[QueueItem]:
        """Get the currently downloading item"""
        return self._current_item

    def start_processing(self) -> None:
        """Start the background download processing task"""
        if not self._is_running:
            self._is_running = True
            self._active_task = asyncio.create_task(self._process_queue())

    def stop_processing(self) -> None:
        """Stop the background download processing task"""
        self._is_running = False
        if self._active_task:
            self._active_task.cancel()

    async def _process_queue(self) -> None:
        """Process the download queue sequentially"""
        # Loop is available via asyncio.run() in main context

        while self._is_running:
            try:
                # Check if we have a pending item
                item = self._queue.get_next_pending()
                if item is None:
                    await asyncio.sleep(1)  # Wait for new items
                    continue

                # Mark item as current
                self._current_item = item
                self._cancel_requested = False

                # Start downloading
                await self._download_item(item)

                # Clear current item after completion
                self._current_item = None

            except asyncio.CancelledError:
                # Task was cancelled
                console.print("[yellow]Download manager cancelled[/yellow]")
                break
            except Exception as e:
                console.print(f"[yellow]Error in download manager: {e}[/yellow]")
                await asyncio.sleep(1)

    async def _download_item(self, item: QueueItem) -> None:
        """Download a single queue item"""
        # Update status to downloading
        item.start_download()
        if item.video:
            self._queue.update_item_status(str(item.id), QueueStatus.DOWNLOADING)

        try:
            # Determine if it's a playlist or video
            is_playlist = item.video and (
                "/playlist" in item.video.url.lower()
                or "list=" in item.video.url.lower()
            )

            # Download with progress callback
            async def progress_callback(percentage: int) -> bool:
                """Progress callback that checks for cancellation"""
                # Update progress
                if item.video:
                    self._queue.update_progress(str(item.id), percentage)
                item.update_progress(percentage)

                # Check for cancellation
                if self._cancel_requested:
                    raise asyncio.CancelledError("Download cancelled by user")

                return True

            if is_playlist and item.video:
                success = await self._youtube_service.download_playlist(
                    item.video,
                    category=item.category,
                    network_folder=item.network_folder,
                    progress_callback=progress_callback,
                )
            elif item.video:
                success = await self._youtube_service.download_video(
                    item.video,
                    category=item.category,
                    network_folder=item.network_folder,
                    progress_callback=progress_callback,
                )
            else:
                # No video to download
                if item.video is None:
                    item.fail(error_message="No video data available")
                success = False

            # Check final status
            if self._cancel_requested:
                # Download was cancelled
                if item.video:
                    self._queue.update_item_status(str(item.id), QueueStatus.CANCELLED)
                item.cancel()
            elif success:
                # Download succeeded
                if item.video:
                    self._queue.update_item_status(str(item.id), QueueStatus.COMPLETED)
                item.complete()
            else:
                # Download failed
                if item.video:
                    self._queue.update_item_status(str(item.id), QueueStatus.FAILED)
                item.fail(error_message="Download failed")

        except asyncio.CancelledError:
            # Task was cancelled
            if item.video:
                self._queue.update_item_status(str(item.id), QueueStatus.CANCELLED)
            item.cancel()
        except Exception as e:
            console.print(f"[red]Download error: {e}[/red]")
            if item.video:
                self._queue.update_item_status(str(item.id), QueueStatus.FAILED)
            item.fail(error_message=str(e))

    def is_processing(self) -> bool:
        """Check if download manager is processing queue"""
        return self._is_running