"""
Abstract base class for worker pools.

Provides a unified interface for CPU and GPU worker pools with proper
initialization, task submission, and resource cleanup.
"""

from __future__ import annotations

import logging
import multiprocessing as mp
from abc import ABC, abstractmethod
from concurrent.futures import Future, ProcessPoolExecutor
from dataclasses import dataclass, field
from typing import Any, Callable, Optional

logger = logging.getLogger(__name__)


@dataclass
class TaskResult:
    """Container for task execution results."""

    task_id: str
    success: bool
    data: Any
    error: Optional[str] = None
    processing_time: float = 0.0
    pool_type: str = ""
    extra: dict = field(default_factory=dict)


class WorkerPool(ABC):
    """
    Abstract base class for worker pools.

    Provides a common interface for ProcessPoolExecutor-based worker pools
    with proper initialization using the 'spawn' start method for CUDA
    compatibility.

    Attributes:
        max_workers: Maximum number of worker processes.
        use_gpu: Whether this pool uses GPU resources.
        gpu_id: GPU device ID (only relevant if use_gpu=True).
    """

    def __init__(
        self,
        max_workers: int,
        use_gpu: bool = False,
        gpu_id: int = 0,
    ) -> None:
        """
        Initialize the worker pool configuration.

        Args:
            max_workers: Maximum number of worker processes.
            use_gpu: Whether this pool uses GPU resources.
            gpu_id: GPU device ID for GPU pools.
        """
        self.max_workers = max_workers
        self.use_gpu = use_gpu
        self.gpu_id = gpu_id
        self._executor: Optional[ProcessPoolExecutor] = None
        self._started = False

    @property
    def name(self) -> str:
        """Return the pool name for logging."""
        return self.__class__.__name__

    @abstractmethod
    def get_initializer(self) -> Optional[Callable[..., None]]:
        """
        Return the worker initialization function.

        This function is called once per worker process when it starts.
        Use it to load models, set environment variables, etc.

        Returns:
            Callable to initialize each worker, or None if no initialization needed.
        """
        pass

    @abstractmethod
    def get_init_args(self) -> tuple:
        """
        Return arguments for the initializer function.

        Returns:
            Tuple of arguments to pass to the initializer.
        """
        pass

    def start(self) -> None:
        """
        Start the worker pool.

        Creates a ProcessPoolExecutor with the 'spawn' start method
        for CUDA compatibility.

        Raises:
            RuntimeError: If the pool is already started.
        """
        if self._started:
            raise RuntimeError(f"{self.name} is already started")

        # Use 'spawn' for CUDA compatibility
        ctx = mp.get_context("spawn")

        initializer = self.get_initializer()
        initargs = self.get_init_args()

        logger.info(
            f"Starting {self.name} with {self.max_workers} workers "
            f"(GPU: {self.use_gpu}, GPU ID: {self.gpu_id})"
        )

        self._executor = ProcessPoolExecutor(
            max_workers=self.max_workers,
            mp_context=ctx,
            initializer=initializer,
            initargs=initargs if initializer else (),
        )
        self._started = True

    def submit(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Future:
        """
        Submit a task to the worker pool.

        Args:
            fn: Function to execute.
            *args: Positional arguments for the function.
            **kwargs: Keyword arguments for the function.

        Returns:
            Future representing the pending result.

        Raises:
            RuntimeError: If the pool is not started.
        """
        if not self._started or self._executor is None:
            raise RuntimeError(f"{self.name} is not started. Call start() first.")

        return self._executor.submit(fn, *args, **kwargs)

    def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
        """
        Shutdown the worker pool.

        Args:
            wait: If True, wait for all pending futures to complete.
            cancel_futures: If True, cancel all pending futures.
        """
        if self._executor is not None:
            logger.info(f"Shutting down {self.name} (wait={wait})")
            self._executor.shutdown(wait=wait, cancel_futures=cancel_futures)
            self._executor = None
        self._started = False

    @property
    def is_running(self) -> bool:
        """Check if the pool is currently running."""
        return self._started and self._executor is not None

    def __enter__(self) -> "WorkerPool":
        """Context manager entry - start the pool."""
        self.start()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit - shutdown the pool."""
        self.shutdown(wait=True)

    def __repr__(self) -> str:
        status = "running" if self.is_running else "stopped"
        return (
            f"{self.__class__.__name__}("
            f"workers={self.max_workers}, "
            f"gpu={self.use_gpu}, "
            f"status={status})"
        )
