"""Variable assignment stage for flow execution."""

from __future__ import annotations

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

logger = get_logger(__name__)


class VariableAssignmentStage:
    """Assigns step result data to configured variables.

    Runs: Post-execution only
    Modifies: Result variables_updated

    Many executors support `resultVariable`/`responseVariable` fields. This
    stage ensures the executor result contains a corresponding variable update.
    """

    RESULT_KEYS = ("resultVariable", "responseVariable")

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Assign result data to configured variable names.

        Args:
            state: Current execution state

        Returns:
            New state with updated result variables when needed
        """
        # Only run after execution
        if state.is_pre_execution:
            return state

        # Skip if no data to assign
        if state.result.data is None:
            return state

        # Determine configured variable alias
        assigned_key = self._get_result_variable(state)
        if assigned_key is None:
            return state

        logger.debug(
            "Assigning result to variable",
            extra={
                "step_id": state.step.id,
                "variable": assigned_key,
            },
        )

        # Mutate result to add variable assignment
        result = state.result
        if result.variables_updated is None:
            result.variables_updated = {}
        result.variables_updated[assigned_key] = result.data

        return state.with_result(result)

    def _get_result_variable(self, state: ExecutionState) -> str | None:
        """Extract configured result variable name from step config."""
        config = state.step.config or {}
        for key in self.RESULT_KEYS:
            alias = config.get(key)
            if isinstance(alias, str) and alias.strip():
                return alias.strip()
        return None
