"""
OCR Worker Pool Manager

Manages a ProcessPoolExecutor with persistent OCR engine initialization.
Key features:
- ProcessPoolExecutor with configurable max_workers (from OCR_WORKERS env)
- Configurable max_tasks_per_child (from OCR_MAX_TASKS_PER_CHILD env, 0=no restart)
- mp_context='spawn' for Windows IIS compatibility
- docTR/PaddleOCR loaded ONCE at worker spawn (not 30s per request)
- atexit + signal handlers for cleanup
- Health check with auto-respawn
- Orphan process cleanup on Windows

Architecture:
    Main Process          │      Worker Process (PERSISTENT)
    ──────────────────────│──────────────────────────────────
    OCRWorkerPool         │      Worker initialized once
        ↓                 │           ↓
    submit_task() ────────│────→  process_ocr()
        ↓                 │           ↓
    Future.result() ←─────│────  Return result
"""

import asyncio
import atexit
import gc
import logging
import multiprocessing as mp
import os
import signal
import sys
import time
from concurrent.futures import ProcessPoolExecutor, Future, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Callable, Optional

logger = logging.getLogger(__name__)

# Try to import psutil for orphan process cleanup
try:
    import psutil
    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False
    logger.warning("[OCRWorkerPool] psutil not available - orphan cleanup disabled")


