"""Enhanced scheduler service for flexible task scheduling with timezone support.""" from collections.abc import Callable from datetime import datetime, timedelta from typing import Any, Dict, List, Optional import pytz from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.cron import CronTrigger from apscheduler.triggers.date import DateTrigger from apscheduler.triggers.interval import IntervalTrigger from sqlmodel.ext.asyncio.session import AsyncSession from app.core.logging import get_logger from app.models.scheduled_task import ( RecurrenceType, ScheduledTask, TaskStatus, TaskType, ) from app.repositories.scheduled_task import ScheduledTaskRepository from app.services.credit import CreditService from app.services.player import PlayerService from app.services.task_handlers import TaskHandlerRegistry logger = get_logger(__name__) class SchedulerService: """Enhanced service for managing scheduled tasks with timezone support.""" def __init__( self, db_session_factory: Callable[[], AsyncSession], player_service: PlayerService, ) -> None: """Initialize the scheduler service. Args: db_session_factory: Factory function to create database sessions player_service: Player service for audio playback tasks """ self.db_session_factory = db_session_factory self.scheduler = AsyncIOScheduler(timezone=pytz.UTC) self.credit_service = CreditService(db_session_factory) self.player_service = player_service self._running_tasks: set[str] = set() async def start(self) -> None: """Start the scheduler and load all active tasks.""" logger.info("Starting enhanced scheduler service...") self.scheduler.start() # Schedule system tasks initialization for after startup self.scheduler.add_job( self._initialize_system_tasks, "date", run_date=datetime.utcnow() + timedelta(seconds=2), id="initialize_system_tasks", name="Initialize System Tasks", replace_existing=True, ) # Schedule periodic cleanup and maintenance self.scheduler.add_job( self._maintenance_job, "interval", minutes=5, id="scheduler_maintenance", name="Scheduler Maintenance", replace_existing=True, ) logger.info("Enhanced scheduler service started successfully") async def stop(self) -> None: """Stop the scheduler.""" logger.info("Stopping scheduler service...") self.scheduler.shutdown(wait=True) logger.info("Scheduler service stopped") async def create_task( self, name: str, task_type: TaskType, scheduled_at: datetime, parameters: Optional[Dict[str, Any]] = None, user_id: Optional[int] = None, timezone: str = "UTC", recurrence_type: RecurrenceType = RecurrenceType.NONE, cron_expression: Optional[str] = None, recurrence_count: Optional[int] = None, expires_at: Optional[datetime] = None, ) -> ScheduledTask: """Create a new scheduled task.""" async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) # Convert scheduled_at to UTC if it's in a different timezone if timezone != "UTC": tz = pytz.timezone(timezone) if scheduled_at.tzinfo is None: # Assume the datetime is in the specified timezone scheduled_at = tz.localize(scheduled_at) scheduled_at = scheduled_at.astimezone(pytz.UTC).replace(tzinfo=None) task_data = { "name": name, "task_type": task_type, "scheduled_at": scheduled_at, "timezone": timezone, "parameters": parameters or {}, "user_id": user_id, "recurrence_type": recurrence_type, "cron_expression": cron_expression, "recurrence_count": recurrence_count, "expires_at": expires_at, } created_task = await repo.create(task_data) await self._schedule_apscheduler_job(created_task) logger.info(f"Created scheduled task: {created_task.name} ({created_task.id})") return created_task async def cancel_task(self, task_id: int) -> bool: """Cancel a scheduled task.""" async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) task = await repo.get_by_id(task_id) if not task: return False task.status = TaskStatus.CANCELLED task.is_active = False await repo.update(task) # Remove from APScheduler try: self.scheduler.remove_job(str(task_id)) except Exception: pass # Job might not exist in scheduler logger.info(f"Cancelled task: {task.name} ({task_id})") return True async def get_user_tasks( self, user_id: int, status: Optional[TaskStatus] = None, task_type: Optional[TaskType] = None, limit: Optional[int] = None, offset: Optional[int] = None, ) -> List[ScheduledTask]: """Get tasks for a specific user.""" async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) return await repo.get_user_tasks(user_id, status, task_type, limit, offset) async def _initialize_system_tasks(self) -> None: """Initialize system tasks and load active tasks from database.""" logger.info("Initializing system tasks...") try: # Create system tasks if they don't exist await self._ensure_system_tasks() # Load all active tasks from database await self._load_active_tasks() logger.info("System tasks initialized successfully") except Exception: logger.exception("Failed to initialize system tasks") async def _ensure_system_tasks(self) -> None: """Ensure required system tasks exist.""" async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) # Check if daily credit recharge task exists system_tasks = await repo.get_system_tasks( task_type=TaskType.CREDIT_RECHARGE ) daily_recharge_exists = any( task.recurrence_type == RecurrenceType.DAILY and task.is_active for task in system_tasks ) if not daily_recharge_exists: # Create daily credit recharge task tomorrow_midnight = datetime.utcnow().replace( hour=0, minute=0, second=0, microsecond=0 ) + timedelta(days=1) task_data = { "name": "Daily Credit Recharge", "task_type": TaskType.CREDIT_RECHARGE, "scheduled_at": tomorrow_midnight, "recurrence_type": RecurrenceType.DAILY, "parameters": {}, } await repo.create(task_data) logger.info("Created system daily credit recharge task") async def _load_active_tasks(self) -> None: """Load all active tasks from database into scheduler.""" async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) active_tasks = await repo.get_active_tasks() for task in active_tasks: await self._schedule_apscheduler_job(task) logger.info(f"Loaded {len(active_tasks)} active tasks into scheduler") async def _schedule_apscheduler_job(self, task: ScheduledTask) -> None: """Schedule a task in APScheduler.""" job_id = str(task.id) # Remove existing job if it exists try: self.scheduler.remove_job(job_id) except Exception: pass # Don't schedule if task is not active or already completed/failed if not task.is_active or task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]: return # Create trigger based on recurrence type trigger = self._create_trigger(task) if not trigger: logger.warning(f"Could not create trigger for task {task.id}") return # Schedule the job self.scheduler.add_job( self._execute_task, trigger=trigger, args=[task.id], id=job_id, name=task.name, replace_existing=True, ) logger.debug(f"Scheduled APScheduler job for task {task.id}") def _create_trigger(self, task: ScheduledTask): """Create APScheduler trigger based on task configuration.""" tz = pytz.timezone(task.timezone) if task.recurrence_type == RecurrenceType.NONE: return DateTrigger(run_date=task.scheduled_at, timezone=tz) elif task.recurrence_type == RecurrenceType.CRON and task.cron_expression: return CronTrigger.from_crontab(task.cron_expression, timezone=tz) elif task.recurrence_type == RecurrenceType.HOURLY: return IntervalTrigger(hours=1, start_date=task.scheduled_at, timezone=tz) elif task.recurrence_type == RecurrenceType.DAILY: return IntervalTrigger(days=1, start_date=task.scheduled_at, timezone=tz) elif task.recurrence_type == RecurrenceType.WEEKLY: return IntervalTrigger(weeks=1, start_date=task.scheduled_at, timezone=tz) elif task.recurrence_type == RecurrenceType.MONTHLY: # Use cron trigger for monthly (more reliable than interval) scheduled_time = task.scheduled_at return CronTrigger( day=scheduled_time.day, hour=scheduled_time.hour, minute=scheduled_time.minute, timezone=tz ) elif task.recurrence_type == RecurrenceType.YEARLY: scheduled_time = task.scheduled_at return CronTrigger( month=scheduled_time.month, day=scheduled_time.day, hour=scheduled_time.hour, minute=scheduled_time.minute, timezone=tz ) return None async def _execute_task(self, task_id: int) -> None: """Execute a scheduled task.""" task_id_str = str(task_id) # Prevent concurrent execution of the same task if task_id_str in self._running_tasks: logger.warning(f"Task {task_id} is already running, skipping execution") return self._running_tasks.add(task_id_str) try: async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) # Get fresh task data task = await repo.get_by_id(task_id) if not task: logger.warning(f"Task {task_id} not found") return # Check if task is still active and pending if not task.is_active or task.status != TaskStatus.PENDING: logger.info(f"Task {task_id} is not active or not pending, skipping") return # Check if task has expired if task.is_expired(): logger.info(f"Task {task_id} has expired, marking as cancelled") task.status = TaskStatus.CANCELLED task.is_active = False await repo.update(task) return # Mark task as running await repo.mark_as_running(task) # Execute the task try: handler_registry = TaskHandlerRegistry( session, self.credit_service, self.player_service ) await handler_registry.execute_task(task) # Calculate next execution time for recurring tasks next_execution_at = None if task.should_repeat(): next_execution_at = self._calculate_next_execution(task) # Mark as completed await repo.mark_as_completed(task, next_execution_at) # Reschedule if recurring if next_execution_at and task.should_repeat(): # Refresh task to get updated data await session.refresh(task) await self._schedule_apscheduler_job(task) except Exception as e: await repo.mark_as_failed(task, str(e)) logger.exception(f"Task {task_id} execution failed: {str(e)}") finally: self._running_tasks.discard(task_id_str) def _calculate_next_execution(self, task: ScheduledTask) -> Optional[datetime]: """Calculate the next execution time for a recurring task.""" now = datetime.utcnow() if task.recurrence_type == RecurrenceType.HOURLY: return now + timedelta(hours=1) elif task.recurrence_type == RecurrenceType.DAILY: return now + timedelta(days=1) elif task.recurrence_type == RecurrenceType.WEEKLY: return now + timedelta(weeks=1) elif task.recurrence_type == RecurrenceType.MONTHLY: # Add approximately one month return now + timedelta(days=30) elif task.recurrence_type == RecurrenceType.YEARLY: return now + timedelta(days=365) return None async def _maintenance_job(self) -> None: """Periodic maintenance job to clean up expired tasks and handle scheduling issues.""" try: async with self.db_session_factory() as session: repo = ScheduledTaskRepository(session) # Handle expired tasks expired_tasks = await repo.get_expired_tasks() for task in expired_tasks: task.status = TaskStatus.CANCELLED task.is_active = False await repo.update(task) # Remove from scheduler try: self.scheduler.remove_job(str(task.id)) except Exception: pass if expired_tasks: logger.info(f"Cleaned up {len(expired_tasks)} expired tasks") # Handle any missed recurring tasks due_recurring = await repo.get_recurring_tasks_due_for_next_execution() for task in due_recurring: if task.should_repeat(): task.status = TaskStatus.PENDING task.scheduled_at = task.next_execution_at or datetime.utcnow() await repo.update(task) await self._schedule_apscheduler_job(task) if due_recurring: logger.info(f"Rescheduled {len(due_recurring)} recurring tasks") except Exception: logger.exception("Maintenance job failed")