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 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]: """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", "-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/best" f"/bestvideo[height<={settings.video_max_height}]+bestaudio/best" f"/best[filesize<{max_size}]" f"/best" ), "--quiet", "--no-warnings", ]) proc = None stderr = b"" try: proc = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) _, stderr = await asyncio.wait_for( proc.communicate(), timeout=settings.yt_dlp_timeout_sec ) except asyncio.TimeoutError: if proc: try: proc.kill() await proc.communicate() except Exception: pass output_path.unlink(missing_ok=True) await unregister_active_path(output_path) return None, "video download timeout", False 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 = (stderr or b"").decode("utf-8", errors="ignore").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", True 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, ) ) 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, ) ) 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