class OCRWorkerPool:
    """
    Singleton manager for OCR ProcessPoolExecutor.

    Ensures OCR engines are loaded once and reused for all requests.
    Uses max_tasks_per_child=5 to restart worker every 5 tasks (prevents memory leak).
    """

    _instance: Optional["OCRWorkerPool"] = None
    _initialized: bool = False

    def __new__(cls) -> "OCRWorkerPool":
        """Singleton pattern - only one pool instance."""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        """Initialize worker pool (runs only once due to singleton)."""
        if self._initialized:
            return

        self._executor: Optional[ProcessPoolExecutor] = None
        self._worker_pid: Optional[int] = None
        self._is_warming: bool = False
        self._is_shutdown: bool = False
        self._lock = asyncio.Lock() if asyncio.get_event_loop_policy() else None
        self._sync_lock = mp.Lock()

        # Register cleanup handlers
        # NOTE: Only use atexit, NOT signal handlers!
        # Signal handlers interfere with FastAPI's shutdown handling.
        # FastAPI's shutdown event calls stop_job_worker() which calls shutdown().
        atexit.register(self._cleanup_on_exit)

        self._initialized = True
        logger.info("[OCRWorkerPool] Singleton instance created")

    def initialize(self) -> bool:
        """
        Initialize the ProcessPoolExecutor.

        Creates executor with spawn context for Windows compatibility.
        Uses max_tasks_per_child=5 to restart worker periodically (prevents memory leak).

        Returns:
            True if initialization successful
        """
        if self._executor is not None:
            logger.warning("[OCRWorkerPool] Already initialized")
            return True

        if self._is_shutdown:
            logger.error("[OCRWorkerPool] Cannot initialize - pool is shutdown")
            return False

        try:
            # Cleanup any orphan workers from previous runs
            self._cleanup_orphan_workers()

            # Read configuration from environment
            max_workers = int(os.getenv('OCR_WORKERS', '2'))
            max_tasks_raw = os.getenv('OCR_MAX_TASKS_PER_CHILD', '0')
            # 0 means no restart (None in ProcessPoolExecutor)
            max_tasks_per_child = int(max_tasks_raw) if max_tasks_raw and int(max_tasks_raw) > 0 else None

            # Create executor with spawn context (Windows compatible)
            # Use mp_context='spawn' explicitly for cross-platform consistency
            mp_context = mp.get_context('spawn')

            # max_tasks_per_child only available in Python 3.11+
            executor_kwargs = {
                'max_workers': max_workers,
                'mp_context': mp_context,
                'initializer': _worker_initializer,
            }
            if sys.version_info >= (3, 11) and max_tasks_per_child is not None:
                executor_kwargs['max_tasks_per_child'] = max_tasks_per_child
            else:
                logger.info(f"[OCRWorkerPool] max_tasks_per_child not supported (Python {sys.version_info.major}.{sys.version_info.minor})")

            self._executor = ProcessPoolExecutor(**executor_kwargs)

            logger.info(f"[OCRWorkerPool] ProcessPoolExecutor created (spawn context, max_workers={max_workers}, max_tasks_per_child={max_tasks_per_child})")
            return True

        except Exception as e:
            logger.error(f"[OCRWorkerPool] Initialization failed: {e}")
            return False

    async def prewarm(self, timeout: float = 60.0) -> bool:
        """
        Pre-warm the worker by loading PaddleOCR before first request.

        This is called at FastAPI startup to avoid 30s delay on first request.
        Submits a dummy task that triggers PaddleOCR initialization.

        Args:
            timeout: Maximum seconds to wait for warmup (default 60s)

        Returns:
            True if warmup successful, False if timeout or error
        """
        if self._executor is None:
            logger.error("[OCRWorkerPool] Cannot prewarm - not initialized")
            return False

        if self._is_warming:
            logger.warning("[OCRWorkerPool] Already warming up")
            return False

        self._is_warming = True
        logger.info("[OCRWorkerPool] Starting pre-warm (loading PaddleOCR in worker)...")
        start_time = time.time()

        try:
            # Submit warmup task that initializes PaddleOCR
            loop = asyncio.get_event_loop()
            future = self._executor.submit(_warmup_task)

            # Wait with timeout
            result = await loop.run_in_executor(None, future.result, timeout)

            elapsed = time.time() - start_time
            if result.get("success"):
                logger.info(f"[OCRWorkerPool] Pre-warm complete in {elapsed:.1f}s - PaddleOCR ready")
                self._worker_pid = result.get("pid")
                return True
            else:
                logger.error(f"[OCRWorkerPool] Pre-warm failed: {result.get('error')}")
                return False

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error(f"[OCRWorkerPool] Pre-warm failed after {elapsed:.1f}s: {e}")
            return False
        finally:
            self._is_warming = False

    async def submit_task(
        self,
        image_bytes: bytes,
        engine: str = "doctr_plus",
        preprocessing: str = "auto",
        timeout: float = 120.0
    ) -> dict:
        """
        Submit OCR task to worker process.

        Args:
            image_bytes: Raw image bytes
            engine: OCR engine ('tesseract', 'doctr', 'doctr_plus', 'paddleocr')
            preprocessing: Preprocessing mode ('light', 'medium', 'heavy', 'auto')
            timeout: Maximum processing time in seconds

        Returns:
            Dict with extraction results

        Raises:
            RuntimeError: If pool not initialized or task fails
        """
        if self._executor is None:
            raise RuntimeError("OCR worker pool not initialized")

        if self._is_shutdown:
            raise RuntimeError("OCR worker pool is shutdown")

        logger.info(f"[OCRWorkerPool] Submitting task: engine={engine}, preprocessing={preprocessing}, size={len(image_bytes)} bytes")

        try:
            loop = asyncio.get_event_loop()
            future = self._executor.submit(
                _process_ocr_task,
                image_bytes,
                engine,
                preprocessing
            )

            # Wait for result with timeout
            result = await loop.run_in_executor(None, future.result, timeout)

            logger.info(f"[OCRWorkerPool] Task complete: success={result.get('success')}")
            return result

        except TimeoutError:
            logger.error(f"[OCRWorkerPool] Task timed out after {timeout}s")
            raise RuntimeError(f"OCR processing timed out after {timeout}s")

        except Exception as e:
            logger.error(f"[OCRWorkerPool] Task failed: {e}")
            raise RuntimeError(f"OCR processing failed: {e}")

    def is_healthy(self) -> bool:
        """
        Check if worker pool is healthy.

        Returns:
            True if pool is ready to accept tasks
        """
        if self._executor is None:
            return False
        if self._is_shutdown:
            return False

        # Check if worker process is still alive
        if self._worker_pid and PSUTIL_AVAILABLE:
            try:
                proc = psutil.Process(self._worker_pid)
                if not proc.is_running():
                    logger.warning("[OCRWorkerPool] Worker process died, needs respawn")
                    return False
            except psutil.NoSuchProcess:
                logger.warning("[OCRWorkerPool] Worker process not found")
                return False

        return True

    def shutdown(self, wait: bool = True, timeout: float = 10.0) -> None:
        """
        Shutdown the worker pool gracefully.

        Args:
            wait: Wait for pending tasks to complete
            timeout: Maximum wait time in seconds
        """
        if self._executor is None:
            return

        logger.info("[OCRWorkerPool] Shutting down...")
        self._is_shutdown = True

        try:
            self._executor.shutdown(wait=wait, cancel_futures=True)
            logger.info("[OCRWorkerPool] Executor shutdown complete")
        except Exception as e:
            logger.error(f"[OCRWorkerPool] Shutdown error: {e}")

        self._executor = None
        self._worker_pid = None

        # Final orphan cleanup
        self._cleanup_orphan_workers()
        logger.info("[OCRWorkerPool] Shutdown complete")

    def _cleanup_orphan_workers(self) -> int:
        """
        Clean up orphan Python processes from previous runs.

        On Windows with NSSM, orphan processes may remain after service restart.
        This finds and kills any python.exe processes that were OCR workers.

        Returns:
            Number of processes killed
        """
        if not PSUTIL_AVAILABLE:
            return 0

        killed = 0
        current_pid = os.getpid()

        try:
            for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
                try:
                    # Skip self
                    if proc.pid == current_pid:
                        continue

                    # Look for Python processes with OCR-related cmdline
                    if proc.name().lower() in ('python.exe', 'python3.exe', 'python', 'python3'):
                        cmdline = ' '.join(proc.cmdline() or [])

                        # Check if this is an OCR worker process
                        if 'ocr_worker_process' in cmdline.lower() or 'process_ocr_task' in cmdline.lower():
                            logger.warning(f"[OCRWorkerPool] Killing orphan worker: PID={proc.pid}")
                            proc.kill()
                            proc.wait(timeout=5)
                            killed += 1

                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    continue

        except Exception as e:
            logger.error(f"[OCRWorkerPool] Orphan cleanup error: {e}")

        if killed > 0:
            logger.info(f"[OCRWorkerPool] Cleaned up {killed} orphan worker(s)")

        return killed

    def _cleanup_on_exit(self) -> None:
        """atexit handler for cleanup."""
        logger.info("[OCRWorkerPool] atexit cleanup triggered")
        self.shutdown(wait=False)

    def _signal_handler(self, signum: int, frame: Any) -> None:
        """Signal handler for SIGTERM/SIGINT."""
        logger.info(f"[OCRWorkerPool] Received signal {signum}, shutting down...")
        self.shutdown(wait=False)


