SpotifyRecAlg/swingmusic/services/download_progress.py

"""
Download Progress Tracking using DragonflyDB.

Provides real-time download progress tracking using DragonflyDB pub/sub
and sorted sets. This enables live progress updates for downloads without
polling the database.
"""

import json
import logging
import time
from dataclasses import dataclass
from enum import StrEnum
from typing import Any

from swingmusic.db.dragonfly_client import get_dragonfly_client

logger = logging.getLogger(__name__)

# Progress update expiry (1 hour)
PROGRESS_TTL = 3600


class DownloadStatus(StrEnum):
    QUEUED = "queued"
    DOWNLOADING = "downloading"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class DownloadProgress:
    """Represents the progress of a single download."""

    download_id: str
    trackhash: str
    title: str
    artist: str
    status: DownloadStatus
    progress_percent: float
    bytes_downloaded: int
    total_bytes: int
    speed_bps: int
    eta_seconds: int
    started_at: int
    updated_at: int
    error_message: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "download_id": self.download_id,
            "trackhash": self.trackhash,
            "title": self.title,
            "artist": self.artist,
            "status": self.status.value,
            "progress_percent": self.progress_percent,
            "bytes_downloaded": self.bytes_downloaded,
            "total_bytes": self.total_bytes,
            "speed_bps": self.speed_bps,
            "eta_seconds": self.eta_seconds,
            "started_at": self.started_at,
            "updated_at": self.updated_at,
            "error_message": self.error_message,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DownloadProgress":
        return cls(
            download_id=data["download_id"],
            trackhash=data["trackhash"],
            title=data["title"],
            artist=data["artist"],
            status=DownloadStatus(data["status"]),
            progress_percent=data["progress_percent"],
            bytes_downloaded=data["bytes_downloaded"],
            total_bytes=data["total_bytes"],
            speed_bps=data["speed_bps"],
            eta_seconds=data["eta_seconds"],
            started_at=data["started_at"],
            updated_at=data["updated_at"],
            error_message=data.get("error_message"),
        )


