"""
Process-based executor for running CPU-intensive pipeline jobs.

This module provides a way to run document ingestion pipelines in separate
processes to avoid blocking the main server thread. It follows the pattern
established in adri_agents.lib.weasyprint_subprocess.
"""

from __future__ import annotations

import asyncio
import atexit
import functools
import logging
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from threading import Lock
from typing import Any, Mapping, Optional

from idxr.pipelines.document_ingest import (
    DocumentIngestContext,
    DocumentIngestRequest,
    DocumentIngestResult,
)


class PipelineProcessError(RuntimeError):
    """Raised when a pipeline subprocess fails."""


@dataclass(frozen=True)
class PipelineBuilder:
    """
    Configuration for building a DocumentIngestPipeline in subprocess.

    This allows the subprocess to construct the pipeline without tight coupling
    to specific agent implementations. The subprocess will import the specified
    module and call the factory function with the provided kwargs.

    Example:
        PipelineBuilder(
            module_path="adri_agents.app.routers.change_req_research_agent.extensions.pipeline_extension",
            function_name="build_pipeline_extension",
            kwargs={"conversation_id": uuid4()}
        )
    """

    module_path: str  # e.g., "myapp.pipeline_factory"
    function_name: str  # e.g., "create_pipeline"
    kwargs: Mapping[str, Any] = field(
        default_factory=dict
    )  # Arguments to pass to factory

    def __post_init__(self) -> None:
        """Ensure immutability of mutable fields."""
        object.__setattr__(self, "kwargs", dict(self.kwargs))


