telegram-sticker-downloader/sticker_downloader/cli.py

"""Command-line interface."""

from __future__ import annotations

import argparse
import asyncio
import os
import sys
from collections.abc import Callable, Sequence
from pathlib import Path

import aiohttp
from dotenv import load_dotenv
from telegram import Bot
from telegram.error import InvalidToken, TelegramError

from sticker_downloader import __version__
from sticker_downloader.config import (
    ALL_FILE_TYPES,
    DEFAULT_CONCURRENCY,
    DEFAULT_OUTPUT_DIR,
    DEFAULT_RETRIES,
    DEFAULT_URL_FILE,
    FILE_TYPE_HELP,
    DownloadConfig,
    parse_file_types,
)
from sticker_downloader.downloader import StickerDownloader
from sticker_downloader.errors import StickerDownloaderError
from sticker_downloader.fetcher import AiohttpFetcher
from sticker_downloader.progress import ConsoleProgress, NullProgress, ProgressReporter
from sticker_downloader.results import PackResult
from sticker_downloader.urls import parse_pack_reference_lines, read_pack_references

EXIT_OK = 0
EXIT_FAILURES = 1
EXIT_USAGE = 2

TOKEN_ENV_VAR = "TELEGRAM_BOT_TOKEN"

_TOKEN_INSTRUCTIONS = f"""\
Create a bot with @BotFather on Telegram, then either:
  * put {TOKEN_ENV_VAR}=<your token> in a .env file next to this project, or
  * export {TOKEN_ENV_VAR}=<your token>, or
  * pass --token <your token>."""

_TOKEN_MISSING = f"No Telegram bot token found.\n\n{_TOKEN_INSTRUCTIONS}"
_TOKEN_REJECTED = f"Telegram rejected the bot token.\n\n{_TOKEN_INSTRUCTIONS}"


class UsageError(StickerDownloaderError):
    """Raised for bad invocations; reported without a traceback."""