class DownloadProgressTracker:
    """
    Tracks download progress in real-time using DragonflyDB.

    Uses Redis sorted sets for ordering by time and hash maps
    for storing progress data. Supports pub/sub for live updates.
    """

    def __init__(self):
        self._client = None

    @property
    def client(self):
        if self._client is None:
            self._client = get_dragonfly_client()
        return self._client

    def _get_progress_key(self, download_id: str) -> str:
        """Get the key for a download's progress data."""
        return f"download_progress:{download_id}"

    def _get_user_downloads_key(self, userid: int) -> str:
        """Get the key for a user's active downloads list."""
        return f"downloads:user:{userid}"

    def _get_channel_name(self, userid: int) -> str:
        """Get the pub/sub channel name for a user."""
        return f"downloads:updates:{userid}"

    def start_download(
        self,
        userid: int,
        download_id: str,
        trackhash: str,
        title: str,
        artist: str,
        total_bytes: int = 0,
    ) -> DownloadProgress:
        """
        Start tracking a new download.

        Args:
            userid: The user ID
            download_id: Unique download identifier
            trackhash: Track hash being downloaded
            title: Track title
            artist: Artist name
            total_bytes: Expected total bytes (0 if unknown)

        Returns:
            The created DownloadProgress object
        """
        now = int(time.time())

        progress = DownloadProgress(
            download_id=download_id,
            trackhash=trackhash,
            title=title,
            artist=artist,
            status=DownloadStatus.DOWNLOADING,
            progress_percent=0.0,
            bytes_downloaded=0,
            total_bytes=total_bytes,
            speed_bps=0,
            eta_seconds=0,
            started_at=now,
            updated_at=now,
        )

        if self.client.is_available():
            try:
                # Store progress data
                key = self._get_progress_key(download_id)
                self.client.set(key, json.dumps(progress.to_dict()), ex=PROGRESS_TTL)

                # Add to user's active downloads (sorted by start time)
                user_key = self._get_user_downloads_key(userid)
                self.client.client.zadd(user_key, {download_id: now})
                self.client.client.expire(user_key, PROGRESS_TTL)

                # Publish update
                self._publish_update(userid, progress)

            except Exception as e:
                logger.error(f"Failed to start download tracking: {e}")

        return progress

    def update_progress(
        self,
        userid: int,
        download_id: str,
        bytes_downloaded: int,
        total_bytes: int = 0,
        speed_bps: int = 0,
    ) -> DownloadProgress | None:
        """
        Update download progress.

        Args:
            userid: The user ID
            download_id: Download identifier
            bytes_downloaded: Bytes downloaded so far
            total_bytes: Total bytes (0 if unknown)
            speed_bps: Current download speed in bytes per second

        Returns:
            Updated DownloadProgress or None if not found
        """
        if not self.client.is_available():
            return None

        try:
            key = self._get_progress_key(download_id)
            data = self.client.get(key)

            if not data:
                return None

            progress = DownloadProgress.from_dict(json.loads(data))

            # Update progress
            progress.bytes_downloaded = bytes_downloaded
            progress.updated_at = int(time.time())

            if total_bytes > 0:
                progress.total_bytes = total_bytes
                progress.progress_percent = (bytes_downloaded / total_bytes) * 100

                if speed_bps > 0:
                    progress.speed_bps = speed_bps
                    remaining_bytes = total_bytes - bytes_downloaded
                    progress.eta_seconds = remaining_bytes // speed_bps

            # Store updated progress
            self.client.set(key, json.dumps(progress.to_dict()), ex=PROGRESS_TTL)

            # Publish update
            self._publish_update(userid, progress)

            return progress

        except Exception as e:
            logger.error(f"Failed to update download progress: {e}")
            return None

    def complete_download(
        self,
        userid: int,
        download_id: str,
        success: bool = True,
        error_message: str | None = None,
    ) -> bool:
        """
        Mark a download as completed or failed.

        Args:
            userid: The user ID
            download_id: Download identifier
            success: Whether download succeeded
            error_message: Error message if failed

        Returns:
            True if successful, False otherwise
        """
        if not self.client.is_available():
            return False

        try:
            key = self._get_progress_key(download_id)
            data = self.client.get(key)

            if not data:
                return False

            progress = DownloadProgress.from_dict(json.loads(data))
            progress.status = (
                DownloadStatus.COMPLETED if success else DownloadStatus.FAILED
            )
            progress.progress_percent = 100.0 if success else progress.progress_percent
            progress.updated_at = int(time.time())
            progress.error_message = error_message

            # Store final progress
            self.client.set(key, json.dumps(progress.to_dict()), ex=PROGRESS_TTL)

            # Remove from active downloads
            user_key = self._get_user_downloads_key(userid)
            self.client.client.zrem(user_key, download_id)

            # Publish final update
            self._publish_update(userid, progress)

            return True

        except Exception as e:
            logger.error(f"Failed to complete download tracking: {e}")
            return False

    def cancel_download(self, userid: int, download_id: str) -> bool:
        """Cancel a download."""
        if not self.client.is_available():
            return False

        try:
            key = self._get_progress_key(download_id)
            data = self.client.get(key)

            if not data:
                return False

            progress = DownloadProgress.from_dict(json.loads(data))
            progress.status = DownloadStatus.CANCELLED
            progress.updated_at = int(time.time())

            # Store cancelled status
            self.client.set(key, json.dumps(progress.to_dict()), ex=PROGRESS_TTL)

            # Remove from active downloads
            user_key = self._get_user_downloads_key(userid)
            self.client.client.zrem(user_key, download_id)

            # Publish update
            self._publish_update(userid, progress)

            return True

        except Exception as e:
            logger.error(f"Failed to cancel download: {e}")
            return False

    def get_progress(self, download_id: str) -> DownloadProgress | None:
        """Get the current progress for a download."""
        if not self.client.is_available():
            return None

        try:
            key = self._get_progress_key(download_id)
            data = self.client.get(key)

            if not data:
                return None

            return DownloadProgress.from_dict(json.loads(data))

        except Exception:
            return None

    def get_active_downloads(self, userid: int) -> list[DownloadProgress]:
        """Get all active downloads for a user."""
        if not self.client.is_available():
            return []

        try:
            user_key = self._get_user_downloads_key(userid)

            # Get all download IDs (most recent first)
            download_ids = self.client.client.zrevrange(user_key, 0, -1)

            downloads = []
            for download_id in download_ids:
                progress = self.get_progress(download_id)
                if progress and progress.status in (
                    DownloadStatus.QUEUED,
                    DownloadStatus.DOWNLOADING,
                ):
                    downloads.append(progress)

            return downloads

        except Exception as e:
            logger.error(f"Failed to get active downloads: {e}")
            return []

    def _publish_update(self, userid: int, progress: DownloadProgress) -> None:
        """Publish a progress update to the user's channel."""
        try:
            channel = self._get_channel_name(userid)
            self.client.client.publish(channel, json.dumps(progress.to_dict()))
        except Exception as e:
            logger.debug(f"Failed to publish progress update: {e}")


# Global instance
download_progress_tracker = DownloadProgressTracker()


def get_download_progress_tracker() -> DownloadProgressTracker:
    """Get the global download progress tracker instance."""
    return download_progress_tracker