"""Execution pipeline infrastructure for Agent Flows.

This module provides the core pipeline components for executing flow steps:
- ExecutionState: Immutable state passed through pipeline stages
- Stage: Protocol defining the interface for all pipeline stages
- ExecutionPipeline: Orchestrator that runs stages in sequence

The pipeline architecture enables clean separation of concerns where each stage
performs a single, well-defined transformation on the execution state.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
from typing import TYPE_CHECKING, Protocol

from agent_flows.exceptions import ExecutorError
from agent_flows.utils.logging import get_logger

if TYPE_CHECKING:
    from agent_flows.models.execution import ExecutionContext, ExecutorResult
    from agent_flows.models.flow import FlowStep

logger = get_logger(__name__)


@dataclass(frozen=True)
class ExecutionState:
    """Immutable state passed through pipeline stages.

    Each stage receives state and returns new state.
    This makes the data flow explicit and debuggable.

    State progresses through two phases:
    1. Pre-execution: result is None
    2. Post-execution: result is populated
    """

    step: FlowStep
    """The step being executed."""

    context: ExecutionContext
    """Current execution context with variables, credentials, etc."""

    result: ExecutorResult | None = None
    """Execution result (None before execution, set after)."""

    should_continue: bool = True
    """If False, pipeline stops (short-circuit)."""

    # ========================================
    # Helper methods for creating new states
    # ========================================

    def with_context(self, context: ExecutionContext) -> ExecutionState:
        """Create new state with updated context."""
        return replace(self, context=context)

    def with_result(self, result: ExecutorResult) -> ExecutionState:
        """Create new state with result."""
        return replace(self, result=result)

    def with_step(self, step: FlowStep) -> ExecutionState:
        """Create new state with updated step."""
        return replace(self, step=step)

    def stop_with_result(self, result: ExecutorResult) -> ExecutionState:
        """Stop pipeline with final result (short-circuit)."""
        return replace(self, result=result, should_continue=False)

    @property
    def is_pre_execution(self) -> bool:
        """True if step hasn't been executed yet."""
        return self.result is None

    @property
    def is_post_execution(self) -> bool:
        """True if step has been executed."""
        return self.result is not None


class Stage(Protocol):
    """A pipeline stage.

    Stages receive execution state and return new state.
    Each stage is independent and self-contained.

    Stages can:
    - Modify context (pre-execution)
    - Execute the step (execution)
    - Modify result (post-execution)
    - Short-circuit with result
    - Pass through unchanged

    The stage itself decides when to act based on state.is_pre_execution
    and state.is_post_execution properties.
    """

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Process execution state.

        Args:
            state: Current execution state

        Returns:
            New execution state (may be modified or unchanged)
        """
        ...


class ExecutionPipeline:
    """Linear pipeline that executes stages in order.

    The pipeline is just a loop that:
    1. Passes state through each stage
    2. Stops if a stage short-circuits
    3. Returns final result

    No complex logic, no hidden behavior, just linear processing.
    """

    def __init__(self, stages: list[Stage]):
        """Initialize pipeline with stages.

        Args:
            stages: List of stages to execute in order
        """
        self.stages = stages

    async def execute(
        self,
        step: FlowStep,
        context: ExecutionContext,
    ) -> ExecutorResult:
        """Execute pipeline for a step.

        Args:
            step: Flow step to execute
            context: Execution context

        Returns:
            Final execution result

        Raises:
            ExecutorError: If pipeline completes without producing result
        """
        # Initialize state
        state = ExecutionState(step=step, context=context)

        # Process through each stage (simple loop!)
        for stage in self.stages:
            # Stage processes and returns new state
            state = await stage.process(state)

            # Check for short-circuit
            if not state.should_continue:
                logger.info(
                    "Pipeline short-circuited",
                    extra={
                        "step_id": step.id,
                        "stage": stage.__class__.__name__,
                    },
                )
                break

            logger.debug(
                "Stage completed",
                extra={
                    "step_id": step.id,
                    "stage": stage.__class__.__name__,
                    "phase": "post-execution" if state.is_post_execution else "pre-execution",
                },
            )

        # Return final result
        if state.result is None:
            raise ExecutorError(f"Pipeline completed without producing result for step: {step.id}")

        return state.result
