telegram-sticker-downloader/sticker_downloader/downloader.py

"""The async sticker-pack downloader."""

from __future__ import annotations

import asyncio
import json
import os
import re
from collections.abc import Awaitable, Callable, Iterable, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, TypeVar

from telegram.error import BadRequest, TelegramError

from sticker_downloader.config import (
    ALL_FILE_TYPES,
    ANIMATED,
    METADATA_FILE_NAME,
    PACK_THUMBNAIL_STEM,
    STATIC,
    THUMBNAIL,
    VIDEO,
    DownloadConfig,
)
from sticker_downloader.errors import InvalidPackReference, StickerDownloaderError
from sticker_downloader.fetcher import Fetcher
from sticker_downloader.progress import NullProgress, ProgressReporter
from sticker_downloader.results import FileOutcome, FileStatus, PackResult
from sticker_downloader.retry import with_retries
from sticker_downloader.urls import parse_pack_name

T = TypeVar("T")

_UNSAFE_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
_PART_SUFFIX = ".part"


@dataclass(frozen=True)
class _FileJob:
    """One file to put on disk."""

    file_id: str
    stem: Path
    #: ``None`` means "derive the extension from Telegram's own file path",
    #: which is how pack thumbnails are handled (they can be webp, tgs or webm).
    suffix: str | None


class StickerDownloader:
    """Downloads sticker packs concurrently into per-pack directories."""

    def __init__(
        self,
        bot: Any,
        fetcher: Fetcher,
        config: DownloadConfig | None = None,
        progress: ProgressReporter | None = None,
        *,
        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
        now: Callable[[], datetime] | None = None,
    ) -> None:
        self._bot = bot
        self._fetcher = fetcher
        self._config = config or DownloadConfig()
        self._progress: ProgressReporter = progress or NullProgress()
        self._sleep = sleep
        self._now = now or (lambda: datetime.now(timezone.utc))

    @property
    def config(self) -> DownloadConfig:
        return self._config

    # -- public API -------------------------------------------------------

    async def download_all(self, references: Iterable[str]) -> list[PackResult]:
        """Download every reference in turn.

        Packs run one after another so progress stays readable; the files inside a
        pack are downloaded concurrently.
        """
        results: list[PackResult] = []
        for reference in references:
            results.append(await self.download_pack(reference))
        self._progress.run_finished(results)
        return results

    async def download_pack(self, reference: str) -> PackResult:
        """Download a single pack, never raising for an expected failure."""
        try:
            pack_name = parse_pack_name(reference)
        except InvalidPackReference as error:
            result = PackResult(reference=reference, error=str(error))
            self._progress.pack_finished(result)
            return result

        try:
            sticker_set = await self._call(lambda: self._bot.get_sticker_set(pack_name))
        except Exception as error:
            result = PackResult(
                reference=reference, name=pack_name, error=_describe(error)
            )
            self._progress.pack_finished(result)
            return result

        stickers: Sequence[Any] = list(getattr(sticker_set, "stickers", None) or [])
        name = getattr(sticker_set, "name", None) or pack_name
        directory = self._config.output_dir / name

        jobs = self._plan_jobs(sticker_set, stickers, directory)
        self._progress.pack_started(reference, name, len(stickers), len(jobs))

        if not self._config.dry_run:
            await asyncio.to_thread(directory.mkdir, parents=True, exist_ok=True)

        semaphore = asyncio.Semaphore(self._config.concurrency)
        outcomes = list(
            await asyncio.gather(*(self._run_job(job, semaphore) for job in jobs))
        )

        result = PackResult(
            reference=reference,
            name=name,
            title=getattr(sticker_set, "title", None),
            directory=directory,
            total_stickers=len(stickers),
            outcomes=outcomes,
        )

        if self._config.write_metadata and not self._config.dry_run:
            try:
                await self._write_metadata(sticker_set, stickers, directory, result)
            except OSError as error:
                self._progress.warn(f"could not write {METADATA_FILE_NAME}: {error}")

        self._progress.pack_finished(result)
        return result

    # -- planning ---------------------------------------------------------

    def _plan_jobs(
        self, sticker_set: Any, stickers: Sequence[Any], directory: Path
    ) -> list[_FileJob]:
        jobs: list[_FileJob] = []

        set_thumbnail = getattr(sticker_set, "thumbnail", None)
        if self._config.pack_thumbnail and set_thumbnail is not None:
            jobs.append(
                _FileJob(
                    file_id=set_thumbnail.file_id,
                    stem=directory / PACK_THUMBNAIL_STEM,
                    suffix=None,
                )
            )

        for index, sticker in enumerate(stickers, start=1):
            base = self._sticker_stem(index, sticker)

            primary = primary_file_type(sticker)
            if self._config.wants(primary):
                jobs.append(
                    _FileJob(
                        file_id=sticker.file_id,
                        stem=directory / base,
                        suffix=f".{primary}",
                    )
                )

            thumbnail = getattr(sticker, "thumbnail", None)
            if self._config.wants(THUMBNAIL) and thumbnail is not None:
                jobs.append(
                    _FileJob(
                        file_id=thumbnail.file_id,
                        stem=directory / f"{base}_thumb",
                        suffix=f".{THUMBNAIL}",
                    )
                )

        return jobs

    def _sticker_stem(self, index: int, sticker: Any) -> str:
        base = f"{index:03d}"
        if not self._config.emoji_names:
            return base
        emoji = sanitize_filename_part(getattr(sticker, "emoji", None) or "")
        return f"{base}_{emoji}" if emoji else base

    # -- execution --------------------------------------------------------

    async def _run_job(self, job: _FileJob, semaphore: asyncio.Semaphore) -> FileOutcome:
        async with semaphore:
            outcome = await self._process_job(job)
        self._progress.file_finished(outcome)
        return outcome

    async def _process_job(self, job: _FileJob) -> FileOutcome:
        planned_path = self._planned_path(job)
        try:
            existing = None if self._config.overwrite else self._existing_path(job)
            if existing is not None:
                return FileOutcome(existing, FileStatus.SKIPPED_EXISTING)

            if self._config.dry_run:
                return FileOutcome(planned_path, FileStatus.PLANNED)

            file = await self._call(lambda: self._bot.get_file(job.file_id))
            remote_path = getattr(file, "file_path", None)
            if not remote_path:
                raise StickerDownloaderError("Telegram returned no file path")

            destination = self._resolve_path(job, remote_path)
            data = await self._call(lambda: self._fetcher.fetch(remote_path))
            await asyncio.to_thread(_write_atomically, destination, data)
            return FileOutcome(destination, FileStatus.DOWNLOADED, size=len(data))
        except Exception as error:
            return FileOutcome(planned_path, FileStatus.FAILED, error=_describe(error))

    async def _call(self, operation: Callable[[], Awaitable[T]]) -> T:
        """Run one network operation with this run's retry policy."""
        return await with_retries(
            operation,
            attempts=self._config.retries,
            on_retry=self._on_retry,
            sleep=self._sleep,
        )

    def _on_retry(self, error: BaseException, attempt: int, delay: float) -> None:
        self._progress.warn(
            f"{_describe(error)} — retrying in {delay:.1f}s (attempt {attempt + 1})"
        )

    # -- paths ------------------------------------------------------------

    def _planned_path(self, job: _FileJob) -> Path:
        """Best guess at the destination, used for dry runs and error messages."""
        return self._resolve_path(job, f"unknown.{STATIC}")

    def _resolve_path(self, job: _FileJob, remote_path: str) -> Path:
        suffix = job.suffix if job.suffix is not None else _suffix_from_url(remote_path)
        return job.stem.parent / f"{job.stem.name}{suffix}"

    def _existing_path(self, job: _FileJob) -> Path | None:
        """The already-downloaded file for ``job``, if there is one.

        Checking before calling ``get_file`` means resuming a large batch costs no
        API calls for the parts that are already on disk.
        """
        if job.suffix is not None:
            candidate = job.stem.parent / f"{job.stem.name}{job.suffix}"
            return candidate if candidate.is_file() else None
        for file_type in ALL_FILE_TYPES:
            candidate = job.stem.parent / f"{job.stem.name}.{file_type}"
            if candidate.is_file():
                return candidate
        return None

    # -- metadata ---------------------------------------------------------

    async def _write_metadata(
        self,
        sticker_set: Any,
        stickers: Sequence[Any],
        directory: Path,
        result: PackResult,
    ) -> None:
        present = {
            outcome.path
            for outcome in result.outcomes
            if outcome.status in (FileStatus.DOWNLOADED, FileStatus.SKIPPED_EXISTING)
        }

        entries = []
        for index, sticker in enumerate(stickers, start=1):
            base = self._sticker_stem(index, sticker)
            files = sorted(
                path.name
                for path in present
                if path.stem == base or path.stem == f"{base}_thumb"
            )
            entries.append(
                {
                    "index": index,
                    "emoji": getattr(sticker, "emoji", None),
                    "file_unique_id": getattr(sticker, "file_unique_id", None),
                    "type": _enum_value(getattr(sticker, "type", None)),
                    "width": getattr(sticker, "width", None),
                    "height": getattr(sticker, "height", None),
                    "is_animated": bool(getattr(sticker, "is_animated", False)),
                    "is_video": bool(getattr(sticker, "is_video", False)),
                    "files": files,
                }
            )

        from sticker_downloader import __version__

        payload = {
            "pack": {
                "reference": result.reference,
                "name": result.name,
                "title": result.title,
                "sticker_type": _enum_value(getattr(sticker_set, "sticker_type", None)),
            },
            "generated_at": self._now().isoformat(),
            "downloader_version": __version__,
            "requested_file_types": self._config.sorted_file_types,
            "sticker_count": len(stickers),
            "stickers": entries,
        }

        text = json.dumps(payload, indent=2, ensure_ascii=False) + "\n"
        await asyncio.to_thread(
            (directory / METADATA_FILE_NAME).write_text, text, encoding="utf-8"
        )


