ttvdrops/twitch/management/commands/download_box_art.py

from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse

import httpx
from django.conf import settings
from django.core.files.base import ContentFile
from django.core.management import call_command
from django.core.management.base import BaseCommand
from PIL import Image

from twitch.models import Game
from twitch.utils import is_twitch_box_art_url
from twitch.utils import normalize_twitch_box_art_url

if TYPE_CHECKING:
    from urllib.parse import ParseResult

    from django.core.management.base import CommandParser
    from django.db.models import QuerySet


class Command(BaseCommand):
    """Download and cache Twitch game box art locally."""

    help = "Download and cache Twitch game box art locally."

    def add_arguments(self, parser: CommandParser) -> None:
        """Register command arguments."""
        parser.add_argument(
            "--limit",
            type=int,
            default=None,
            help="Limit the number of games to process.",
        )
        parser.add_argument(
            "--force",
            action="store_true",
            help="Re-download even if a local box art file already exists.",
        )

    def handle(  # noqa: PLR0914, PLR0915
        self,
        *_args: str,
        **options: str | bool | int | None,
    ) -> None:
        """Download Twitch box art images for all games."""
        limit_value: str | bool | int | None = options.get("limit")
        limit: int | None = limit_value if isinstance(limit_value, int) else None
        force: bool = bool(options.get("force"))

        queryset: QuerySet[Game] = Game.objects.all().order_by("twitch_id")
        if limit:
            queryset = queryset[:limit]

        total: int = queryset.count()
        downloaded: int = 0
        skipped: int = 0
        failed: int = 0
        placeholders_404: int = 0

        with httpx.Client(timeout=20, follow_redirects=True) as client:
            for game in queryset:
                if not game.box_art:
                    skipped += 1
                    continue
                if not is_twitch_box_art_url(game.box_art):
                    skipped += 1
                    continue
                if (
                    game.box_art_file
                    and getattr(game.box_art_file, "name", "")
                    and not force
                ):
                    skipped += 1
                    continue

                normalized_url: str = normalize_twitch_box_art_url(game.box_art)
                parsed_url: ParseResult = urlparse(normalized_url)
                suffix: str = Path(parsed_url.path).suffix or ".jpg"
                file_name: str = f"{game.twitch_id}{suffix}"

                try:
                    response: httpx.Response = client.get(normalized_url)
                    response.raise_for_status()
                except httpx.HTTPError as exc:
                    failed += 1
                    self.stdout.write(
                        self.style.WARNING(
                            f"Failed to download box art for {game.twitch_id}: {exc}",
                        ),
                    )
                    continue

                if response.url.path.endswith("/ttv-static/404_boxart.jpg"):
                    placeholders_404 += 1
                    skipped += 1
                    continue

                if game.box_art_file is None:
                    failed += 1
                    continue

                game.box_art_file.save(
                    file_name,
                    ContentFile(response.content),
                    save=True,
                )

                # Auto-convert to WebP and AVIF
                box_art_path: str | None = getattr(game.box_art_file, "path", None)
                if box_art_path:
                    self._convert_to_modern_formats(box_art_path)

                downloaded += 1

        self.stdout.write(
            self.style.SUCCESS(
                f"Processed {total} games. Downloaded: {downloaded}, skipped: {skipped}, "
                f"404 placeholders: {placeholders_404}, failed: {failed}.",
            ),
        )
        box_art_dir: Path = Path(settings.MEDIA_ROOT) / "games" / "box_art"
        self.stdout.write(self.style.SUCCESS(f"Saved box art to: {box_art_dir}"))

        # Convert downloaded images to modern formats (WebP, AVIF)
        if downloaded > 0:
            self.stdout.write(
                self.style.MIGRATE_HEADING(
                    "\nConverting downloaded images to modern formats...",
                ),
            )
            call_command("convert_images_to_modern_formats")

    def _convert_to_modern_formats(self, image_path: str) -> None:
        """Convert downloaded image to WebP and AVIF formats.

        Args:
            image_path: Absolute path to the downloaded image file
        """
        try:
            source_path = Path(image_path)
            if not source_path.exists() or source_path.suffix.lower() not in {
                ".jpg",
                ".jpeg",
                ".png",
            }:
                return

            base_path = source_path.with_suffix("")
            webp_path = base_path.with_suffix(".webp")
            avif_path = base_path.with_suffix(".avif")

            with Image.open(source_path) as img:
                # Convert to RGB if needed
                if img.mode in {"RGBA", "LA"} or (
                    img.mode == "P" and "transparency" in img.info
                ):
                    background = Image.new("RGB", img.size, (255, 255, 255))
                    rgba_img = img.convert("RGBA") if img.mode == "P" else img
                    background.paste(
                        rgba_img,
                        mask=rgba_img.split()[-1]
                        if rgba_img.mode in {"RGBA", "LA"}
                        else None,
                    )
                    rgb_img = background
                elif img.mode != "RGB":
                    rgb_img = img.convert("RGB")
                else:
                    rgb_img = img

                # Save WebP
                rgb_img.save(webp_path, "WEBP", quality=85, method=6)

                # Save AVIF
                rgb_img.save(avif_path, "AVIF", quality=85, speed=4)

        except (OSError, ValueError) as e:
            # Don't fail the download if conversion fails
            self.stdout.write(
                self.style.WARNING(f"Failed to convert {image_path}: {e}"),
            )