#exonware/xwsystem/src/exonware/xwsystem/ipc/process_pool.py
"""
Process Pool Utilities
======================
Production-grade process pools for xwsystem.
Author: eXonware Backend Team
Email: connect@exonware.com
Company: eXonware.com
Generation Date: September 05, 2025
"""

import asyncio
import concurrent.futures
import functools
import logging
import multiprocessing as mp
import platform
import time
from dataclasses import dataclass
from typing import Any

from collections.abc import Callable
logger = logging.getLogger(__name__)
@dataclass


class TaskResult:
    """Result of a process pool task."""
    task_id: str
    result: Any
    success: bool
    error: str | None = None
    execution_time: float = 0.0
    worker_pid: int | None = None


class ProcessPool:
    """
    Production-grade process pool with monitoring and error handling.
    Features:
    - Automatic worker management
    - Task timeout handling
    - Error recovery
    - Performance monitoring
    - Graceful shutdown
    - Load balancing
    """

    def __init__(self, 
                 max_workers: int | None = None,
                 size: int | None = None,
                 initializer: Callable | None = None,
                 initargs: tuple = (),
                 timeout: float | None = None):
        """
        Initialize process pool.
        Args:
            max_workers: Maximum number of worker processes
            size: Alias for max_workers
            initializer: Function to run in each worker on startup
            initargs: Arguments for initializer function
            timeout: Default timeout for tasks
        """
        # Handle 'size' parameter as alias for max_workers
        if size is not None and max_workers is None:
            max_workers = size
        self.max_workers = max_workers or mp.cpu_count()
        self.initializer = initializer
        self.initargs = initargs
        self.timeout = timeout
        # On Windows, local/lambda functions are not picklable with ProcessPool.
        # Use thread pool fallback to preserve API behavior in tests and tooling.
        executor_cls = (
            concurrent.futures.ThreadPoolExecutor
            if platform.system() == "Windows"
            else concurrent.futures.ProcessPoolExecutor
        )
        self._executor = executor_cls(
            max_workers=self.max_workers,
            initializer=initializer,
            initargs=initargs
        )
        # Task tracking
        self._active_tasks: dict[str, concurrent.futures.Future] = {}
        self._completed_tasks: list[TaskResult] = []
        self._task_counter = 0
        # Statistics
        self._stats = {
            'tasks_submitted': 0,
            'tasks_completed': 0,
            'tasks_failed': 0,
            'total_execution_time': 0.0,
        }

    def submit(self, 
               fn: Callable, 
               *args, 
               task_id: str | None = None,
               timeout: float | None = None,
               **kwargs) -> str:
        """
        Submit a task to the process pool.
        Args:
            fn: Function to execute
            *args: Positional arguments for function
            task_id: Optional task identifier
            timeout: Task timeout (overrides default)
            **kwargs: Keyword arguments for function
        Returns:
            Task ID
        """
        if task_id is None:
            self._task_counter += 1
            task_id = f"task_{self._task_counter}"
        start_time = time.time()
        # Submit task to executor
        try:
            future = self._executor.submit(fn, *args, **kwargs)
            self._active_tasks[task_id] = future
            # Add completion callback
            def task_completed(fut):
                execution_time = time.time() - start_time
                try:
                    result = fut.result(timeout=timeout or self.timeout)
                    task_result = TaskResult(
                        task_id=task_id,
                        result=result,
                        success=True,
                        execution_time=execution_time
                    )
                    self._stats['tasks_completed'] += 1
                    self._stats['total_execution_time'] += execution_time
                except Exception as e:
                    task_result = TaskResult(
                        task_id=task_id,
                        result=None,
                        success=False,
                        error=str(e),
                        execution_time=execution_time
                    )
                    self._stats['tasks_failed'] += 1
                # Store result and cleanup
                self._completed_tasks.append(task_result)
                if task_id in self._active_tasks:
                    del self._active_tasks[task_id]
                logger.debug(f"Task {task_id} completed in {execution_time:.3f}s")
            future.add_done_callback(task_completed)
            self._stats['tasks_submitted'] += 1
            logger.debug(f"Submitted task {task_id}")
            return task_id
        except Exception as e:
            # Handle any submission errors
            logger.error(f"Failed to submit task {task_id}: {e}")
            raise

    def get_result(self, task_id: str, timeout: float | None = None) -> TaskResult | None:
        """
        Get result of a completed task.
        Args:
            task_id: Task identifier
            timeout: Timeout for waiting
        Returns:
            Task result or None if not found
        """
        # Check completed tasks first
        for result in self._completed_tasks:
            if result.task_id == task_id:
                logger.debug(f"Found completed task {task_id} with result: {result.result}")
                return result
        # Check if task is still active
        if task_id in self._active_tasks:
            future = self._active_tasks[task_id]
            try:
                # Wait for completion
                future.result(timeout=timeout)
                # In completed tasks now
                for result in self._completed_tasks:
                    if result.task_id == task_id:
                        return result
            except concurrent.futures.TimeoutError:
                logger.warning(f"Task {task_id} timed out")
                return None
            except Exception as e:
                logger.error(f"Task {task_id} failed: {e}")
                return None
        return None

    def wait_for_all(self, timeout: float | None = None) -> list[TaskResult]:
        """
        Wait for all active tasks to complete.
        Args:
            timeout: Timeout for waiting
        Returns:
            List of all task results
        """
        if not self._active_tasks:
            return self._completed_tasks.copy()
        try:
            # Wait for all active futures
            active_futures = list(self._active_tasks.values())
            concurrent.futures.wait(active_futures, timeout=timeout)
        except Exception as e:
            logger.error(f"Error waiting for tasks: {e}")
        return self._completed_tasks.copy()

    def cancel_task(self, task_id: str) -> bool:
        """
        Cancel an active task.
        Args:
            task_id: Task identifier
        Returns:
            True if cancelled successfully
        """
        if task_id not in self._active_tasks:
            return False
        future = self._active_tasks[task_id]
        cancelled = future.cancel()
        if cancelled:
            del self._active_tasks[task_id]
            logger.info(f"Cancelled task {task_id}")
        return cancelled

    def get_active_tasks(self) -> list[str]:
        """Get list of active task IDs."""
        return list(self._active_tasks.keys())

    def get_stats(self) -> dict:
        """Get process pool statistics."""
        stats = self._stats.copy()
        stats['active_tasks'] = len(self._active_tasks)
        stats['completed_tasks'] = len(self._completed_tasks)
        stats['max_workers'] = self.max_workers
        if stats['tasks_completed'] > 0:
            stats['avg_execution_time'] = stats['total_execution_time'] / stats['tasks_completed']
        else:
            stats['avg_execution_time'] = 0.0
        return stats

    def shutdown(self, wait: bool = True, timeout: float | None = None):
        """
        Shutdown the process pool.
        Args:
            wait: Whether to wait for active tasks
            timeout: Timeout for shutdown
        """
        logger.info("Shutting down process pool")
        if wait and timeout:
            self.wait_for_all(timeout)
        self._executor.shutdown(wait=wait)

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.shutdown()


