"""
Task pool for managing multiple tasks in memory
Provides task storage and status tracking (not actual multiprocessing)
"""
import logging
from concurrent.futures import ProcessPoolExecutor, Future
from multiprocessing import get_context
from .task import Task, TaskState

logger = logging.getLogger(__name__)


class TaskPool:
    """Manages task storage and status tracking"""

    def __init__(self, max_workers: int = 2):
        """Initialize task pool with maximum workers"""
        self.max_workers = max_workers
        self._tasks: dict[str, Task] = {}
        self.executor = ProcessPoolExecutor(
            max_workers=max_workers, mp_context=get_context("spawn")
        )

    def submit(self, task_id: str, task_type: str) -> Task:
        """Submit a new task to the pool"""
        task = Task(task_id=task_id, task_type=task_type)
        self._tasks[task_id] = task
        return task

    def submit_work(self, task_id: str, func, *args) -> Future:
        """Submit work to the process pool executor"""
        logger.debug("submit_work task_id=%s", task_id)
        return self.executor.submit(func, *args)

    def update_stage(self, task_id: str, stage: str):
        """Update the stage of a task"""
        t = self._tasks.get(task_id)
        if t:
            t.stage = stage

    def get_status(self, task_id: str) -> TaskState | None:
        """Get task status by ID"""
        task = self._tasks.get(task_id)
        return task.status if task else None

    def get_stage(self, task_id: str) -> str | None:
        """Get task stage by ID"""
        t = self._tasks.get(task_id)
        return t.stage if t else None

    def get_task(self, task_id: str) -> Task | None:
        """Get task object by ID"""
        return self._tasks.get(task_id)