"""
Pipeline Executor - The Async DAG Runner.

Stateless execution engine that persists state via checkpoints.
Designed for crash recovery and distributed execution.
"""

from __future__ import annotations

import asyncio
import importlib
import logging
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Protocol, runtime_checkable

from pydantic import BaseModel, Field

from aperion_flow.definitions.pipeline import (
    LLMAssistanceMode,
    Pipeline,
    PipelineStage,
    PipelineStatus,
    StageStatus,
)
from aperion_flow.engine.context import (
    Checkpoint,
    CheckpointStore,
    CheckpointType,
    InMemoryCheckpointStore,
    PipelineContext,
)
from aperion_flow.engine.recovery import (
    RecoveryConfig,
    RecoveryEngine,
    RecoveryResult,
)

logger = logging.getLogger(__name__)


@runtime_checkable
class EventEmitter(Protocol):
    """Protocol for event emission to The Nervous System."""

    async def emit(
        self,
        event_type: str,
        payload: dict[str, Any],
        *,
        source: str | None = None,
        correlation_id: str | None = None,
    ) -> str:
        """Emit an event."""
        ...


@runtime_checkable
class LLMAssistant(Protocol):
    """Protocol for LLM assistance integration with The Switchboard."""

    async def pre_validate(
        self,
        stage: PipelineStage,
        context: PipelineContext,
        timeout: float,
    ) -> dict[str, Any]:
        """Validate stage inputs before execution."""
        ...

    async def analyze_output(
        self,
        stage: PipelineStage,
        output: dict[str, Any],
        context: PipelineContext,
        timeout: float,
    ) -> dict[str, Any]:
        """Analyze and interpret stage output."""
        ...

    async def diagnose_error(
        self,
        stage: PipelineStage,
        error: str,
        context: PipelineContext,
        timeout: float,
    ) -> dict[str, Any]:
        """Diagnose a stage failure and suggest fixes."""
        ...


class StageExecution(BaseModel):
    """Tracks the execution of a single stage."""

    stage_id: str
    stage_name: str
    status: StageStatus = StageStatus.PENDING
    started_at: datetime | None = None
    completed_at: datetime | None = None
    output: dict[str, Any] = Field(default_factory=dict)
    error: str | None = None
    llm_analysis: dict[str, Any] | None = None


class ExecutionResult(BaseModel):
    """Result of a pipeline execution."""

    execution_id: str
    pipeline_id: str
    pipeline_name: str
    status: PipelineStatus
    data: dict[str, Any] = Field(default_factory=dict)

    # Stages
    stage_count: int = 0
    stages_completed: int = 0
    stages_failed: int = 0
    stages_skipped: int = 0

    # Timing
    started_at: datetime | None = None
    completed_at: datetime | None = None

    # Errors
    error: str | None = None
    error_stage: str | None = None

    # Recovery
    recovery_attempts: int = 0
    resumed_from_checkpoint: bool = False

    @property
    def duration_seconds(self) -> float | None:
        """Calculate execution duration."""
        if self.started_at and self.completed_at:
            return (self.completed_at - self.started_at).total_seconds()
        return None


