"""Executor stage for running the current flow step."""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

from agent_flows.utils.logging import get_logger

if TYPE_CHECKING:
    from agent_flows.core.flow_engine.pipeline import ExecutionState
    from agent_flows.executors.registry import ExecutorRegistry

logger = get_logger(__name__)


class ExecutorStage:
    """Executes the current step through the registered executor.

    Runs: Execution phase (only when result is not yet populated)
    Modifies: Result

    This stage is the only place where the actual executor logic runs.
    """

    def __init__(self, executor_registry: ExecutorRegistry) -> None:
        """Initialize executor stage.

        Args:
            executor_registry: Registry used to resolve executors by type
        """
        self.executor_registry = executor_registry

    async def process(self, state: ExecutionState):
        """Execute the current step if it has not been executed yet.

        Args:
            state: Current execution state

        Returns:
            State with populated result after executor runs
        """
        # Skip if step already executed
        if state.is_post_execution:
            return state

        executor = self.executor_registry.get_executor(state.step.type)
        logger.debug(
            "Executing step via executor",
            extra={
                "step_id": state.step.id,
                "step_type": state.step.type,
                "executor": executor.__class__.__name__,
            },
        )

        start_time = time.monotonic()
        result = await executor.execute(state.step.config, state.context)

        # Ensure execution_time is populated
        if getattr(result, "execution_time", 0) in (None, 0):
            result.execution_time = time.monotonic() - start_time

        return state.with_result(result)
