youtube-tui/screens/download.py

#!/usr/bin/env python3
"""
Download Screen for YouTube TUI
"""

import asyncio
from typing import Optional

from textual.app import ComposeResult
from textual.containers import Container
from textual.screen import Screen
from textual.widgets import (
    Footer,
    Header,
    ProgressBar,
    Static,
)

from youtube_tui.models.video import Video
from youtube_tui.services.youtube import YouTubeService


class DownloadScreen(Screen):
    """Screen for displaying download progress"""

    CSS = """
    DownloadScreen {
        align: center middle;
    }
    
    #download-container {
        width: 70%;
        height: auto;
        border: double #555555;
        padding: 2 3;
        margin: 2 0;
    }
    
    #video-title {
        width: 100%;
        height: 3;
        content-align: center middle;
        background: $surface;
        margin-bottom: 1;
    }
    
    #progress-container {
        width: 100%;
        height: 5;
        margin: 2 0;
    }
    
    #status-message {
        width: 100%;
        height: auto;
        content-align: center middle;
        margin: 1 0;
    }
    
    #actions {
        width: 100%;
        height: auto;
        dock: bottom;
        margin-top: 1;
    }
    
    Button {
        width: 15;
        margin: 1 1;
    }
    
    #status-bar {
        dock: bottom;
        height: 1;
        background: $surface;
        color: $text-muted;
        padding: 0 1;
    }
    """

    BINDINGS = [
        ("escape", "cancel", "Cancel"),
        ("ctrl+r", "refresh_screen", "Refresh"),
    ]

    def __init__(self, video: Video, category: Optional[str] = None):
        super().__init__()
        self.youtube_service = YouTubeService()
        self.video = video
        self.category = category
        self.download_complete = False
        self.download_error = False
        self.download_task: Optional[asyncio.Task] = None

    def compose(self) -> ComposeResult:
        """Compose the download screen"""
        yield Header()
        yield Container(
            Static(f"[bold]{self.video.display_title}[/bold]", id="video-title"),
            Static("Preparing download...", id="status-message"),
            Container(
                ProgressBar(total=100, id="progress-bar"),
                id="progress-container",
            ),
            id="download-container",
        )
        yield Static(id="status-bar")
        yield Footer()

    def on_mount(self) -> None:
        """Called when screen is mounted"""
        self.update_status("[blue]Starting download...[/blue]")
        # Use asyncio.create_task instead of app.run_background
        self.download_task = asyncio.create_task(self.start_download())

    def action_refresh_screen(self) -> None:
        """Refresh the screen"""
        # For download screen, refresh just updates the status
        self.update_status("[blue]Status: Download in progress...[/blue]")

    async def start_download(self) -> None:
        """Start the download process"""
        try:
            async def progress_callback(percentage: int) -> bool:
                """Update progress bar during download"""
                self.update_progress(percentage)
                self.update_status(
                    f"[blue]Downloading... {percentage}%[/blue]"
                )
                return True

            success = await self.youtube_service.download_video(
                self.video, self.category, progress_callback=progress_callback
            )

            if success:
                self.download_complete = True
                self.update_status("[green]Download completed![/green]")
                self.update_progress(100)
                # Wait briefly before returning
                await asyncio.sleep(2)
                self.app.pop_screen()
            else:
                self.download_error = True
                self.update_status("[red]Download failed![/red]")
                # Wait briefly before returning
                await asyncio.sleep(2)
                self.app.pop_screen()

        except asyncio.CancelledError:
            # Task was cancelled
            self.download_error = True
            self.update_status("[yellow]Download cancelled[/yellow]")
            self.app.pop_screen()
        except Exception as e:
            self.download_error = True
            self.update_status(f"[red]Error: {e}[/red]")
            # Wait briefly before returning
            await asyncio.sleep(2)
            self.app.pop_screen()

    def update_progress(self, percentage: int) -> None:
        """Update the progress bar"""
        progress_bar = self.query_one("#progress-bar", ProgressBar)
        progress_bar.progress = percentage

    def update_status(self, message: str) -> None:
        """Update the status message"""
        status_message = self.query_one("#status-message", Static)
        status_message.update(message)

        status_bar = self.query_one("#status-bar", Static)
        status_bar.update(f"[bold white]{message}[/bold white]")

    def action_cancel(self) -> None:
        """Cancel the download"""
        # Check if there's a download manager and queue to cancel
        if hasattr(self.app, "download_manager") and self.app.download_manager:
            # Cancel via the download manager
            self.app.download_manager.cancel_active_download()

        if self.download_task:
            self.download_task.cancel()
        self.download_error = True
        self.update_status("[yellow]Download cancelled[/yellow]")
        self.app.pop_screen()

    def on_unload(self) -> None:
        """Called when screen is unloaded"""
        if self.download_complete:
            # Show success message briefly before returning
            self.app.notify(
                f"Downloaded: {self.video.display_title}",
                title="Success",
                severity="information",
                timeout=3,
            )
        elif self.download_error:
            self.app.notify(
                f"Failed to download: {self.video.display_title}",
                title="Error",
                severity="error",
                timeout=3,
            )