def build_parser() -> argparse.ArgumentParser:
    types_list = ", ".join(f"{name} ({FILE_TYPE_HELP[name]})" for name in ALL_FILE_TYPES)
    parser = argparse.ArgumentParser(
        prog="tg-stickers",
        description="Download Telegram sticker packs into per-pack folders.",
        epilog=(
            "examples:\n"
            "  tg-stickers https://t.me/addstickers/SomePack\n"
            "  tg-stickers SomePack AnotherPack --types webp,png\n"
            f"  tg-stickers --from-file {DEFAULT_URL_FILE} --concurrency 16\n"
            "  tg-stickers --from-file - < my-packs.txt\n"
            "  tg-stickers            # interactive prompts\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "packs",
        nargs="*",
        metavar="PACK",
        help="sticker pack URLs, tg:// links or bare pack names",
    )
    parser.add_argument(
        "-f",
        "--from-file",
        metavar="PATH",
        help=f"read references from a file, one per line ('-' for stdin); "
        f"blank lines and # comments are ignored (default: {DEFAULT_URL_FILE} "
        f"when chosen interactively)",
    )
    parser.add_argument(
        "-t",
        "--types",
        metavar="LIST",
        help=f"comma-separated file types to save, or 'all'. Supported: {types_list}",
    )
    parser.add_argument(
        "-o",
        "--output",
        metavar="DIR",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help=f"directory to write packs into (default: {DEFAULT_OUTPUT_DIR})",
    )
    parser.add_argument(
        "-c",
        "--concurrency",
        type=int,
        default=DEFAULT_CONCURRENCY,
        metavar="N",
        help=f"files to download in parallel per pack (default: {DEFAULT_CONCURRENCY})",
    )
    parser.add_argument(
        "--retries",
        type=int,
        default=DEFAULT_RETRIES,
        metavar="N",
        help=f"attempts per network call (default: {DEFAULT_RETRIES})",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="re-download files that are already on disk (default: skip them)",
    )
    parser.add_argument(
        "--emoji-names",
        action="store_true",
        help="include each sticker's emoji in its file name",
    )
    parser.add_argument(
        "--no-metadata",
        dest="write_metadata",
        action="store_false",
        help="do not write pack.json alongside the stickers",
    )
    parser.add_argument(
        "--no-pack-thumbnail",
        dest="pack_thumbnail",
        action="store_false",
        help="do not download the pack's own cover thumbnail",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="list what would be downloaded without writing anything",
    )
    parser.add_argument(
        "--token",
        metavar="TOKEN",
        help=f"bot token (default: ${TOKEN_ENV_VAR}, also read from --env-file)",
    )
    parser.add_argument(
        "--env-file",
        metavar="PATH",
        default=".env",
        help="dotenv file to load the token from (default: .env)",
    )
    verbosity = parser.add_mutually_exclusive_group()
    verbosity.add_argument(
        "-v", "--verbose", action="store_true", help="log every file individually"
    )
    verbosity.add_argument(
        "-q", "--quiet", action="store_true", help="only report fatal errors"
    )
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    return parser


def _resolve_references(
    args: argparse.Namespace, progress: ProgressReporter
) -> list[str]:
    """Work out which packs to download, prompting only when nothing was given."""
    if args.packs and args.from_file:
        raise UsageError("pass either PACK arguments or --from-file, not both")

    if args.packs:
        references, problems = parse_pack_reference_lines(args.packs, source="argument")
        for problem in problems:
            progress.warn(problem.split(": ", 1)[-1])
        if not references:
            raise UsageError("no valid sticker pack references given")
        return references

    if args.from_file:
        return _references_from_file(args.from_file, progress)

    return []


def _references_from_file(source: str, progress: ProgressReporter) -> list[str]:
    if source == "-":
        references, problems = parse_pack_reference_lines(
            sys.stdin.read().splitlines(), source="<stdin>"
        )
    else:
        path = Path(source)
        if not path.is_file():
            raise UsageError(f"{path} not found")
        references, problems = read_pack_references(path)

    for problem in problems:
        progress.warn(problem)
    if not references:
        raise UsageError(f"no valid sticker pack references in {source}")
    return references


def collect_interactively(
    *,
    types_given: str | None,
    url_file: Path = DEFAULT_URL_FILE,
    ask: Callable[[str], str] = input,
    say: Callable[[str], None] = print,
) -> tuple[str | None, list[str]]:
    """Ask for the file types and pack references.

    Returns ``(types_text, references)``; ``types_text`` is ``None`` when the
    caller already supplied ``--types``.
    """
    say(f"--- Telegram Sticker Pack Downloader {__version__} ---")

    if types_given is None:
        say("")
        say("Available file types:")
        for name in ALL_FILE_TYPES:
            say(f"  {name:<5} {FILE_TYPE_HELP[name]}")
        types_given = ask("File types (comma-separated, or 'all') [all]: ").strip()

    # Validate now so a typo does not surface after the bot has connected.
    parse_file_types(types_given)

    say("")
    say("Choose an option:")
    say("  1: download one or more packs from URLs you type")
    say(f"  2: download every pack listed in {url_file}")
    choice = ask("Your choice (1 or 2) [1]: ").strip() or "1"

    if choice == "1":
        answer = ask("Sticker pack URL(s), separated by spaces: ").strip()
        references = answer.split()
        if not references:
            raise UsageError("no sticker pack URL given")
        return types_given, references

    if choice == "2":
        if not url_file.is_file():
            raise UsageError(f"{url_file} not found")
        references, problems = read_pack_references(url_file)
        for problem in problems:
            say(f"warning: {problem}")
        if not references:
            raise UsageError(f"no valid sticker pack references in {url_file}")
        say(f"Found {len(references)} pack(s) in {url_file}.")
        return types_given, references

    raise UsageError(f"invalid choice {choice!r}; expected 1 or 2")


def _make_progress(args: argparse.Namespace) -> ProgressReporter:
    if args.quiet:
        return NullProgress()
    return ConsoleProgress(sys.stderr, verbose=args.verbose)


def _resolve_token(args: argparse.Namespace) -> str:
    env_file = Path(args.env_file)
    if env_file.is_file():
        load_dotenv(env_file)
    else:
        load_dotenv()
    token = (args.token or os.getenv(TOKEN_ENV_VAR) or "").strip()
    if not token:
        raise UsageError(_TOKEN_MISSING)
    return token


async def _download(
    token: str,
    references: Sequence[str],
    config: DownloadConfig,
    progress: ProgressReporter,
) -> list[PackResult]:
    timeout = aiohttp.ClientTimeout(total=300, sock_connect=30, sock_read=60)
    async with Bot(token) as bot:
        me = await bot.get_me()
        progress.info(f"Connected as @{me.username}.")
        async with aiohttp.ClientSession(timeout=timeout) as session:
            downloader = StickerDownloader(
                bot=bot,
                fetcher=AiohttpFetcher(session),
                config=config,
                progress=progress,
            )
            return await downloader.download_all(references)


def main(argv: Sequence[str] | None = None) -> int:
    """Run the CLI. Returns the process exit code."""
    args = build_parser().parse_args(argv)
    progress = _make_progress(args)

    try:
        references = _resolve_references(args, progress)
        types_text = args.types

        if not references:
            if not sys.stdin.isatty():
                raise UsageError(
                    "no packs given. Pass PACK arguments, --from-file PATH, "
                    "or run interactively from a terminal."
                )
            types_text, references = collect_interactively(types_given=args.types)

        config = DownloadConfig(
            file_types=parse_file_types(types_text),
            output_dir=args.output,
            concurrency=args.concurrency,
            retries=args.retries,
            overwrite=args.overwrite,
            write_metadata=args.write_metadata,
            pack_thumbnail=args.pack_thumbnail,
            emoji_names=args.emoji_names,
            dry_run=args.dry_run,
        )
        token = _resolve_token(args)
    except (UsageError, StickerDownloaderError, ValueError) as error:
        print(f"error: {error}", file=sys.stderr)
        return EXIT_USAGE
    except (EOFError, KeyboardInterrupt):
        print("\nAborted.", file=sys.stderr)
        return EXIT_USAGE

    if config.dry_run:
        progress.info("Dry run: nothing will be written to disk.")

    try:
        results = asyncio.run(_download(token, references, config, progress))
    except InvalidToken:
        print(f"error: {_TOKEN_REJECTED}", file=sys.stderr)
        return EXIT_USAGE
    except TelegramError as error:
        print(f"error: could not reach Telegram: {error}", file=sys.stderr)
        return EXIT_FAILURES
    except KeyboardInterrupt:
        print("\nInterrupted.", file=sys.stderr)
        return EXIT_FAILURES

    return EXIT_OK if all(result.ok for result in results) else EXIT_FAILURES


def run() -> int:  # pragma: no cover - thin console-script wrapper
    return main()


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())