"""Pin check stage for test executions."""

from __future__ import annotations

import copy
from typing import TYPE_CHECKING

from agent_flows.core.runtime.variable_interpolator import VariableInterpolator
from agent_flows.utils.dict_utils import deep_update
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 PinCheckStage:
    """Checks for pinned results and applies overrides or short-circuits.

    Runs: Pre-execution only
    Can short-circuit: Yes
    Modifies: Step config (when applying overrides), result (when returning pin)

    Pins allow tests to override step configuration or provide canned outputs.
    """

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

        Args:
            test_manifest: Optional test manifest containing pin definitions
        """
        self.test_manifest = test_manifest
        self._interpolator = VariableInterpolator()

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Apply pin overrides or short-circuit with pinned result.

        Args:
            state: Current execution state

        Returns:
            New state reflecting pin handling
        """
        # Only run before execution
        if state.is_post_execution:
            return state

        # Skip if test manifest not provided or no pin for this step
        if not self.test_manifest or not state.step.id:
            return state

        pin = self.test_manifest.get_pin(state.step.id)
        if pin is None:
            return state

        logger.info(
            "Processing test pin",
            extra={"step_id": state.step.id, "has_override": bool(pin.override_config)},
        )

        step_config = copy.deepcopy(state.step.config or {})
        updated_step = state.step

        # Apply configuration overrides if provided
        if pin.override_config:
            merged_config = deep_update(step_config, pin.override_config)
            updated_step = state.step.model_copy(update={"config": merged_config})
            step_config = merged_config
            logger.info(
                "Applied pin configuration override",
                extra={"step_id": state.step.id},
            )

        # If pin provides variables, short-circuit with canned result
        if pin.variables_updated is not None:
            interpolated_config = self._interpolator.interpolate_object(
                step_config, state.context.variables
            )
            result = pin.to_executor_result(interpolated_config)

            logger.info(
                "Using pinned result",
                extra={
                    "step_id": state.step.id,
                    "variables_updated": len(pin.variables_updated),
                },
            )

            return state.stop_with_result(result)

        # Only override config? Continue with updated step
        if updated_step is not state.step:
            return state.with_step(updated_step)

        return state
