mirror of
https://github.com/pacnpal/Pac-cogs.git
synced 2025-12-20 10:51:05 -05:00
b1eafbb01d
- Created new queue module with separate components: - models.py: QueueItem and QueueMetrics data classes - persistence.py: Queue state persistence - monitoring.py: Health monitoring and metrics - cleanup.py: Cleanup operations - manager.py: Main queue management - __init__.py: Package exports - Updated imports in video_archiver.py and processor.py - Removed old enhanced_queue.py - Updated README with new queue system details This refactoring improves code organization and maintainability through better separation of concerns while maintaining all existing functionality.
380 lines
15 KiB
Python
380 lines
15 KiB
Python
"""VideoArchiver cog for Red-DiscordBot"""
|
|
from __future__ import annotations
|
|
|
|
import discord
|
|
from redbot.core import commands, Config, data_manager
|
|
from pathlib import Path
|
|
import logging
|
|
import asyncio
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime
|
|
import sys
|
|
import traceback
|
|
|
|
from videoarchiver.config_manager import ConfigManager
|
|
from videoarchiver.update_checker import UpdateChecker
|
|
from videoarchiver.processor import VideoProcessor
|
|
from videoarchiver.utils.video_downloader import VideoDownloader
|
|
from videoarchiver.utils.message_manager import MessageManager
|
|
from videoarchiver.utils.file_ops import cleanup_downloads
|
|
from videoarchiver.queue import EnhancedVideoQueueManager # Updated import
|
|
from videoarchiver.ffmpeg.ffmpeg_manager import FFmpegManager
|
|
from videoarchiver.utils.exceptions import (
|
|
VideoArchiverError as ProcessingError,
|
|
ConfigurationError as ConfigError,
|
|
VideoVerificationError as UpdateError,
|
|
QueueError,
|
|
FileCleanupError as FileOperationError
|
|
)
|
|
|
|
logger = logging.getLogger("VideoArchiver")
|
|
|
|
# Constants for timeouts
|
|
UNLOAD_TIMEOUT = 30 # seconds
|
|
CLEANUP_TIMEOUT = 15 # seconds
|
|
|
|
class VideoArchiver(commands.Cog):
|
|
"""Archive videos from Discord channels"""
|
|
|
|
def __init__(self, bot: commands.Bot) -> None:
|
|
"""Initialize the cog with proper error handling"""
|
|
self.bot = bot
|
|
self.ready = asyncio.Event()
|
|
self._init_task: Optional[asyncio.Task] = None
|
|
self._cleanup_task: Optional[asyncio.Task] = None
|
|
self._unloading = False
|
|
|
|
# Start initialization
|
|
self._init_task = asyncio.create_task(self._initialize())
|
|
self._init_task.add_done_callback(self._init_callback)
|
|
|
|
async def _initialize(self) -> None:
|
|
"""Initialize all components with proper error handling"""
|
|
try:
|
|
# Initialize config first as other components depend on it
|
|
config = Config.get_conf(self, identifier=855847, force_registration=True)
|
|
self.config_manager = ConfigManager(config)
|
|
|
|
# Set up paths
|
|
self.data_path = Path(data_manager.cog_data_path(self))
|
|
self.download_path = self.data_path / "downloads"
|
|
self.download_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Clean existing downloads
|
|
await cleanup_downloads(str(self.download_path))
|
|
|
|
# Initialize shared FFmpeg manager
|
|
self.ffmpeg_mgr = FFmpegManager()
|
|
logger.info("Initialized shared FFmpeg manager")
|
|
|
|
# Initialize components dict first
|
|
self.components: Dict[int, Dict[str, Any]] = {}
|
|
|
|
# Initialize components for existing guilds
|
|
for guild in self.bot.guilds:
|
|
try:
|
|
await self.initialize_guild_components(guild.id)
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize guild {guild.id}: {str(e)}")
|
|
continue
|
|
|
|
# Initialize queue manager after components are ready
|
|
queue_path = self.data_path / "queue_state.json"
|
|
queue_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.queue_manager = EnhancedVideoQueueManager(
|
|
max_retries=3,
|
|
retry_delay=5,
|
|
max_queue_size=1000,
|
|
cleanup_interval=1800,
|
|
max_history_age=86400,
|
|
persistence_path=str(queue_path),
|
|
)
|
|
|
|
# Initialize update checker
|
|
self.update_checker = UpdateChecker(self.bot, self.config_manager)
|
|
|
|
# Initialize processor with queue manager and shared FFmpeg manager
|
|
self.processor = VideoProcessor(
|
|
self.bot,
|
|
self.config_manager,
|
|
self.components,
|
|
queue_manager=self.queue_manager,
|
|
ffmpeg_mgr=self.ffmpeg_mgr,
|
|
)
|
|
|
|
# Start update checker
|
|
await self.update_checker.start()
|
|
|
|
# Set ready flag
|
|
self.ready.set()
|
|
|
|
logger.info("VideoArchiver initialization completed successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Critical error during initialization: {traceback.format_exc()}")
|
|
await self._cleanup()
|
|
raise
|
|
|
|
def _init_callback(self, task: asyncio.Task) -> None:
|
|
"""Handle initialization task completion"""
|
|
try:
|
|
task.result()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"Initialization failed: {str(e)}")
|
|
asyncio.create_task(self._cleanup())
|
|
|
|
async def cog_load(self) -> None:
|
|
"""Handle cog loading"""
|
|
try:
|
|
await asyncio.wait_for(self.ready.wait(), timeout=30)
|
|
except asyncio.TimeoutError:
|
|
await self._cleanup()
|
|
raise ProcessingError("Cog initialization timed out")
|
|
except Exception as e:
|
|
await self._cleanup()
|
|
raise ProcessingError(f"Error during cog load: {str(e)}")
|
|
|
|
async def cog_unload(self) -> None:
|
|
"""Clean up when cog is unloaded with timeout"""
|
|
self._unloading = True
|
|
try:
|
|
# Create cleanup task with timeout
|
|
cleanup_task = asyncio.create_task(self._cleanup())
|
|
try:
|
|
await asyncio.wait_for(cleanup_task, timeout=UNLOAD_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
logger.error("Cog unload timed out, forcing cleanup")
|
|
# Force cleanup of any remaining resources
|
|
await self._force_cleanup()
|
|
except Exception as e:
|
|
logger.error(f"Error during cog unload: {str(e)}")
|
|
await self._force_cleanup()
|
|
finally:
|
|
self._unloading = False
|
|
|
|
async def _force_cleanup(self) -> None:
|
|
"""Force cleanup of resources when timeout occurs"""
|
|
try:
|
|
# Cancel all tasks
|
|
if hasattr(self, "processor"):
|
|
await self.processor.force_cleanup()
|
|
|
|
# Force stop queue manager
|
|
if hasattr(self, "queue_manager"):
|
|
self.queue_manager.force_stop()
|
|
|
|
# Kill any remaining FFmpeg processes
|
|
if hasattr(self, "ffmpeg_mgr"):
|
|
self.ffmpeg_mgr.kill_all_processes()
|
|
|
|
# Clean up download directory
|
|
if hasattr(self, "download_path") and self.download_path.exists():
|
|
try:
|
|
await cleanup_downloads(str(self.download_path))
|
|
self.download_path.rmdir()
|
|
except Exception as e:
|
|
logger.error(f"Error force cleaning download directory: {str(e)}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during force cleanup: {str(e)}")
|
|
finally:
|
|
self.ready.clear()
|
|
|
|
async def _cleanup(self) -> None:
|
|
"""Clean up all resources with proper handling"""
|
|
try:
|
|
# Cancel initialization if still running
|
|
if self._init_task and not self._init_task.done():
|
|
self._init_task.cancel()
|
|
try:
|
|
await asyncio.wait_for(self._init_task, timeout=CLEANUP_TIMEOUT)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
|
|
# Stop update checker
|
|
if hasattr(self, "update_checker"):
|
|
try:
|
|
await asyncio.wait_for(self.update_checker.stop(), timeout=CLEANUP_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
# Clean up processor
|
|
if hasattr(self, "processor"):
|
|
try:
|
|
await asyncio.wait_for(self.processor.cleanup(), timeout=CLEANUP_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
await self.processor.force_cleanup()
|
|
|
|
# Clean up queue manager
|
|
if hasattr(self, "queue_manager"):
|
|
try:
|
|
await asyncio.wait_for(self.queue_manager.cleanup(), timeout=CLEANUP_TIMEOUT)
|
|
except asyncio.TimeoutError:
|
|
self.queue_manager.force_stop()
|
|
|
|
# Clean up components for each guild
|
|
if hasattr(self, "components"):
|
|
for guild_id, components in self.components.items():
|
|
try:
|
|
if "message_manager" in components:
|
|
await components["message_manager"].cancel_all_deletions()
|
|
if "downloader" in components:
|
|
components["downloader"] = None
|
|
if "ffmpeg_mgr" in components:
|
|
components["ffmpeg_mgr"] = None
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up guild {guild_id}: {str(e)}")
|
|
|
|
self.components.clear()
|
|
|
|
# Clean up download directory
|
|
if hasattr(self, "download_path") and self.download_path.exists():
|
|
try:
|
|
await cleanup_downloads(str(self.download_path))
|
|
self.download_path.rmdir()
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up download directory: {str(e)}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during cleanup: {traceback.format_exc()}")
|
|
raise ProcessingError(f"Cleanup failed: {str(e)}")
|
|
finally:
|
|
self.ready.clear()
|
|
|
|
async def initialize_guild_components(self, guild_id: int) -> None:
|
|
"""Initialize or update components for a guild with error handling"""
|
|
try:
|
|
settings = await self.config_manager.get_guild_settings(guild_id)
|
|
|
|
# Ensure download directory exists and is clean
|
|
self.download_path.mkdir(parents=True, exist_ok=True)
|
|
await cleanup_downloads(str(self.download_path))
|
|
|
|
# Clean up old components if they exist
|
|
if guild_id in self.components:
|
|
old_components = self.components[guild_id]
|
|
if "message_manager" in old_components:
|
|
await old_components["message_manager"].cancel_all_deletions()
|
|
if "downloader" in old_components:
|
|
old_components["downloader"] = None
|
|
|
|
# Initialize new components with validated settings
|
|
self.components[guild_id] = {
|
|
"downloader": VideoDownloader(
|
|
str(self.download_path),
|
|
settings["video_format"],
|
|
settings["video_quality"],
|
|
settings["max_file_size"],
|
|
settings["enabled_sites"] if settings["enabled_sites"] else None,
|
|
settings["concurrent_downloads"],
|
|
ffmpeg_mgr=self.ffmpeg_mgr, # Use shared FFmpeg manager
|
|
),
|
|
"message_manager": MessageManager(
|
|
settings["message_duration"], settings["message_template"]
|
|
),
|
|
}
|
|
|
|
logger.info(f"Successfully initialized components for guild {guild_id}")
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to initialize guild {guild_id}: {traceback.format_exc()}"
|
|
)
|
|
raise ProcessingError(f"Guild initialization failed: {str(e)}")
|
|
|
|
@commands.Cog.listener()
|
|
async def on_guild_join(self, guild: discord.Guild) -> None:
|
|
"""Handle bot joining a new guild"""
|
|
if not self.ready.is_set():
|
|
return
|
|
|
|
try:
|
|
await self.initialize_guild_components(guild.id)
|
|
logger.info(f"Initialized components for new guild {guild.id}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize new guild {guild.id}: {str(e)}")
|
|
|
|
@commands.Cog.listener()
|
|
async def on_guild_remove(self, guild: discord.Guild) -> None:
|
|
"""Handle bot leaving a guild"""
|
|
try:
|
|
if guild.id in self.components:
|
|
# Clean up components
|
|
components = self.components[guild.id]
|
|
if "message_manager" in components:
|
|
await components["message_manager"].cancel_all_deletions()
|
|
if "downloader" in components:
|
|
components["downloader"] = None
|
|
if "ffmpeg_mgr" in components:
|
|
components["ffmpeg_mgr"] = None
|
|
|
|
# Remove guild components
|
|
self.components.pop(guild.id)
|
|
|
|
logger.info(f"Cleaned up components for removed guild {guild.id}")
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up removed guild {guild.id}: {str(e)}")
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message) -> None:
|
|
"""Handle new messages for video processing"""
|
|
if not self.ready.is_set() or message.guild is None or message.author.bot:
|
|
return
|
|
|
|
try:
|
|
await self.processor.process_message(message)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error processing message {message.id}: {traceback.format_exc()}"
|
|
)
|
|
try:
|
|
log_channel = await self.config_manager.get_channel(
|
|
message.guild, "log"
|
|
)
|
|
if log_channel:
|
|
await log_channel.send(
|
|
f"Error processing message: {str(e)}\n"
|
|
f"Message ID: {message.id}\n"
|
|
f"Channel: {message.channel.mention}"
|
|
)
|
|
except Exception as log_error:
|
|
logger.error(f"Failed to log error to guild: {str(log_error)}")
|
|
|
|
async def cog_command_error(self, ctx: commands.Context, error: Exception) -> None:
|
|
"""Handle command errors"""
|
|
error_msg = None
|
|
try:
|
|
if isinstance(error, commands.MissingPermissions):
|
|
error_msg = "❌ You don't have permission to use this command."
|
|
elif isinstance(error, commands.BotMissingPermissions):
|
|
error_msg = "❌ I don't have the required permissions to do that."
|
|
elif isinstance(error, commands.MissingRequiredArgument):
|
|
error_msg = f"❌ Missing required argument: {error.param.name}"
|
|
elif isinstance(error, commands.BadArgument):
|
|
error_msg = f"❌ Invalid argument: {str(error)}"
|
|
elif isinstance(error, ConfigError):
|
|
error_msg = f"❌ Configuration error: {str(error)}"
|
|
elif isinstance(error, ProcessingError):
|
|
error_msg = f"❌ Processing error: {str(error)}"
|
|
else:
|
|
logger.error(
|
|
f"Command error in {ctx.command}: {traceback.format_exc()}"
|
|
)
|
|
error_msg = (
|
|
"❌ An unexpected error occurred. Check the logs for details."
|
|
)
|
|
|
|
if error_msg:
|
|
await ctx.send(error_msg)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error handling command error: {str(e)}")
|
|
try:
|
|
await ctx.send(
|
|
"❌ An error occurred while handling another error. Please check the logs."
|
|
)
|
|
except Exception:
|
|
pass # Give up if we can't even send error messages
|