PVideoDl/app/services/downloader.py

"""Логика скачивания.

Две стратегии за общим интерфейсом:
- HttpxDownloader — прямые ссылки на файлы, прогресс по chunk'ам.
- YtDlpDownloader — видео/медиа с сайтов (YouTube и сотни других) через yt-dlp.

`pick_downloader()` выбирает стратегию по URL. Прогресс отдаётся через
async-колбэк on_progress, который дёргается не чаще, чем раз в progress_interval.
"""

from __future__ import annotations

import asyncio
import os
import re
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import unquote, urlparse

import httpx

from app.config import settings

# Расширения, которые качаем напрямую через httpx, а не через yt-dlp.
_DIRECT_EXTENSIONS = {
    ".zip", ".rar", ".7z", ".tar", ".gz", ".tgz", ".bz2", ".xz",
    ".iso", ".dmg", ".exe", ".msi", ".apk", ".deb", ".rpm", ".appimage",
    ".pdf", ".epub", ".mobi", ".djvu",
    ".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg", ".bmp", ".tiff",
    ".mp3", ".flac", ".wav", ".ogg", ".m4a", ".aac",
    ".mp4", ".mkv", ".webm", ".mov", ".avi", ".flv", ".m4v",  # прямые ссылки на медиа
    ".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx", ".csv", ".txt", ".json",
    ".bin", ".img", ".dll", ".so",
}

OnProgress = Callable[["Progress"], Awaitable[None]]


@dataclass
class Progress:
    downloaded_bytes: int
    total_bytes: int | None
    speed: float | None
    eta: float | None
    filename: str | None = None


@dataclass
class DownloadResult:
    filename: str
    size_bytes: int | None
    path: Path


def _safe_filename(name: str) -> str:
    """Чистим имя файла от разделителей пути и опасных символов."""
    name = unquote(name).strip().replace("\\", "/").split("/")[-1]
    name = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", name)
    name = name.strip(". ") or "download"
    return name[:200]


def _unique_path(directory: Path, filename: str) -> Path:
    """Не перезатираем существующие файлы — добавляем (1), (2), ..."""
    candidate = directory / filename
    if not candidate.exists():
        return candidate
    stem, suffix = candidate.stem, candidate.suffix
    i = 1
    while True:
        candidate = directory / f"{stem} ({i}){suffix}"
        if not candidate.exists():
            return candidate
        i += 1


def is_direct_file(url: str) -> bool:
    path = urlparse(url).path.lower()
    return any(path.endswith(ext) for ext in _DIRECT_EXTENSIONS)


class Downloader:
    """Базовый интерфейс стратегии скачивания."""

    async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
        raise NotImplementedError


class HttpxDownloader(Downloader):
    def __init__(self) -> None:
        self._dir = settings.download_dir
        self._chunk = settings.chunk_size
        self._interval = settings.progress_interval

    def _filename_from_response(self, url: str, resp: httpx.Response) -> str:
        cd = resp.headers.get("content-disposition", "")
        match = re.search(r"filename\*=(?:UTF-8'')?([^;]+)|filename=\"?([^\";]+)\"?", cd)
        if match:
            raw = match.group(1) or match.group(2)
            if raw:
                return _safe_filename(raw)
        name = urlparse(str(resp.url)).path
        return _safe_filename(name) if name and name != "/" else "download"

    async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
        self._dir.mkdir(parents=True, exist_ok=True)
        timeout = httpx.Timeout(30.0, read=None)
        async with httpx.AsyncClient(follow_redirects=True, timeout=timeout) as client:
            async with client.stream("GET", url) as resp:
                resp.raise_for_status()
                filename = self._filename_from_response(url, resp)
                total = int(resp.headers["content-length"]) if "content-length" in resp.headers else None
                target = _unique_path(self._dir, filename)

                downloaded = 0
                start = time.monotonic()
                last_emit = 0.0
                try:
                    with target.open("wb") as fh:
                        async for chunk in resp.aiter_bytes(self._chunk):
                            fh.write(chunk)
                            downloaded += len(chunk)
                            now = time.monotonic()
                            if now - last_emit >= self._interval:
                                elapsed = now - start
                                speed = downloaded / elapsed if elapsed > 0 else None
                                eta = (
                                    (total - downloaded) / speed
                                    if total and speed and speed > 0
                                    else None
                                )
                                await on_progress(
                                    Progress(downloaded, total, speed, eta, target.name)
                                )
                                last_emit = now
                except BaseException:
                    target.unlink(missing_ok=True)
                    raise

        elapsed = time.monotonic() - start
        speed = downloaded / elapsed if elapsed > 0 else None
        await on_progress(Progress(downloaded, total or downloaded, speed, 0, target.name))
        return DownloadResult(target.name, downloaded, target)


class YtDlpDownloader(Downloader):
    """Скачивание через yt-dlp. yt-dlp синхронный, поэтому крутим его в потоке,
    а progress-хуки прокидываем обратно в event loop через call_soon_threadsafe."""

    def __init__(self) -> None:
        self._dir = settings.download_dir
        self._interval = settings.progress_interval

    async def download(self, url: str, on_progress: OnProgress) -> DownloadResult:
        self._dir.mkdir(parents=True, exist_ok=True)
        loop = asyncio.get_running_loop()
        last_emit = 0.0
        result_holder: dict[str, object] = {}

        def hook(d: dict) -> None:
            nonlocal last_emit
            status = d.get("status")
            if status == "downloading":
                now = time.monotonic()
                if now - last_emit < self._interval:
                    return
                last_emit = now
                downloaded = d.get("downloaded_bytes") or 0
                total = d.get("total_bytes") or d.get("total_bytes_estimate")
                progress = Progress(
                    downloaded_bytes=downloaded,
                    total_bytes=total,
                    speed=d.get("speed"),
                    eta=d.get("eta"),
                    filename=os.path.basename(d.get("filename") or "") or None,
                )
                asyncio.run_coroutine_threadsafe(on_progress(progress), loop)
            elif status == "finished":
                result_holder["path"] = d.get("filename")

        def run_blocking() -> DownloadResult:
            # Импортируем лениво, чтобы httpx-only сценарий не тянул yt-dlp.
            from yt_dlp import YoutubeDL

            ydl_opts = {
                "outtmpl": str(self._dir / "%(title)s [%(id)s].%(ext)s"),
                "progress_hooks": [hook],
                "noprogress": True,
                "quiet": True,
                "no_warnings": True,
                "noplaylist": True,
            }
            with YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                final_path = result_holder.get("path")
                if not final_path:
                    final_path = ydl.prepare_filename(info)
            path = Path(str(final_path))
            size = path.stat().st_size if path.exists() else None
            return DownloadResult(path.name, size, path)

        return await loop.run_in_executor(None, run_blocking)


def pick_downloader(url: str) -> Downloader:
    """Прямые ссылки на файлы — httpx, всё остальное (страницы сайтов) — yt-dlp."""
    return HttpxDownloader() if is_direct_file(url) else YtDlpDownloader()