@dataclass(frozen=True)
class PipelineExecutionRequest:
    """
    Serializable request for running a pipeline in a subprocess.

    Contains all the information needed to execute a document ingestion
    pipeline in an isolated process.
    """

    # Pipeline request parameters
    ingest_request: DocumentIngestRequest
    ingest_context: DocumentIngestContext

    # Pipeline builder configuration
    pipeline_builder: PipelineBuilder

    # Environment configuration (optional overrides)
    env_overrides: Mapping[str, str] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Ensure immutability of mutable fields."""
        object.__setattr__(self, "env_overrides", dict(self.env_overrides))


@dataclass(frozen=True)
class PipelineExecutionResult:
    """Result returned by the pipeline subprocess."""

    ingest_result: DocumentIngestResult
    worker_pid: int  # For debugging/monitoring
    stdout_logs: str  # Captured stdout from subprocess
    stderr_logs: str  # Captured stderr from subprocess


# Module-level executor singleton
_EXECUTOR_LOCK = Lock()
_EXECUTOR: Optional[ProcessPoolExecutor] = None
_ATEXIT_REGISTERED = False


def _resolve_max_workers() -> int:
    """
    Determine the number of worker processes for the executor.

    Reads from ADRI_PIPELINE_MAX_WORKERS environment variable,
    or defaults to min(cpu_count, 4) to avoid oversubscription.
    """
    env_value = os.getenv("ADRI_PIPELINE_MAX_WORKERS")
    if env_value:
        try:
            parsed = int(env_value)
            if parsed > 0:
                return parsed
        except ValueError:
            pass

    cpu_count = os.cpu_count() or 1
    # Cap to avoid oversubscription: up to 4 workers, at least 1.
    return max(1, min(4, cpu_count))


def _create_executor(max_workers: Optional[int] = None) -> ProcessPoolExecutor:
    """
    Create a new ProcessPoolExecutor with the specified number of workers.

    Uses 'spawn' multiprocessing context for better isolation and
    cross-platform compatibility.
    """
    max_workers = max_workers or _resolve_max_workers()
    ctx = multiprocessing.get_context("spawn")
    return ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx)


def get_pipeline_executor() -> ProcessPoolExecutor:
    """
    Return a module-level ProcessPoolExecutor for pipeline tasks.

    The executor is created lazily on first use and reused afterwards.
    This singleton pattern ensures efficient resource usage across
    multiple pipeline executions.

    Automatically registers an atexit handler to ensure proper cleanup
    of multiprocessing resources when the process exits.
    """
    global _EXECUTOR, _ATEXIT_REGISTERED
    if _EXECUTOR is None:
        with _EXECUTOR_LOCK:
            if _EXECUTOR is None:
                _EXECUTOR = _create_executor()

                # Register atexit handler to clean up resources on process exit
                # This prevents semaphore leaks when the server shuts down
                if not _ATEXIT_REGISTERED:
                    atexit.register(shutdown_executor)
                    _ATEXIT_REGISTERED = True

    return _EXECUTOR


def _run_pipeline_in_process(
    request: PipelineExecutionRequest,
    log_queue: Optional[multiprocessing.Queue] = None,
) -> PipelineExecutionResult:
    """
    Worker function that runs in the subprocess.

    This function:
    1. Applies environment overrides
    2. Configures logging to capture output
    3. Builds a fresh pipeline extension
    4. Runs the ingestion pipeline
    5. Returns the result with captured logs

    Args:
        request: The pipeline execution request with all necessary parameters
        log_queue: Optional queue for real-time log streaming to parent process

    Note: This function must be pickle-able, so it cannot reference
    non-serializable objects from the parent process.
    """
    import asyncio
    import io
    import logging
    import os
    import sys

    # Capture stdout and stderr
    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()

    # Save original streams
    original_stdout = sys.stdout
    original_stderr = sys.stderr

    try:
        # Redirect stdout and stderr
        sys.stdout = stdout_capture
        sys.stderr = stderr_capture

        # Configure logging to write to stderr (which is now captured)
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            force=True,  # Reconfigure even if already configured
        )

        logger = logging.getLogger(__name__)

        # Add queue handler for real-time streaming if queue provided
        if log_queue is not None:
            # Capture queue in closure for type checking
            queue_ref = log_queue

            class QueueHandler(logging.Handler):
                """Handler that sends log records to a multiprocessing queue."""

                def emit(self, record: logging.LogRecord) -> None:
                    try:
                        # Send formatted message to queue (non-blocking)
                        queue_ref.put_nowait(self.format(record))
                    except Exception:
                        # Silently ignore queue errors to avoid breaking pipeline
                        pass

            queue_handler = QueueHandler()
            queue_handler.setFormatter(
                logging.Formatter(
                    "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
                )
            )
            logger.addHandler(queue_handler)

        logger.info(
            f"Starting pipeline execution in subprocess {os.getpid()} "
            f"for collection: {request.ingest_request.collection_name}"
        )

        # Apply environment overrides (if any)
        for key, value in request.env_overrides.items():
            os.environ[key] = value
            logger.debug(f"Set env var: {key}={value}")

        # Build pipeline using the provided builder configuration
        # This keeps the subprocess worker agent-agnostic
        logger.info("Building pipeline using factory")
        import importlib

        builder = request.pipeline_builder
        try:
            # Dynamically import the factory module
            module = importlib.import_module(builder.module_path)
            factory_func = getattr(module, builder.function_name)

            # Call the factory to build the pipeline extension
            pipeline_extension = factory_func(**builder.kwargs)
        except Exception as e:
            logger.error(f"Failed to build pipeline: {e}")
            raise RuntimeError(f"Pipeline builder failed: {e}") from e

        # Run the pipeline (need to create event loop for async execution)
        logger.info("Running document ingestion pipeline")
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            ingest_result = loop.run_until_complete(
                pipeline_extension.pipeline.run(
                    request.ingest_request,
                    request.ingest_context,
                )
            )
        finally:
            loop.close()

        logger.info(
            f"Pipeline execution completed: "
            f"{ingest_result.inserted_chunks} chunks inserted"
        )

        # Send completion sentinel to queue
        if log_queue is not None:
            try:
                log_queue.put_nowait(None)
            except Exception:
                pass

        return PipelineExecutionResult(
            ingest_result=ingest_result,
            worker_pid=os.getpid(),
            stdout_logs=stdout_capture.getvalue(),
            stderr_logs=stderr_capture.getvalue(),
        )

    except Exception as e:
        # Log the error before re-raising
        logging.error(f"Pipeline execution failed: {e}", exc_info=True)

        # Send error sentinel to queue
        if log_queue is not None:
            try:
                log_queue.put_nowait(None)
            except Exception:
                pass

        raise

    finally:
        # Restore original streams
        sys.stdout = original_stdout
        sys.stderr = original_stderr


async def _stream_logs_from_queue(
    log_queue: multiprocessing.Queue, logger: logging.Logger, worker_pid: int
) -> None:
    """
    Background task that consumes log messages from subprocess queue.

    Args:
        log_queue: Queue from which to read log messages
        logger: Logger to write messages to
        worker_pid: PID of worker process (for log prefixing)
    """
    while True:
        try:
            # Check for messages with a short timeout
            await asyncio.sleep(0.05)  # 50ms polling interval

            # Try to get message (non-blocking)
            try:
                message = log_queue.get_nowait()
            except Exception:
                # Queue empty, continue polling
                continue

            # None is the sentinel indicating completion
            if message is None:
                break

            # Log the message from subprocess
            logger.info(f"[subprocess {worker_pid}] {message}")

        except Exception as e:
            # Log any unexpected errors but don't crash
            logger.warning(f"Error streaming logs from subprocess: {e}")
            break


async def run_pipeline_in_subprocess(
    ingest_request: DocumentIngestRequest,
    ingest_context: DocumentIngestContext,
    pipeline_builder: PipelineBuilder,
    *,
    env_overrides: Optional[Mapping[str, str]] = None,
    executor: Optional[ProcessPoolExecutor] = None,
    stream_logs: bool = True,
) -> DocumentIngestResult:
    """
    Run a document ingestion pipeline in a separate process with real-time logging.

    This function offloads CPU-intensive pipeline work (downloading,
    text extraction, chunking, embedding) to a subprocess, preventing
    it from blocking the main server thread.

    The pipeline is constructed in the subprocess using the provided builder,
    keeping this executor agent-agnostic and allowing different agents to use
    different pipeline configurations.

    When stream_logs is enabled, log messages from the subprocess are streamed
    to the parent process in real-time via a multiprocessing queue, allowing
    you to monitor progress during long-running ingestions.

    Args:
        ingest_request: The document ingest request parameters.
        ingest_context: The ingestion context (conversation ID, metadata).
        pipeline_builder: Configuration for building the pipeline in subprocess.
        env_overrides: Optional environment variable overrides for the subprocess.
            Caller is responsible for providing necessary env vars for their pipeline.
        executor: Optional custom executor. Defaults to module singleton.
        stream_logs: If True (default), stream logs from subprocess in real-time.
            If False, logs are only available after completion.

    Returns:
        DocumentIngestResult from the pipeline execution.

    Raises:
        PipelineProcessError: If the pipeline execution fails in the subprocess.

    Example:
        >>> from uuid import UUID
        >>> builder = PipelineBuilder(
        ...     module_path="myapp.pipeline_factory",
        ...     function_name="build_pipeline",
        ...     kwargs={"conversation_id": UUID("...")}
        ... )
        >>> request = DocumentIngestRequest(
        ...     source_url="https://example.com/doc.pdf",
        ...     collection_name="my_collection",
        ...     mime_type="application/pdf",
        ... )
        >>> context = DocumentIngestContext(
        ...     conversation_id=UUID("..."),
        ...     default_metadata={"agent": "research"},
        ... )
        >>> result = await run_pipeline_in_subprocess(request, context, builder)
        >>> print(f"Inserted {result.inserted_chunks} chunks")
    """
    # Use provided env overrides or empty dict
    if env_overrides is None:
        env_overrides = {}

    # Build the execution request
    execution_request = PipelineExecutionRequest(
        ingest_request=ingest_request,
        ingest_context=ingest_context,
        pipeline_builder=pipeline_builder,
        env_overrides=env_overrides,
    )

    # Submit to executor
    chosen_executor = executor or get_pipeline_executor()
    loop = asyncio.get_running_loop()

    logger = logging.getLogger(__name__)

    # Create queue for log streaming if enabled
    log_queue: Optional[multiprocessing.Queue] = None
    log_stream_task: Optional[asyncio.Task] = None

    if stream_logs:
        # Use Manager().Queue() which creates a proxy that can be pickled
        # and passed to subprocess via ProcessPoolExecutor
        manager = multiprocessing.Manager()
        log_queue = manager.Queue()  # type: ignore[assignment]

    try:
        # Prepare worker function with queue
        if log_queue is not None:
            worker_fn = functools.partial(_run_pipeline_in_process, log_queue=log_queue)
        else:
            worker_fn = _run_pipeline_in_process  # type: ignore[assignment]

        # Start log streaming task if queue exists
        # We need a temporary PID placeholder; will update from result
        if log_queue is not None:
            log_stream_task = asyncio.create_task(
                _stream_logs_from_queue(log_queue, logger, worker_pid=0)
            )

        # Submit work to subprocess
        result = await loop.run_in_executor(
            chosen_executor,
            worker_fn,
            execution_request,
        )

        # Wait for log streaming to complete (queue will have sentinel)
        if log_stream_task is not None:
            await log_stream_task

        # If not streaming, log the captured output now (fallback)
        if not stream_logs:
            if result.stdout_logs:
                logger.info(
                    f"Pipeline subprocess {result.worker_pid} stdout:\n"
                    f"{result.stdout_logs}"
                )

            if result.stderr_logs:
                logger.info(
                    f"Pipeline subprocess {result.worker_pid} stderr:\n"
                    f"{result.stderr_logs}"
                )

        # Log completion summary
        logger.info(
            f"Pipeline subprocess {result.worker_pid} completed successfully: "
            f"{result.ingest_result.inserted_chunks} chunks inserted into "
            f"collection '{result.ingest_result.collection_name}'"
        )

        return result.ingest_result

    except Exception as exc:
        # Cancel log streaming task if running
        if log_stream_task is not None and not log_stream_task.done():
            log_stream_task.cancel()
            try:
                await log_stream_task
            except asyncio.CancelledError:
                pass

        # Try to get logs from the failed execution if available
        error_msg = f"Pipeline subprocess failed: {exc}"

        # Note: We can't access result.stdout_logs here because the exception
        # occurred during execution. The logs are already printed via logging
        # inside the subprocess before the exception was raised.
        logger.error(error_msg)

        raise PipelineProcessError(error_msg) from exc

    finally:
        # Close the queue if it was created
        if log_queue is not None:
            try:
                log_queue.close()
                log_queue.join_thread()
            except Exception:
                pass


def shutdown_executor() -> None:
    """
    Shutdown the shared executor and cleanup resources.

    This is useful for:
    - Clean application shutdown (called automatically via atexit)
    - Unit tests that need to reset state
    - Resource cleanup in long-running applications

    Safe to call multiple times - subsequent calls are no-ops.
    """
    global _EXECUTOR, _ATEXIT_REGISTERED
    with _EXECUTOR_LOCK:
        if _EXECUTOR is not None:
            _EXECUTOR.shutdown(wait=False, cancel_futures=True)
            _EXECUTOR = None
        # Reset the flag so a new executor can register atexit again
        _ATEXIT_REGISTERED = False
