Pac-cogs/video-archive/video_archiver/utils.py

import os
import shutil
import logging
import asyncio
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Set
import yt_dlp
import ffmpeg
from datetime import datetime, timedelta
from concurrent.futures import ThreadPoolExecutor
from .ffmpeg_manager import FFmpegManager

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("VideoArchiver")

# Initialize FFmpeg manager
ffmpeg_mgr = FFmpegManager()

# Global thread pool for concurrent downloads
download_pool = ThreadPoolExecutor(max_workers=3)

class VideoDownloader:
    def __init__(self, download_path: str, video_format: str, max_quality: int, max_file_size: int, enabled_sites: Optional[List[str]] = None):
        self.download_path = download_path
        self.video_format = video_format
        self.max_quality = max_quality
        self.max_file_size = max_file_size
        self.enabled_sites = enabled_sites
        self.url_patterns = self._get_url_patterns()

        # Configure yt-dlp options
        self.ydl_opts = {
            'format': f'bestvideo[height<={max_quality}]+bestaudio/best[height<={max_quality}]',
            'outtmpl': os.path.join(download_path, '%(title)s.%(ext)s'),
            'merge_output_format': video_format,
            'quiet': True,
            'no_warnings': True,
            'extract_flat': False,
            'concurrent_fragment_downloads': 3,
            'postprocessor_hooks': [self._check_file_size],
            'progress_hooks': [self._progress_hook],
            'ffmpeg_location': ffmpeg_mgr.get_ffmpeg_path(),
        }

    def _get_url_patterns(self) -> List[str]:
        """Get URL patterns for supported sites"""
        patterns = []
        with yt_dlp.YoutubeDL() as ydl:
            for extractor in ydl._ies:
                if hasattr(extractor, '_VALID_URL') and extractor._VALID_URL:
                    if not self.enabled_sites or any(site.lower() in extractor.IE_NAME.lower() for site in self.enabled_sites):
                        patterns.append(extractor._VALID_URL)
        return patterns

    def _check_file_size(self, info):
        """Check if file size is within limits"""
        if info.get('filepath') and os.path.exists(info['filepath']):
            size = os.path.getsize(info['filepath'])
            if size > (self.max_file_size * 1024 * 1024):
                logger.info(f"File exceeds size limit, will compress: {info['filepath']}")

    def _progress_hook(self, d):
        """Handle download progress"""
        if d['status'] == 'finished':
            logger.info(f"Download completed: {d['filename']}")

    async def download_video(self, url: str) -> Tuple[bool, str, str]:
        """Download and process a video"""
        try:
            # Configure yt-dlp for this download
            ydl_opts = self.ydl_opts.copy()
            
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                # Run download in executor to prevent blocking
                info = await asyncio.get_event_loop().run_in_executor(
                    download_pool, lambda: ydl.extract_info(url, download=True)
                )

                if info is None:
                    return False, "", "Failed to extract video information"

                file_path = os.path.join(self.download_path, ydl.prepare_filename(info))

                if not os.path.exists(file_path):
                    return False, "", "Download completed but file not found"

                # Check file size and compress if needed
                file_size = os.path.getsize(file_path)
                if file_size > (self.max_file_size * 1024 * 1024):
                    logger.info(f"Compressing video: {file_path}")
                    try:
                        # Get optimal compression parameters
                        params = ffmpeg_mgr.get_compression_params(file_path, self.max_file_size)
                        output_path = file_path + ".compressed." + self.video_format

                        # Configure ffmpeg with optimal parameters
                        stream = ffmpeg.input(file_path)
                        stream = ffmpeg.output(stream, output_path, **params)

                        # Run compression in executor
                        await asyncio.get_event_loop().run_in_executor(
                            None,
                            lambda: ffmpeg.run(
                                stream,
                                capture_stdout=True,
                                capture_stderr=True,
                                overwrite_output=True,
                            ),
                        )

                        if os.path.exists(output_path):
                            compressed_size = os.path.getsize(output_path)
                            if compressed_size <= (self.max_file_size * 1024 * 1024):
                                os.remove(file_path)  # Remove original
                                return True, output_path, ""
                            else:
                                os.remove(output_path)
                                return False, "", "Failed to compress to target size"
                    except Exception as e:
                        logger.error(f"Compression error: {str(e)}")
                        return False, "", f"Compression error: {str(e)}"

                return True, file_path, ""

        except Exception as e:
            logger.error(f"Download error: {str(e)}")
            return False, "", str(e)

    def is_supported_url(self, url: str) -> bool:
        """Check if URL is supported"""
        try:
            with yt_dlp.YoutubeDL() as ydl:
                # Try to extract info without downloading
                ie = ydl.extract_info(url, download=False, process=False)
                return ie is not None
        except:
            return False


class MessageManager:
    def __init__(self, message_duration: int, message_template: str):
        self.message_duration = message_duration
        self.message_template = message_template
        self.scheduled_deletions: Dict[int, asyncio.Task] = {}

    def format_archive_message(self, author: str, url: str, original_message: str) -> str:
        return self.message_template.format(
            author=author, url=url, original_message=original_message
        )

    async def schedule_message_deletion(self, message_id: int, delete_func) -> None:
        if self.message_duration <= 0:
            return

        if message_id in self.scheduled_deletions:
            self.scheduled_deletions[message_id].cancel()

        async def delete_later():
            await asyncio.sleep(self.message_duration * 3600)  # Convert hours to seconds
            try:
                await delete_func()
            except Exception as e:
                logger.error(f"Failed to delete message {message_id}: {str(e)}")
            finally:
                self.scheduled_deletions.pop(message_id, None)

        self.scheduled_deletions[message_id] = asyncio.create_task(delete_later())


def secure_delete_file(file_path: str, passes: int = 3) -> bool:
    if not os.path.exists(file_path):
        return True

    try:
        file_size = os.path.getsize(file_path)
        for _ in range(passes):
            with open(file_path, "wb") as f:
                f.write(os.urandom(file_size))
                f.flush()
                os.fsync(f.fileno())

        os.remove(file_path)

        if os.path.exists(file_path) or Path(file_path).exists():
            os.unlink(file_path)

        return not (os.path.exists(file_path) or Path(file_path).exists())

    except Exception as e:
        logger.error(f"Error during secure delete: {str(e)}")
        return False


def cleanup_downloads(download_path: str) -> None:
    try:
        if os.path.exists(download_path):
            for file_path in Path(download_path).glob("*"):
                secure_delete_file(str(file_path))

            shutil.rmtree(download_path, ignore_errors=True)
            Path(download_path).mkdir(parents=True, exist_ok=True)
    except Exception as e:
        logger.error(f"Error during cleanup: {str(e)}")