Click/managers/background_tasks.py

"""Background task manager for async task tracking."""

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Dict, Optional, Callable, Coroutine, Any

from core.constants import TaskStatus

logger = logging.getLogger(__name__)


@dataclass
class TaskInfo:
    """Information about a running task."""
    task_id: str
    task_type: str
    coroutine: Coroutine = None
    status: TaskStatus = TaskStatus.RUNNING
    created_at: float = field(default_factory=lambda: asyncio.get_running_loop().time())
    metadata: dict = field(default_factory=dict)


class BackgroundTaskManager:
    """
    Manages background async tasks with tracking and cancellation support.
    
    Features:
    - Start and track async tasks
    - Cancel individual or all tasks
    - Get task status and metadata
    - Prevent duplicate tasks
    """
    
    def __init__(self):
        self._tasks: Dict[str, asyncio.Task] = {}
        self._task_info: Dict[str, TaskInfo] = {}
        self._lock = asyncio.Lock()
    
    async def start_task(
        self,
        task_id: str,
        coro: Coroutine,
        task_type: str = "unknown",
        metadata: dict = None
    ) -> bool:
        """
        Start a background task.
        
        Args:
            task_id: Unique task identifier
            coro: Coroutine to run
            task_type: Type of task
            metadata: Additional task metadata
        
        Returns:
            True if started successfully, False if already running
        """
        async with self._lock:
            if task_id in self._tasks and not self._tasks[task_id].done():
                logger.warning(f"Task {task_id} is already running")
                return False
            
            task = asyncio.create_task(coro)
            self._tasks[task_id] = task
            
            self._task_info[task_id] = TaskInfo(
                task_id=task_id,
                task_type=task_type,
                coroutine=coro,
                metadata=metadata or {}
            )
            
            task.add_done_callback(lambda t, tid=task_id: self._on_task_done(tid, t))
            
            logger.info(f"Started background task: {task_id} ({task_type})")
            return True
    
    def _on_task_done(self, task_id: str, task: asyncio.Task) -> None:
        """Callback когда задача завершилась."""
        try:
            exc = task.exception()
            if exc:
                if isinstance(exc, asyncio.CancelledError):
                    logger.info(f"Task {task_id} cancelled (normal)")
                    # Не считаем ошибкой
                else:
                    logger.error(f"Task {task_id} failed: {exc}")
            else:
                logger.info(f"Task {task_id} completed successfully")
        except asyncio.CancelledError: 
            logger.info(f"Task {task_id} was cancelled")
        except Exception as e:
            logger.error(f"Error in task callback: {e}")
    
    async def cancel_task(self, task_id: str) -> bool:
        """
        Cancel a specific task.
        
        Args:
            task_id: Task identifier to cancel
        
        Returns:
            True if task was cancelled, False if not found
        """
        async with self._lock:
            task = self._tasks.get(task_id)
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
                
                logger.info(f"Task cancelled: {task_id}")
                return True
            return False
    
    async def cancel_all(self) -> int:
        """
        Cancel all running tasks.
        
        Returns:
            Number of tasks cancelled
        """
        cancelled = 0
        
        async with self._lock:
            for task_id, task in list(self._tasks.items()):
                if not task.done():
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
                    cancelled += 1
            
            self._tasks.clear()
            self._task_info.clear()
        
        logger.info(f"All tasks cancelled: {cancelled}")
        return cancelled
    
    def get_active_tasks(self) -> Dict[str, dict]:
        """
        Get all active tasks with their metadata.
        
        Returns:
            Dictionary of active tasks
        """
        active = {}
        for task_id, task in self._tasks.items():
            if not task.done():
                info = self._task_info.get(task_id)
                active[task_id] = {
                    "info": info.metadata if info else {},
                    "status": TaskStatus.RUNNING.value,
                    "type": info.task_type if info else "unknown",
                    "created_at": info.created_at if info else 0
                }
        return active
    
    @property
    def active_count(self) -> int:
        """Number of currently active tasks."""
        return sum(1 for t in self._tasks.values() if not t.done())
    
    def has_active_tasks(self) -> bool:
        """Check if there are any active tasks."""
        return self.active_count > 0