"""
CPU Worker Pool for text PDF processing.

This pool handles CPU-bound tasks like text extraction from native PDFs
that don't require OCR.
"""

from __future__ import annotations

import logging
import os
from typing import Callable, Optional

from training.processing.worker_pool import WorkerPool

logger = logging.getLogger(__name__)

# Global resources for CPU workers (initialized once per process)
_cpu_initialized: bool = False


def _init_cpu_worker() -> None:
    """
    Initialize a CPU worker process.

    Disables GPU access and sets up CPU-only environment.
    """
    global _cpu_initialized

    # Disable GPU access for CPU workers
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

    # Set threading limits for better CPU utilization
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")

    _cpu_initialized = True
    logger.debug(f"CPU worker initialized in process {os.getpid()}")


class CPUWorkerPool(WorkerPool):
    """
    Worker pool for CPU-bound tasks.

    Handles text PDF processing that doesn't require OCR.
    Each worker is initialized with CUDA disabled to prevent
    accidental GPU memory consumption.

    Example:
        with CPUWorkerPool(max_workers=4) as pool:
            future = pool.submit(process_text_pdf, pdf_path)
            result = future.result()
    """

    def __init__(self, max_workers: int = 4) -> None:
        """
        Initialize CPU worker pool.

        Args:
            max_workers: Number of CPU worker processes.
                        Defaults to 4 for balanced performance.
        """
        super().__init__(max_workers=max_workers, use_gpu=False, gpu_id=-1)

    def get_initializer(self) -> Optional[Callable[..., None]]:
        """Return the CPU worker initializer."""
        return _init_cpu_worker

    def get_init_args(self) -> tuple:
        """Return empty args for CPU initializer."""
        return ()
