"""Worker pool manager for parallel gremlin execution.

This module provides the WorkerPool class that manages a pool of worker processes
for running mutation tests in parallel.

Note: This module uses ProcessPoolExecutor which internally uses pickle for
inter-process communication. The data being serialized is entirely under our
control (WorkerResult dataclass with primitive types) - no untrusted external
content is ever deserialized.
"""

from __future__ import annotations

from concurrent.futures import (
    Future,
    ProcessPoolExecutor,
)
from dataclasses import dataclass
import logging
import os
from pathlib import Path
import subprocess
import time
from typing import Self

from pytest_gremlins.parallel.lightweight import build_lightweight_command
from pytest_gremlins.reporting.results import GremlinResultStatus

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class WorkerResult:
    """Result from a worker process.

    This is a simplified result that can be passed between processes.
    Unlike GremlinResult, it doesn't contain AST nodes which cannot be
    serialized for multiprocessing.

    Attributes:
        gremlin_id: The ID of the gremlin that was tested.
        status: The outcome of testing the gremlin.
        killing_test: Name of test that killed the gremlin (if any).
        execution_time_ms: Time taken to test this gremlin.
        error_output: Captured stderr or exception message when status is ERROR.
    """

    gremlin_id: str
    status: GremlinResultStatus
    killing_test: str | None = None
    execution_time_ms: float | None = None
    error_output: str = ''


def _run_gremlin_test(  # pragma: no cover
    gremlin_id: str,
    test_command: list[str],
    rootdir: str,
    env_vars: dict[str, str],
    timeout: int,
) -> WorkerResult:
    """Execute tests for a single gremlin in a worker process.

    This function runs in a separate process and executes the test command
    with the ACTIVE_GREMLIN environment variable set.

    Args:
        gremlin_id: The ID of the gremlin to test.
        test_command: Command to run tests.
        rootdir: Root directory for test execution.
        env_vars: Additional environment variables to set.
        timeout: Timeout in seconds.

    Returns:
        WorkerResult with the outcome of testing.
    """
    start_time = time.monotonic()

    env = os.environ.copy()
    env.update(env_vars)
    env['ACTIVE_GREMLIN'] = gremlin_id
    env['GREMLIN_ROOTDIR'] = rootdir

    # Use lightweight runner if available (skips full pytest startup)
    lightweight_cmd = build_lightweight_command(test_command, env_vars)
    effective_command = lightweight_cmd if lightweight_cmd is not None else test_command

    try:
        result = subprocess.run(  # Intentional: runs pytest test commands
            effective_command,
            cwd=rootdir,
            env=env,
            capture_output=True,
            timeout=timeout,
            check=False,
        )

        execution_time_ms = (time.monotonic() - start_time) * 1000

        # Only return code 1 (tests failed) should be treated as a mutation
        # being zapped. Other non-zero return codes indicate pytest errors
        # (collection/import/internal) and should not inflate the score.
        if result.returncode == 0:
            return WorkerResult(
                gremlin_id=gremlin_id,
                status=GremlinResultStatus.SURVIVED,
                execution_time_ms=execution_time_ms,
            )
        if result.returncode == 1:
            return WorkerResult(
                gremlin_id=gremlin_id,
                status=GremlinResultStatus.ZAPPED,
                killing_test='unknown',
                execution_time_ms=execution_time_ms,
            )
        error_output = ''
        if result.stderr:
            error_output = result.stderr.decode(errors='replace')[:2000]
        return WorkerResult(
            gremlin_id=gremlin_id,
            status=GremlinResultStatus.ERROR,
            execution_time_ms=execution_time_ms,
            error_output=error_output,
        )
    except subprocess.TimeoutExpired:
        execution_time_ms = (time.monotonic() - start_time) * 1000
        return WorkerResult(
            gremlin_id=gremlin_id,
            status=GremlinResultStatus.TIMEOUT,
            execution_time_ms=execution_time_ms,
        )
    except Exception as exc:
        logger.warning('Error testing gremlin %s: %s', gremlin_id, exc)
        execution_time_ms = (time.monotonic() - start_time) * 1000
        return WorkerResult(
            gremlin_id=gremlin_id,
            status=GremlinResultStatus.ERROR,
            execution_time_ms=execution_time_ms,
            error_output=str(exc)[:2000],
        )


class WorkerPool:
    """Manages a pool of worker processes for parallel mutation testing.

    The worker pool wraps a ProcessPoolExecutor and provides lifecycle management
    for parallel gremlin execution. Workers are isolated processes that each set
    their own ACTIVE_GREMLIN environment variable.

    Attributes:
        max_workers: Maximum number of worker processes.
        timeout: Timeout in seconds for individual gremlin tests.

    Example:
        >>> with WorkerPool(max_workers=4) as pool:
        ...     # Submit work to pool
        ...     pass
    """

    def __init__(
        self,
        max_workers: int | None = None,
        timeout: int = 30,
    ) -> None:
        """Initialize the worker pool.

        Args:
            max_workers: Maximum number of worker processes. Defaults to CPU count.
            timeout: Timeout in seconds for individual tests. Defaults to 30.
        """
        self._max_workers = max_workers if max_workers is not None else (os.cpu_count() or 4)
        self._timeout = timeout
        self._executor: ProcessPoolExecutor | None = None
        self._shutdown_called = False

    @property
    def max_workers(self) -> int:
        """Return the maximum number of workers."""
        return self._max_workers

    @property
    def timeout(self) -> int:
        """Return the timeout in seconds."""
        return self._timeout

    def __enter__(self) -> Self:
        """Enter the context manager, starting the worker pool."""
        self._executor = ProcessPoolExecutor(max_workers=self._max_workers)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        """Exit the context manager, shutting down the pool."""
        self.shutdown(wait=True)

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

        Args:
            wait: If True, wait for pending work to complete. If False, cancel
                  pending work immediately.
        """
        if self._shutdown_called:
            return

        self._shutdown_called = True
        if self._executor is not None:
            self._executor.shutdown(wait=wait, cancel_futures=not wait)
            self._executor = None

    def submit(
        self,
        gremlin_id: str,
        test_command: list[str],
        rootdir: str,
        instrumented_dir: str | None,
        env_vars: dict[str, str],
    ) -> Future[WorkerResult]:
        """Submit a gremlin test for execution.

        Args:
            gremlin_id: The ID of the gremlin to test.
            test_command: Command to run tests.
            rootdir: Root directory for test execution.
            instrumented_dir: Directory with instrumented sources (or None).
            env_vars: Additional environment variables to set.

        Returns:
            Future that will contain the WorkerResult when complete.

        Raises:
            RuntimeError: If the pool is not active (not in context).
        """
        if self._executor is None:
            msg = 'WorkerPool is not active. Use as context manager.'
            raise RuntimeError(msg)

        # Add instrumented dir to env vars if provided
        all_env_vars = dict(env_vars)
        if instrumented_dir is not None:
            all_env_vars['PYTEST_GREMLINS_SOURCES_FILE'] = str(Path(instrumented_dir) / 'sources.json')

        # Suppress subprocess coverage tracking: sitecustomize.py fires at Python startup
        # before --no-cov takes effect, so we explicitly clear the env var here.
        all_env_vars['COVERAGE_PROCESS_START'] = ''

        return self._executor.submit(
            _run_gremlin_test,
            gremlin_id,
            test_command,
            rootdir,
            all_env_vars,
            self._timeout,
        )
