"""Target mode stage for test mode 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
    from agent_flows.models.test import TestManifest

logger = get_logger(__name__)


class TargetModeStage:
    """Sets target mode flag for composite executors.

    Runs: Pre-execution only
    Modifies: Context flags

    In test mode, when a composite executor (conditional, switch, loop) is
    itself the target step, this stage sets a flag in the context. The
    executor can check this flag to return its own result without executing
    child blocks.
    """

    COMPOSITE_TYPES = {"conditional", "switch", "loop"}

    def __init__(self, test_manifest: TestManifest | None = None) -> None:
        """Initialize target mode stage.

        Args:
            test_manifest: Optional test manifest for target step configuration
        """
        self.test_manifest = test_manifest

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Set target mode flag if step is composite target.

        Args:
            state: Current execution state

        Returns:
            New state with target mode flag in context if applicable
        """
        # Only run before execution
        if state.is_post_execution:
            return state

        # Skip if no test manifest
        if not self.test_manifest:
            return state

        # Check if this step is the target and is a composite executor
        is_target = self.test_manifest.target == state.step.id
        is_composite = state.step.type in self.COMPOSITE_TYPES

        if is_target and is_composite:
            logger.info(
                "Enabling composite target mode",
                extra={"step_id": state.step.id, "step_type": state.step.type},
            )
            new_context = state.context.with_flag("composite_target_mode", True)
            return state.with_context(new_context)

        return state
