from __future__ import annotations from loguru import logger import asyncio import os import sys import tempfile import time import uuid from dataclasses import dataclass from pathlib import Path from typing import Any, Optional, Set import aiohttp try: from .config import settings except (ImportError, ValueError): from config import settings # Global registry of paths currently being written / downloaded / processed _ACTIVE_LOCK = asyncio.Lock() _ACTIVE_PATHS: Set[str] = set() async def register_active_path(path: Path) -> None: async with _ACTIVE_LOCK: _ACTIVE_PATHS.add(str(path.resolve())) async def unregister_active_path(path: Path) -> None: async with _ACTIVE_LOCK: _ACTIVE_PATHS.discard(str(path.resolve())) async def is_path_active(path: Path) -> bool: async with _ACTIVE_LOCK: return str(path.resolve()) in _ACTIVE_PATHS @dataclass class ProcessedMedia: media_type: str # "photo" or "video" local_path: Optional[Path] original_url: str attachment_id: str size_bytes: int = 0 duration_sec: Optional[int] = None width: Optional[int] = None height: Optional[int] = None error: Optional[str] = None is_link_only: bool = False thumbnail_url: Optional[str] = None class MediaProcessor: def __init__(self, is_local_tg_api: bool = False) -> None: self.is_local_tg_api = is_local_tg_api self.cache_dir = settings.cache_path @property def max_video_bytes(self) -> int: max_mb = ( settings.video_max_size_mb_local if self.is_local_tg_api else settings.video_max_size_mb_cloud ) return max_mb * 1024 * 1024 async def download_photo( self, session: aiohttp.ClientSession, url: str, prefix: str = "photo_" ) -> Optional[Path]: if not url.startswith(("http://", "https://")): return None suffix = Path(url.split("?", 1)[0]).suffix[:6] or ".jpg" tmp = tempfile.NamedTemporaryFile( dir=self.cache_dir, prefix=prefix, suffix=suffix, delete=False ) tmp_path = Path(tmp.name) tmp.close() await register_active_path(tmp_path) try: async with session.get( url, timeout=aiohttp.ClientTimeout(total=settings.media_download_timeout_sec) ) as resp: if resp.status != 200: logger.warning("Photo download failed with status {}: {}", resp.status, url) tmp_path.unlink(missing_ok=True) await unregister_active_path(tmp_path) return None with open(tmp_path, "wb") as f: async for chunk in resp.content.iter_chunked(64 * 1024): f.write(chunk) return tmp_path except Exception as exc: logger.warning("Error downloading photo {}: {}", url, exc) tmp_path.unlink(missing_ok=True) await unregister_active_path(tmp_path) return None async def download_video_ytdlp( self, url: str, prefix: str = "video_" ) -> tuple[Optional[Path], Optional[str], bool]: """Retries transient failures (stalls, timeouts, one-off yt-dlp/network hiccups) a couple of times before giving up - a single flaky attempt used to fall straight through to the link-only fallback even though the video was perfectly downloadable.""" last_result: tuple[Optional[Path], Optional[str], bool] = (None, "video download failed", False) for attempt in range(1, settings.yt_dlp_retry_attempts + 2): path, err, is_permanent = await self._download_video_ytdlp_attempt(url, prefix) if path or is_permanent: return path, err, is_permanent last_result = (path, err, is_permanent) if attempt <= settings.yt_dlp_retry_attempts: # A fresh upload's lower-quality renditions can still be mid-transcode on # VK's side - the "720p" URL then briefly serves the full-res master until # the ladder catches up, which trips our size cap. That's not a permanent # rejection, just needs more time than a stalled connection does. delay = 60.0 if "exceeds size limit" in (err or "") else 3.0 logger.warning( "yt-dlp attempt {}/{} failed transiently for {}: {} - retrying in {}s", attempt, settings.yt_dlp_retry_attempts + 1, url, err, delay, ) await asyncio.sleep(delay) return last_result async def _download_video_ytdlp_attempt( self, url: str, prefix: str = "video_" ) -> tuple[Optional[Path], Optional[str], bool]: """Returns (path, error_message, is_permanent_error)""" output_path = self.cache_dir / f"{prefix}{uuid.uuid4().hex[:12]}.mp4" await register_active_path(output_path) max_size = self.max_video_bytes netrc_path = None cmd = [sys.executable, "-m", "yt_dlp"] if "vk.com" in url and settings.vk_access_token: netrc_path = f"{output_path}.netrc" try: with open(netrc_path, "w", encoding="utf-8") as fh: fh.write(f"machine vk.com login vk_token password {settings.vk_access_token}\n") cmd.extend(["--netrc-location", netrc_path]) except Exception as exc: logger.warning("Failed to write netrc for yt-dlp: {}", exc) cmd.extend([ url, "-o", str(output_path), "--no-playlist", "--match-filter", f"duration <= {settings.video_max_duration_sec}", "--merge-output-format", "mp4", # Every fallback keeps the height cap: without it yt-dlp can pull an # arbitrarily large/high-res stream only to have it discarded afterwards # by the size check below, burning bandwidth and the whole timeout budget. "-f", ( f"best[height<={settings.video_max_height}][filesize<{max_size}]" f"/best[height<={settings.video_max_height}]" f"/bestvideo[height<={settings.video_max_height}][filesize<{max_size}]+bestaudio" f"/bestvideo[height<={settings.video_max_height}]+bestaudio" ), # No --quiet: we need yt-dlp's own progress lines to tell a slow-but-alive # download (fine, however long it takes) apart from a genuinely hung one - # a flat wall-clock timeout can't tell those apart and was killing large # videos that just needed more time (same class of bug as the TG upload # timeout). --newline makes each progress update its own line instead of # overwriting via \r, so we can read it with readline(). "--newline", "--no-warnings", ]) proc = None output_lines: list[str] = [] deadline = asyncio.get_running_loop().time() + settings.yt_dlp_timeout_sec stall_reason: Optional[str] = None try: proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, ) while True: remaining_total = deadline - asyncio.get_running_loop().time() if remaining_total <= 0: stall_reason = "video download exceeded overall timeout" break wait_for = min(settings.yt_dlp_stall_timeout_sec, remaining_total) try: line = await asyncio.wait_for(proc.stdout.readline(), timeout=wait_for) except asyncio.TimeoutError: stall_reason = "video download stalled (no progress from yt-dlp)" break if not line: break # stdout closed - process is finishing up output_lines.append(line.decode("utf-8", errors="ignore")) if stall_reason: proc.kill() await proc.communicate() output_path.unlink(missing_ok=True) await unregister_active_path(output_path) return None, stall_reason, False await proc.wait() finally: if netrc_path: Path(netrc_path).unlink(missing_ok=True) if proc.returncode != 0 or not output_path.exists() or output_path.stat().st_size == 0: err_msg = "".join(output_lines).lower() is_permanent = any( m in err_msg for m in ( "removed", "unavailable", "private", "access denied", "does not pass filter", "sign in" ) ) output_path.unlink(missing_ok=True) await unregister_active_path(output_path) return None, f"yt-dlp failed: {err_msg[:200]}", is_permanent size = output_path.stat().st_size if size > max_size: logger.warning( "Downloaded video size {} MB exceeds limit {} MB", round(size / (1024 * 1024), 2), round(max_size / (1024 * 1024), 2), ) output_path.unlink(missing_ok=True) await unregister_active_path(output_path) return None, "video exceeds size limit", False return output_path, None, False async def process_media_items( self, items: list[Any] ) -> list[ProcessedMedia]: results: list[ProcessedMedia] = [] async with aiohttp.ClientSession() as session: for item in items: media_type = item.media_type url = item.url att_id = item.attachment_id if media_type == "photo": p = await self.download_photo(session, url, prefix=f"p_{att_id}_") if p: results.append( ProcessedMedia( media_type="photo", local_path=p, original_url=url, attachment_id=att_id, size_bytes=p.stat().st_size, width=item.width, height=item.height, ) ) else: results.append( ProcessedMedia( media_type="photo", local_path=None, original_url=url, attachment_id=att_id, error="failed to download photo", is_link_only=True, ) ) elif media_type == "video": # Check duration if item.duration_sec and item.duration_sec > settings.video_max_duration_sec: results.append( ProcessedMedia( media_type="video", local_path=None, original_url=url, attachment_id=att_id, duration_sec=item.duration_sec, error="video duration exceeds maximum", is_link_only=True, thumbnail_url=item.thumbnail_url, ) ) continue p, err, perm = await self.download_video_ytdlp(url, prefix=f"v_{att_id}_") if p: results.append( ProcessedMedia( media_type="video", local_path=p, original_url=url, attachment_id=att_id, size_bytes=p.stat().st_size, duration_sec=item.duration_sec, width=item.width, height=item.height, ) ) else: results.append( ProcessedMedia( media_type="video", local_path=None, original_url=url, attachment_id=att_id, duration_sec=item.duration_sec, error=err or "failed to download video", is_link_only=True, thumbnail_url=item.thumbnail_url, ) ) return results async def cleanup(self, items: list[ProcessedMedia]) -> None: """Immediately unregisters and removes downloaded local files.""" for item in items: if item.local_path: try: await unregister_active_path(item.local_path) item.local_path.unlink(missing_ok=True) except Exception as exc: logger.warning("Error cleaning up {}: {}", item.local_path, exc) item.local_path = None