# ============================================================================
# WORKER PROCESS FUNCTIONS
# ============================================================================
# These functions run in the child process, not the main FastAPI process.

# Global engines - persist between tasks in worker process
_paddle_engine = None
_tesseract_engine = None
_doctr_engine = None  # docTR engine (PyTorch backend)
_worker_initialized = False


def _worker_initializer() -> None:
    """
    Called once when worker process spawns.

    Initializes global OCR engines IN PARALLEL for faster startup.
    Uses ThreadPoolExecutor to load enabled engines concurrently.
    Respects OCR_ENABLE_PADDLEOCR and OCR_ENABLE_TESSERACT from .env.

    Total warmup time = max(engine_times) instead of sum(engine_times).
    """
    global _paddle_engine, _tesseract_engine, _doctr_engine, _worker_initialized

    if _worker_initialized:
        print(f"[Worker {os.getpid()}] Already initialized", flush=True)
        return

    # Check which engines are enabled via .env
    paddle_enabled = os.getenv("OCR_ENABLE_PADDLEOCR", "true").lower() == "true"
    tesseract_enabled = os.getenv("OCR_ENABLE_TESSERACT", "true").lower() == "true"

    enabled_engines = ["doctr"]  # docTR is always loaded (primary engine)
    if paddle_enabled:
        enabled_engines.append("paddle")
    if tesseract_enabled:
        enabled_engines.append("tesseract")

    print(f"[Worker {os.getpid()}] Initializing OCR engines: {enabled_engines}", flush=True)
    if not paddle_enabled:
        print(f"[Worker {os.getpid()}] PaddleOCR DISABLED - saving ~800MB RAM", flush=True)
    if not tesseract_enabled:
        print(f"[Worker {os.getpid()}] Tesseract DISABLED - saving ~50MB RAM", flush=True)

    start_time = time.time()

    # Define loader functions - each runs in its own thread
    def load_doctr():
        try:
            from backend.modules.data_entry.services.ocr.ocr_worker_process import initialize_doctr_engine
            engine = initialize_doctr_engine()
            return ("doctr", engine, None)
        except Exception as e:
            return ("doctr", None, str(e))

    def load_paddle():
        if not paddle_enabled:
            return ("paddle", None, "disabled via OCR_ENABLE_PADDLEOCR=false")
        try:
            from backend.modules.data_entry.services.ocr.ocr_worker_process import initialize_paddle_engine
            engine = initialize_paddle_engine()
            return ("paddle", engine, None)
        except Exception as e:
            return ("paddle", None, str(e))

    def load_tesseract():
        if not tesseract_enabled:
            return ("tesseract", None, "disabled via OCR_ENABLE_TESSERACT=false")
        try:
            from backend.modules.data_entry.services.ocr.tesseract_engine import TesseractEngine
            engine = TesseractEngine()
            return ("tesseract", engine, None)
        except Exception as e:
            return ("tesseract", None, str(e))

    # Build list of futures for enabled engines only
    futures_to_submit = [load_doctr]  # docTR always loaded
    if paddle_enabled:
        futures_to_submit.append(load_paddle)
    if tesseract_enabled:
        futures_to_submit.append(load_tesseract)

    # Load engines in parallel using ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=len(futures_to_submit)) as executor:
        futures = [executor.submit(fn) for fn in futures_to_submit]

        for future in as_completed(futures):
            name, engine, error = future.result()
            if error and "disabled" not in error:
                print(f"[Worker {os.getpid()}] {name} init failed: {error}", flush=True)
            elif engine:
                print(f"[Worker {os.getpid()}] {name} loaded", flush=True)
                if name == "doctr":
                    _doctr_engine = engine
                elif name == "paddle":
                    _paddle_engine = engine
                elif name == "tesseract":
                    _tesseract_engine = engine

    elapsed = time.time() - start_time
    _worker_initialized = True
    print(f"[Worker {os.getpid()}] Initialization complete in {elapsed:.1f}s (engines: {enabled_engines})", flush=True)