# -- module-level helpers -------------------------------------------------


def primary_file_type(sticker: Any) -> str:
    """The file type Telegram stores this sticker as."""
    if getattr(sticker, "is_animated", False):
        return ANIMATED
    if getattr(sticker, "is_video", False):
        return VIDEO
    return STATIC


def sanitize_filename_part(text: str) -> str:
    """Strip characters that are illegal or awkward inside a file name."""
    cleaned = _UNSAFE_FILENAME_CHARS.sub("", text).strip(" .")
    return cleaned[:32]


def _suffix_from_url(remote_path: str) -> str:
    suffix = Path(remote_path.split("?")[0]).suffix.lower()
    return suffix if suffix else f".{STATIC}"


def _write_atomically(destination: Path, data: bytes) -> None:
    """Write via a temporary file so an interrupted run leaves no partial file."""
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(destination.name + _PART_SUFFIX)
    temporary.write_bytes(data)
    os.replace(temporary, destination)


def _enum_value(value: Any) -> Any:
    return getattr(value, "value", value)


def _describe(error: BaseException) -> str:
    """A short, user-facing description of a failure."""
    if isinstance(error, BadRequest) and "stickerset_invalid" in str(error).lower():
        return "sticker pack not found"
    if isinstance(error, TelegramError):
        return f"{type(error).__name__}: {error}"
    message = str(error).strip()
    return message or type(error).__name__