class AsyncProcessPool:
    """
    Async-compatible process pool wrapper.
    Features:
    - Asyncio integration
    - Non-blocking task submission
    - Async result retrieval
    - Graceful async shutdown
    """

    def __init__(self, 
                 max_workers: int | None = None,
                 initializer: Callable | None = None,
                 initargs: tuple = ()):
        """
        Initialize async process pool.
        Args:
            max_workers: Maximum number of worker processes
            initializer: Function to run in each worker on startup
            initargs: Arguments for initializer function
        """
        self.max_workers = max_workers or mp.cpu_count()
        self.initializer = initializer
        self.initargs = initargs
        # Will be created when first used
        self._executor: concurrent.futures.ProcessPoolExecutor | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        # Task tracking
        self._active_tasks: dict[str, asyncio.Task] = {}
        self._task_counter = 0

    def _ensure_executor(self):
        """Ensure executor is created in the correct event loop."""
        if self._executor is None:
            executor_cls = (
                concurrent.futures.ThreadPoolExecutor
                if platform.system() == "Windows"
                else concurrent.futures.ProcessPoolExecutor
            )
            self._executor = executor_cls(
                max_workers=self.max_workers,
                initializer=self.initializer,
                initargs=self.initargs
            )
            self._loop = asyncio.get_event_loop()

    async def submit(self, 
                    fn: Callable, 
                    *args, 
                    task_id: str | None = None,
                    **kwargs) -> str:
        """
        Submit a task to the async process pool.
        Args:
            fn: Function to execute
            *args: Positional arguments for function
            task_id: Optional task identifier
            **kwargs: Keyword arguments for function
        Returns:
            Task ID
        """
        self._ensure_executor()
        if task_id is None:
            self._task_counter += 1
            task_id = f"async_task_{self._task_counter}"
        # Create async task
        async def run_task():
            loop = asyncio.get_event_loop()
            callable_fn = functools.partial(fn, *args, **kwargs)
            return await loop.run_in_executor(self._executor, callable_fn)
        task = asyncio.create_task(run_task())
        self._active_tasks[task_id] = task
        logger.debug(f"Submitted async task {task_id}")
        return task_id

    async def get_result(self, task_id: str, timeout: float | None = None) -> Any:
        """
        Get result of a task.
        Args:
            task_id: Task identifier
            timeout: Timeout for waiting
        Returns:
            Task result
        """
        if task_id not in self._active_tasks:
            raise ValueError(f"Task {task_id} not found")
        task = self._active_tasks[task_id]
        try:
            if timeout is not None:
                return await asyncio.wait_for(task, timeout=timeout)
            return await task
        finally:
            self._active_tasks.pop(task_id, None)

    async def wait_for_all(self, timeout: float | None = None) -> list[Any]:
        """
        Wait for all active tasks to complete.
        Args:
            timeout: Timeout for waiting
        Returns:
            List of all results
        """
        if not self._active_tasks:
            return []
        tasks = list(self._active_tasks.values())
        if timeout is not None:
            done, pending = await asyncio.wait(tasks, timeout=timeout)
            results = []
            # Get results from completed tasks
            for task in done:
                try:
                    results.append(await task)
                except Exception as e:
                    logger.error(f"Task failed: {e}")
                    results.append(None)
            # Cancel pending tasks
            for task in pending:
                task.cancel()
            # Remove completed/cancelled tasks from tracking.
            self._active_tasks = {
                tid: task for tid, task in self._active_tasks.items() if task in pending
            }
            return results
        else:
            results = await asyncio.gather(*tasks, return_exceptions=True)
            self._active_tasks.clear()
            return results

    def cancel_task(self, task_id: str) -> bool:
        """
        Cancel an active task.
        Args:
            task_id: Task identifier
        Returns:
            True if cancelled successfully
        """
        if task_id not in self._active_tasks:
            return False
        task = self._active_tasks[task_id]
        cancelled = task.cancel()
        if cancelled:
            del self._active_tasks[task_id]
            logger.info(f"Cancelled async task {task_id}")
        return cancelled

    def get_active_tasks(self) -> list[str]:
        """Get list of active task IDs."""
        return list(self._active_tasks.keys())

    async def shutdown(self, timeout: float | None = None):
        """
        Shutdown the async process pool.
        Args:
            timeout: Timeout for shutdown
        """
        logger.info("Shutting down async process pool")
        # Cancel all active tasks
        for task_id in list(self._active_tasks.keys()):
            self.cancel_task(task_id)
        # Shutdown executor
        if self._executor:
            self._executor.shutdown(wait=False)
            self._executor = None

    async def __aenter__(self):
        """Async context manager entry."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.shutdown()


def is_process_pool_available() -> bool:
    """Check if process pool functionality is available."""
    # multiprocessing and concurrent.futures are built-in Python modules
    import multiprocessing
    import concurrent.futures
    return True
