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

"""Background task executor - Phase 10.4"""

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

from .manager import (
    BackgroundTask,
    BackgroundTaskManager,
    BackgroundTaskStatus,
    get_background_task_manager,
)


class BackgroundExecutor:
    """Executes background tasks with error handling and result storage.

    Provides methods to execute tasks synchronously or asynchronously,
    with full integration into BackgroundTaskManager for tracking.
    """

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

        Args:
            task_manager: Optional BackgroundTaskManager instance.
                         If not provided, uses the global singleton.
        """
        self._task_manager = task_manager or get_background_task_manager()
        self._executors: dict[str, asyncio.Task] = {}

    async def execute_task(
        self,
        task_id: str,
        func: Callable[..., Coroutine[Any, Any, Any]],
        *args: Any,
        **kwargs: Any,
    ) -> BackgroundTask:
        """Execute a specific task by ID.

        Args:
            task_id: Unique task identifier
            func: Async function to execute
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function

        Returns:
            The BackgroundTask with result or error
        """
        # Get or create task record
        task = self._task_manager.get_task_status(task_id)
        if task is None:
            # Create a new task record if one doesn't exist
            task = BackgroundTask(
                id=task_id,
                name=f"executor_task_{task_id}",
                status=BackgroundTaskStatus.PENDING,
                created_at=datetime.now(),
            )
            self._task_manager._tasks[task_id] = task

        # Update status to running
        task.status = BackgroundTaskStatus.RUNNING
        task.started_at = datetime.now()

        try:
            # Execute the async function
            result = await func(*args, **kwargs)
            task.status = BackgroundTaskStatus.COMPLETED
            task.result = result
        except Exception as e:
            task.status = BackgroundTaskStatus.FAILED
            task.error = f"{type(e).__name__}: {str(e)}"
            task.result = None
        finally:
            task.completed_at = datetime.now()
            # Clean up executor reference
            if task_id in self._executors:
                del self._executors[task_id]

        return task

    async def execute_async(
        self,
        task_id: str,
        func: Callable[..., Coroutine[Any, Any, Any]],
        *args: Any,
        **kwargs: Any,
    ) -> str:
        """Execute a task asynchronously in the background.

        Args:
            task_id: Unique task identifier
            func: Async function to execute
            *args: Positional arguments for the function
            **kwargs: Keyword arguments for the function

        Returns:
            The task ID
        """
        # Create task record if it doesn't exist
        if self._task_manager.get_task_status(task_id) is None:
            self._task_manager._tasks[task_id] = BackgroundTask(
                id=task_id,
                name=f"async_task_{task_id}",
                status=BackgroundTaskStatus.PENDING,
                created_at=datetime.now(),
            )

        # Create and store the asyncio task
        async_task = asyncio.create_task(self.execute_task(task_id, func, *args, **kwargs))
        self._executors[task_id] = async_task

        return task_id

    def cancel_task(self, task_id: str) -> bool:
        """Cancel a running task.

        Args:
            task_id: The task ID to cancel

        Returns:
            True if cancelled, False if not found or not running
        """
        if task_id not in self._executors:
            return False

        self._executors[task_id].cancel()
        del self._executors[task_id]

        # Update task status
        task = self._task_manager.get_task_status(task_id)
        if task:
            task.status = BackgroundTaskStatus.CANCELLED
            task.completed_at = datetime.now()
            return True
        return False

    def get_task_result(self, task_id: str) -> Any:
        """Get the result of a completed task.

        Args:
            task_id: The task ID

        Returns:
            The task result or None if not found/not completed
        """
        task = self._task_manager.get_task_status(task_id)
        if task and task.status == BackgroundTaskStatus.COMPLETED:
            return task.result
        return None

    def get_task_error(self, task_id: str) -> str | None:
        """Get the error of a failed task.

        Args:
            task_id: The task ID

        Returns:
            The error message or None if not found/not failed
        """
        task = self._task_manager.get_task_status(task_id)
        if task and task.status == BackgroundTaskStatus.FAILED:
            return task.error
        return None

    def is_task_running(self, task_id: str) -> bool:
        """Check if a task is currently running.

        Args:
            task_id: The task ID

        Returns:
            True if running, False otherwise
        """
        return task_id in self._executors

    def wait_for_task(self, task_id: str, timeout: float | None = None) -> BackgroundTask:
        """Wait for a task to complete.

        Args:
            task_id: The task ID to wait for
            timeout: Optional timeout in seconds

        Returns:
            The completed BackgroundTask

        Raises:
            asyncio.TimeoutError: If task doesn't complete within timeout
            asyncio.CancelledError: If task is cancelled
        """
        if task_id not in self._executors:
            task = self._task_manager.get_task_status(task_id)
            if task:
                return task
            raise ValueError(f"Task {task_id} not found")

        async def wait_task() -> BackgroundTask:
            await self._executors[task_id]
            return self._task_manager.get_task_status(task_id)

        return asyncio.run_until_complete(asyncio.wait_for(wait_task(), timeout=timeout))

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


# Global executor instance
_executor: BackgroundExecutor | None = None


def get_background_executor() -> BackgroundExecutor:
    """Get the global BackgroundExecutor instance."""
    global _executor
    if _executor is None:
        _executor = BackgroundExecutor()
    return _executor