def _warmup_task() -> dict:
    """
    Warmup task that ensures engines are loaded.

    Called at FastAPI startup to pre-warm the worker.
    Returns success status and worker PID.
    """
    global _paddle_engine, _tesseract_engine, _doctr_engine, _worker_initialized

    try:
        # Ensure initialization
        if not _worker_initialized:
            _worker_initializer()

        # Quick test - create a small dummy image
        import numpy as np
        dummy_img = np.ones((100, 100, 3), dtype=np.uint8) * 255

        # Test docTR if available (fastest engine)
        if _doctr_engine is not None:
            try:
                _doctr_engine([dummy_img])
                print(f"[Worker {os.getpid()}] docTR warmup OK", flush=True)
            except Exception as e:
                print(f"[Worker {os.getpid()}] docTR warmup error: {e}", flush=True)

        # Test PaddleOCR if available
        if _paddle_engine is not None:
            try:
                _paddle_engine.predict(dummy_img)
                print(f"[Worker {os.getpid()}] PaddleOCR warmup OK", flush=True)
            except Exception as e:
                print(f"[Worker {os.getpid()}] PaddleOCR warmup error: {e}", flush=True)

        # Cleanup
        gc.collect()

        return {
            "success": True,
            "pid": os.getpid(),
            "doctr_available": _doctr_engine is not None,
            "paddle_available": _paddle_engine is not None,
            "tesseract_available": _tesseract_engine is not None
        }

    except Exception as e:
        return {
            "success": False,
            "pid": os.getpid(),
            "error": str(e)
        }


def _process_ocr_task(
    image_bytes: bytes,
    engine: str = "doctr_plus",
    preprocessing: str = "auto"
) -> dict:
    """
    Process OCR task in worker process.

    This is the main work function called for each OCR request.
    Uses persistent global engines loaded at worker init.

    Args:
        image_bytes: Raw image bytes
        engine: OCR engine choice ('tesseract', 'doctr', 'doctr_plus', 'paddleocr')
        preprocessing: Preprocessing mode

    Returns:
        Dict with extraction results
    """
    global _paddle_engine, _tesseract_engine, _doctr_engine, _worker_initialized

    try:
        # Ensure initialization
        if not _worker_initialized:
            _worker_initializer()

        # Import processing function
        from backend.modules.data_entry.services.ocr.ocr_worker_process import process_ocr

        # Run OCR
        result = process_ocr(
            image_bytes=image_bytes,
            paddle_engine=_paddle_engine,
            tesseract_engine=_tesseract_engine,
            engine=engine,
            preprocessing=preprocessing,
            doctr_engine=_doctr_engine
        )

        # Cleanup after each task
        gc.collect()

        return result

    except Exception as e:
        print(f"[Worker {os.getpid()}] Task error: {e}", flush=True)
        import traceback
        traceback.print_exc()
        return {
            "success": False,
            "error": str(e),
            "pid": os.getpid()
        }


# Singleton instance
ocr_worker_pool = OCRWorkerPool()
