JARVIS/backend/app/agents/background/scheduler.py

"""Background task scheduler - Phase 10.4"""

from collections.abc import Callable, Coroutine
from typing import Any

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.base import BaseTrigger

from .manager import BackgroundTaskManager, get_background_task_manager


class BackgroundScheduler:
    """Background task scheduler using APScheduler.

    Integrates with BackgroundTaskManager for task tracking and execution.
    """

    def __init__(self, task_manager: BackgroundTaskManager | None = None):
        """Initialize the scheduler.

        Args:
            task_manager: Optional BackgroundTaskManager instance.
                         If not provided, uses the global singleton.
        """
        self._scheduler = AsyncIOScheduler()
        self._task_manager = task_manager or get_background_task_manager()
        self._job_tasks: dict[str, str] = {}  # Maps APScheduler job_id to task_id

    def add_job(
        self,
        func: Callable[..., Coroutine[Any, Any, Any]],
        trigger: BaseTrigger,
        args: tuple[Any, ...] | None = None,
        kwargs: dict[str, Any] | None = None,
        id: str | None = None,
        name: str | None = None,
        **apscheduler_kwargs: Any,
    ) -> str:
        """Add a job to the scheduler.

        Args:
            func: Async function to execute
            trigger: APScheduler trigger (date, interval, cron, etc.)
            args: Positional arguments for the function
            kwargs: Keyword arguments for the function
            id: Unique job ID (auto-generated if not provided)
            name: Job name for display purposes
            **apscheduler_kwargs: Additional APScheduler options

        Returns:
            The job ID
        """
        job_id = id or f"job_{len(self._job_tasks)}"
        task_name = name or f"scheduled_task_{job_id}"

        # Wrap the async function to integrate with BackgroundTaskManager
        async def wrapped_func() -> None:
            coro = func(*(args or ()), **(kwargs or {}))
            task_id = self._task_manager.submit_task(task_name, coro)
            self._job_tasks[job_id] = task_id

        self._scheduler.add_job(
            wrapped_func,
            trigger=trigger,
            id=job_id,
            name=task_name,
            **apscheduler_kwargs,
        )

        return job_id

    def remove_job(self, job_id: str) -> bool:
        """Remove a job from the scheduler.

        Args:
            job_id: The ID of the job to remove

        Returns:
            True if job was removed, False if job didn't exist
        """
        try:
            self._scheduler.remove_job(job_id)
            # Clean up task mapping if exists
            if job_id in self._job_tasks:
                task_id = self._job_tasks.pop(job_id)
                # Cancel the background task if still running
                self._task_manager.cancel_task(task_id)
            return True
        except Exception:
            return False

    def list_jobs(self) -> list[dict[str, Any]]:
        """List all scheduled jobs.

        Returns:
            List of job information dictionaries
        """
        jobs = self._scheduler.get_jobs()
        return [
            {
                "id": job.id,
                "name": job.name,
                "next_run_time": job.next_run_time,
                "trigger": str(job.trigger),
            }
            for job in jobs
        ]

    def start(self) -> None:
        """Start the scheduler."""
        if not self._scheduler.running:
            self._scheduler.start()

    def shutdown(self, wait: bool = True) -> None:
        """Shutdown the scheduler.

        Args:
            wait: Whether to wait for running jobs to complete
        """
        if self._scheduler.running:
            self._scheduler.shutdown(wait=wait)

    def pause(self) -> None:
        """Pause the scheduler."""
        self._scheduler.pause()

    def resume(self) -> None:
        """Resume the scheduler."""
        self._scheduler.resume()

    @property
    def task_manager(self) -> BackgroundTaskManager:
        """Get the underlying task manager."""
        return self._task_manager


# Global scheduler instance
_scheduler: BackgroundScheduler | None = None


def get_background_scheduler() -> BackgroundScheduler:
    """Get the global BackgroundScheduler instance."""
    global _scheduler
    if _scheduler is None:
        _scheduler = BackgroundScheduler()
    return _scheduler