"""Target marking 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 TargetMarkingStage:
    """Marks result as target if step was the target.

    Runs: Post-execution only
    Modifies: Result target fields

    In test mode, when a step is the target, this stage marks the result
    with target-related fields for test framework processing.
    """

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

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

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

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Mark result as target if step was the target.

        Args:
            state: Current execution state

        Returns:
            New state with marked result if applicable
        """
        # Only run after execution
        if state.is_pre_execution:
            return state

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

        # Check if this was the target
        if self.test_manifest.target != state.step.id:
            return state

        # Mark as target
        result = state.result
        result.is_target_result = True
        result.target_step_id = state.step.id
        result.target_step_type = state.step.type

        # Mark as composite target if applicable
        if state.step.type in self.COMPOSITE_TYPES:
            result.is_composite_target = True

        logger.info("Marked target result", extra={"step_id": state.step.id})

        return state.with_result(result)
