"""Target propagation stage for nested target detection."""

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 TargetPropagationStage:
    """Propagates nested target results to parent.

    Runs: Post-execution only
    Modifies: Result target fields

    For composite executors, if a nested block contained the target step,
    this stage propagates the target information to the parent result.
    This enables target detection to bubble up through nested structures.
    """

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Propagate nested target information to parent result.

        Args:
            state: Current execution state

        Returns:
            New state with propagated target info if applicable
        """
        # Only run after execution
        if state.is_pre_execution:
            return state

        result = state.result

        # Already marked as target? Skip propagation
        if result.is_target_result:
            return state

        # Check if any nested results contain the target
        if result.nested_results:
            for nested in result.nested_results:
                if nested.is_target_result:
                    # Propagate target info from nested to parent
                    result.is_target_result = True
                    result.target_step_id = nested.target_step_id
                    result.target_step_type = nested.target_step_type

                    logger.info(
                        "Propagated nested target",
                        extra={
                            "parent_step": state.step.id,
                            "target_step": nested.target_step_id,
                        },
                    )

                    # Only create new state when we modified the result
                    return state.with_result(result)

        # No propagation occurred - return unchanged state
        return state
