"""
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 concurrent.futures.process import BrokenProcessPool
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.

        worker 被 OOM 击杀后 ProcessPoolExecutor 永久破损(BrokenProcessPool),
        重建池重试一次,避免一次异常让任务中心瘫痪到容器重启。
        """
        logger.debug("submit_work task_id=%s", task_id)
        try:
            return self.executor.submit(func, *args)
        except BrokenProcessPool:
            logger.warning("pool broken (worker died?), rebuilding for %s", task_id)
            self.executor.shutdown(wait=False)
            self.executor = ProcessPoolExecutor(
                max_workers=self.max_workers, mp_context=get_context("spawn")
            )
            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)