class PipelineExecutor:
    """
    Async DAG pipeline executor.

    Features:
    - Parallel stage execution respecting dependencies
    - Checkpoint-based crash recovery
    - LLM assistance integration
    - Event emission for progress tracking
    - Circuit breaker protection
    """

    def __init__(
        self,
        checkpoint_store: CheckpointStore | None = None,
        event_emitter: EventEmitter | None = None,
        llm_assistant: LLMAssistant | None = None,
        recovery_engine: RecoveryEngine | None = None,
        handler_registry: dict[str, Callable[..., Awaitable[dict[str, Any]]]] | None = None,
    ):
        """
        Initialize the executor.

        Args:
            checkpoint_store: Storage for execution checkpoints
            event_emitter: Event emitter for progress notifications
            llm_assistant: LLM assistant for intelligent stage support
            recovery_engine: Recovery engine for retry/backoff
            handler_registry: Pre-registered handlers (alternative to dynamic loading)
        """
        self._checkpoint_store = checkpoint_store or InMemoryCheckpointStore()
        self._event_emitter = event_emitter
        self._llm_assistant = llm_assistant
        self._recovery_engine = recovery_engine or RecoveryEngine()
        self._handler_registry = handler_registry or {}

    async def execute(
        self,
        pipeline: Pipeline,
        initial_data: dict[str, Any] | None = None,
        resume_from: str | None = None,
        correlation_id: str | None = None,
    ) -> ExecutionResult:
        """
        Execute a pipeline.

        Args:
            pipeline: The pipeline to execute
            initial_data: Initial data for the context
            resume_from: Execution ID to resume from (for crash recovery)
            correlation_id: Correlation ID for tracing

        Returns:
            ExecutionResult with final status and data
        """
        # Initialize or resume context
        if resume_from:
            context = await self._resume_context(pipeline, resume_from)
            resumed = True
        else:
            context = self._create_context(pipeline, initial_data, correlation_id)
            resumed = False

        # Emit start event
        await self._emit_event("flow.started", {
            "pipeline_id": pipeline.id,
            "pipeline_name": pipeline.name,
            "execution_id": context.execution_id,
            "resumed": resumed,
        }, context)

        try:
            # Run the pipeline
            await self._run_pipeline(pipeline, context)

            # Final checkpoint
            await self._save_checkpoint(
                context,
                CheckpointType.PIPELINE_COMPLETE if context.status == PipelineStatus.COMPLETED
                else CheckpointType.PIPELINE_FAILED,
            )

        except asyncio.CancelledError:
            context.mark_cancelled()
            await self._emit_event("flow.cancelled", {
                "execution_id": context.execution_id,
            }, context)
            raise

        except Exception as e:
            context.mark_failed(str(e), type(e).__name__)
            await self._emit_event("flow.failed", {
                "execution_id": context.execution_id,
                "error": str(e),
                "error_type": type(e).__name__,
            }, context)
            await self._save_checkpoint(context, CheckpointType.PIPELINE_FAILED)

        # Build result
        return self._build_result(context, pipeline, resumed)

    async def _run_pipeline(self, pipeline: Pipeline, context: PipelineContext) -> None:
        """Run all stages in the pipeline."""
        context.mark_started()
        await self._save_checkpoint(context, CheckpointType.PIPELINE_START)

        # Get execution order
        stages = pipeline.topological_order()
        completed_ids = set(context.completed_stage_ids)
        pending_stages = [s for s in stages if s.id not in completed_ids]

        # Track running stages
        running: set[asyncio.Task[None]] = set()
        stage_tasks: dict[str, asyncio.Task[None]] = {}

        # Process stages respecting dependencies and parallelism
        stages_by_id = {s.id: s for s in stages}
        ready_queue: list[PipelineStage] = []

        def update_ready_queue() -> None:
            """Update the queue of stages ready to run."""
            for stage in pending_stages:
                if stage.id in stage_tasks:
                    continue  # Already running
                if stage.id in completed_ids:
                    continue  # Already done

                # Check dependencies
                deps_satisfied = all(
                    dep_id in completed_ids
                    for dep_id in stage.depends_on
                )

                if deps_satisfied and stage not in ready_queue:
                    ready_queue.append(stage)

        update_ready_queue()

        while ready_queue or running:
            # Start new stages up to max_parallel
            while ready_queue and len(running) < pipeline.max_parallel:
                stage = ready_queue.pop(0)

                # Check if should skip due to failed dependencies
                if stage.skip_on_failure:
                    failed_deps = [
                        dep_id for dep_id in stage.depends_on
                        if dep_id in context.failed_stage_ids
                    ]
                    if failed_deps:
                        context.record_stage_skipped(
                            stage.id,
                            stage.name,
                            f"Skipped due to failed dependencies: {failed_deps}"
                        )
                        completed_ids.add(stage.id)
                        continue

                task = asyncio.create_task(
                    self._run_stage(stage, pipeline, context)
                )
                running.add(task)
                stage_tasks[stage.id] = task

            if not running:
                break

            # Wait for at least one stage to complete
            done, running = await asyncio.wait(running, return_when=asyncio.FIRST_COMPLETED)

            # Process completed stages
            for task in done:
                # Find which stage completed
                for stage_id, t in list(stage_tasks.items()):
                    if t is task:
                        del stage_tasks[stage_id]
                        completed_ids.add(stage_id)

                        # Check for failure with fail_fast
                        if (pipeline.fail_fast and
                            stage_id in context.failed_stage_ids):
                            # Cancel remaining tasks
                            for remaining_task in running:
                                remaining_task.cancel()
                            # Get actual error from stage result
                            stage_result = context.stage_results.get(stage_id)
                            stage_error = stage_result.error if stage_result else "Unknown error"
                            context.mark_failed(
                                stage_error,
                                "StageFailure",
                                stage_id,
                            )
                            return

            update_ready_queue()

        # Check final status
        if context.failed_stage_ids:
            # Collect error messages from failed stages
            errors = []
            for stage_id in context.failed_stage_ids:
                stage_result = context.stage_results.get(stage_id)
                if stage_result and stage_result.error:
                    errors.append(stage_result.error)
            error_msg = "; ".join(errors) if errors else f"{len(context.failed_stage_ids)} stages failed"
            context.mark_failed(error_msg, "PipelineFailure")
        else:
            context.mark_completed()

    async def _run_stage(
        self,
        stage: PipelineStage,
        pipeline: Pipeline,
        context: PipelineContext,
    ) -> None:
        """Run a single stage with recovery and LLM assistance."""
        context.record_stage_start(stage.id, stage.name)

        await self._emit_event("flow.stage.started", {
            "execution_id": context.execution_id,
            "stage_id": stage.id,
            "stage_name": stage.name,
        }, context)

        await self._save_checkpoint(context, CheckpointType.STAGE_START)

        try:
            # Pre-validate with LLM if configured
            if stage.has_llm_mode(LLMAssistanceMode.PRE_VALIDATE):
                await self._llm_pre_validate(stage, context)

            # Get handler
            handler = await self._get_handler(stage.handler)

            # Execute with recovery
            async def execute_handler() -> dict[str, Any]:
                return await asyncio.wait_for(
                    handler(context.data),
                    timeout=stage.timeout,
                )

            if stage.retry_count > 0:
                config = RecoveryConfig(max_retries=stage.retry_count)
                engine = RecoveryEngine(config)
                result = await engine.execute_with_recovery(
                    execute_handler,
                    circuit_key=stage.id,
                )
                if not result.success:
                    raise RuntimeError(result.final_error or "Stage execution failed")
                output = result.data
            else:
                output = await execute_handler()

            # Merge output to context
            context.update(output)

            # Analyze output with LLM if configured
            llm_analysis = None
            if stage.has_llm_mode(LLMAssistanceMode.ANALYZE_OUTPUT):
                llm_analysis = await self._llm_analyze_output(stage, output, context)

            context.record_stage_complete(stage.id, output, llm_analysis)

            await self._emit_event("flow.stage.completed", {
                "execution_id": context.execution_id,
                "stage_id": stage.id,
                "stage_name": stage.name,
            }, context)

            await self._save_checkpoint(context, CheckpointType.STAGE_COMPLETE)

        except Exception as e:
            error_msg = str(e)
            error_type = type(e).__name__

            # Diagnose with LLM if configured
            if stage.has_llm_mode(LLMAssistanceMode.ERROR_DIAGNOSIS):
                diagnosis = await self._llm_diagnose_error(stage, error_msg, context)
                if diagnosis:
                    error_msg = f"{error_msg}\nDiagnosis: {diagnosis.get('diagnosis', '')}"

            context.record_stage_failed(stage.id, error_msg, error_type)

            await self._emit_event("flow.stage.failed", {
                "execution_id": context.execution_id,
                "stage_id": stage.id,
                "stage_name": stage.name,
                "error": error_msg,
                "error_type": error_type,
            }, context)

            await self._save_checkpoint(context, CheckpointType.STAGE_FAILED)

    async def _get_handler(
        self,
        handler_name: str,
    ) -> Callable[..., Awaitable[dict[str, Any]]]:
        """Get a handler function by name."""
        # Check registry first
        if handler_name in self._handler_registry:
            return self._handler_registry[handler_name]

        # Dynamic import
        try:
            module_path, func_name = handler_name.rsplit(".", 1)
            module = importlib.import_module(module_path)
            handler = getattr(module, func_name)
            return handler
        except (ValueError, ImportError, AttributeError) as e:
            raise RuntimeError(f"Cannot load handler '{handler_name}': {e}") from e

    async def _llm_pre_validate(
        self,
        stage: PipelineStage,
        context: PipelineContext,
    ) -> None:
        """Run LLM pre-validation."""
        if not self._llm_assistant:
            return

        config = stage.get_llm_config(LLMAssistanceMode.PRE_VALIDATE)
        if not config:
            return

        try:
            result = await self._llm_assistant.pre_validate(
                stage, context, config.timeout
            )
            if not result.get("valid", True):
                raise ValueError(
                    f"Pre-validation failed: {result.get('reason', 'Unknown')}"
                )
        except Exception as e:
            if not config.fallback_on_error:
                raise
            logger.warning(f"LLM pre-validation failed, continuing: {e}")

    async def _llm_analyze_output(
        self,
        stage: PipelineStage,
        output: dict[str, Any],
        context: PipelineContext,
    ) -> dict[str, Any] | None:
        """Run LLM output analysis."""
        if not self._llm_assistant:
            return None

        config = stage.get_llm_config(LLMAssistanceMode.ANALYZE_OUTPUT)
        if not config:
            return None

        try:
            return await self._llm_assistant.analyze_output(
                stage, output, context, config.timeout
            )
        except Exception as e:
            logger.warning(f"LLM output analysis failed: {e}")
            return None

    async def _llm_diagnose_error(
        self,
        stage: PipelineStage,
        error: str,
        context: PipelineContext,
    ) -> dict[str, Any] | None:
        """Run LLM error diagnosis."""
        if not self._llm_assistant:
            return None

        config = stage.get_llm_config(LLMAssistanceMode.ERROR_DIAGNOSIS)
        if not config:
            return None

        try:
            return await self._llm_assistant.diagnose_error(
                stage, error, context, config.timeout
            )
        except Exception as e:
            logger.warning(f"LLM error diagnosis failed: {e}")
            return None

    def _create_context(
        self,
        pipeline: Pipeline,
        initial_data: dict[str, Any] | None,
        correlation_id: str | None,
    ) -> PipelineContext:
        """Create a new pipeline context."""
        return PipelineContext(
            pipeline_id=pipeline.id,
            pipeline_name=pipeline.name,
            data=initial_data or {},
            correlation_id=correlation_id,
            metadata=pipeline.metadata.copy(),
        )

    async def _resume_context(
        self,
        pipeline: Pipeline,
        execution_id: str,
    ) -> PipelineContext:
        """Resume a context from checkpoint."""
        checkpoint = await self._checkpoint_store.load(execution_id)
        if not checkpoint:
            raise ValueError(f"No checkpoint found for execution: {execution_id}")

        context = PipelineContext(
            pipeline_id=pipeline.id,
            pipeline_name=pipeline.name,
            execution_id=execution_id,
            correlation_id=checkpoint.metadata.get("correlation_id"),
        )
        context.restore_from_checkpoint(checkpoint)

        logger.info(
            f"Resumed execution {execution_id} from checkpoint, "
            f"{len(checkpoint.completed_stages)} stages already completed"
        )

        return context

    async def _save_checkpoint(
        self,
        context: PipelineContext,
        checkpoint_type: CheckpointType,
    ) -> None:
        """Save a checkpoint."""
        checkpoint = context.create_checkpoint(checkpoint_type)
        await self._checkpoint_store.save(checkpoint)

    async def _emit_event(
        self,
        event_type: str,
        payload: dict[str, Any],
        context: PipelineContext,
    ) -> None:
        """Emit an event if emitter is configured."""
        if not self._event_emitter:
            return

        try:
            await self._event_emitter.emit(
                event_type,
                payload,
                source="aperion-flow",
                correlation_id=context.correlation_id,
            )
        except Exception as e:
            logger.warning(f"Failed to emit event {event_type}: {e}")

    def _build_result(
        self,
        context: PipelineContext,
        pipeline: Pipeline,
        resumed: bool,
    ) -> ExecutionResult:
        """Build the execution result from context."""
        return ExecutionResult(
            execution_id=context.execution_id,
            pipeline_id=pipeline.id,
            pipeline_name=pipeline.name,
            status=context.status,
            data=context.data,
            stage_count=len(pipeline.stages),
            stages_completed=len(context.completed_stage_ids),
            stages_failed=len(context.failed_stage_ids),
            stages_skipped=len([
                r for r in context.stage_results.values()
                if r.status == StageStatus.SKIPPED
            ]),
            started_at=context.started_at,
            completed_at=context.completed_at,
            error=context.error,
            error_stage=context.error_stage,
            resumed_from_checkpoint=resumed,
        )

    def register_handler(
        self,
        name: str,
        handler: Callable[..., Awaitable[dict[str, Any]]],
    ) -> None:
        """Register a handler function."""
        self._handler_registry[name] = handler
