"""Interpolation stage for variable replacement in step configuration."""

from __future__ import annotations

import copy
from typing import TYPE_CHECKING, Any

from agent_flows.core.runtime.variable_interpolator import VariableInterpolator
from agent_flows.models.flow import FlowStep

if TYPE_CHECKING:
    from agent_flows.core.flow_engine.pipeline import ExecutionState


class InterpolationStage:
    """Interpolates variables in step configuration.

    Runs: Pre-execution only
    Modifies: Step config in context

    For composite executors (switch, conditional, loop), nested blocks are
    preserved to enable deferred interpolation. This ensures that blocks
    see the latest variable state updated by sibling blocks.
    """

    def __init__(self) -> None:
        """Initialize interpolation stage."""
        self.interpolator = VariableInterpolator()

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Interpolate variables in step configuration.

        Args:
            state: Current execution state

        Returns:
            New state with interpolated step config
        """
        # Only run before execution
        if state.is_post_execution:
            return state

        # For composite executors, preserve nested blocks to enable deferred interpolation
        config_to_interpolate, preserved_blocks = self._preserve_nested_blocks(
            state.step.type, state.step.config
        )

        # Interpolate variables in the (potentially modified) step configuration
        interpolated_config = self.interpolator.interpolate_object(
            config_to_interpolate, state.context.variables
        )

        # Restore un-interpolated nested blocks for composite executors
        final_config = self._restore_nested_blocks(
            state.step.type, interpolated_config, preserved_blocks
        )

        # Return state with updated step
        new_step = FlowStep(
            id=state.step.id,
            type=state.step.type,
            config=final_config,
        )
        return state.with_step(new_step)

    def _preserve_nested_blocks(
        self, step_type: str, config: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Preserve nested blocks from composite executors for deferred interpolation.

        For composite executors (switch, conditional, loop), nested block configs should not be
        interpolated when the parent executor's config is interpolated. Instead, blocks
        are interpolated individually when they execute, ensuring they see the latest
        variable state updated by sibling blocks.

        Args:
            step_type: Type of the executor step
            config: Step configuration dict

        Returns:
            Tuple of (config_without_nested_blocks, preserved_blocks)
        """
        if step_type not in ["switch", "conditional", "loop"]:
            return config, {}

        preserved_blocks: dict[str, Any] = {}
        config_copy = copy.deepcopy(config)

        if step_type == "switch":
            if "cases" in config_copy:
                preserved_blocks["cases"] = []
                for case in config_copy["cases"]:
                    preserved_blocks["cases"].append({"blocks": case.get("blocks", [])})
                    case["blocks"] = []

            if "defaultBlocks" in config_copy:
                preserved_blocks["defaultBlocks"] = config_copy["defaultBlocks"]
                config_copy["defaultBlocks"] = []

        elif step_type == "conditional":
            if "truePath" in config_copy:
                preserved_blocks["truePath"] = config_copy["truePath"]
                config_copy["truePath"] = []

            if "falsePath" in config_copy:
                preserved_blocks["falsePath"] = config_copy["falsePath"]
                config_copy["falsePath"] = []

        elif step_type == "loop":
            if "loopBlocks" in config_copy:
                preserved_blocks["loopBlocks"] = config_copy["loopBlocks"]
                config_copy["loopBlocks"] = []

        return config_copy, preserved_blocks

    def _restore_nested_blocks(
        self, step_type: str, interpolated_config: dict[str, Any], preserved_blocks: dict[str, Any]
    ) -> dict[str, Any]:
        """Restore un-interpolated nested blocks to the interpolated config.

        Args:
            step_type: Type of the executor step
            interpolated_config: Config that has been interpolated
            preserved_blocks: Preserved un-interpolated nested blocks

        Returns:
            Config with restored nested blocks
        """
        if not preserved_blocks:
            return interpolated_config

        if step_type == "switch":
            if "cases" in preserved_blocks and "cases" in interpolated_config:
                for i, case in enumerate(interpolated_config["cases"]):
                    if i < len(preserved_blocks["cases"]):
                        case["blocks"] = preserved_blocks["cases"][i]["blocks"]

            if "defaultBlocks" in preserved_blocks:
                interpolated_config["defaultBlocks"] = preserved_blocks["defaultBlocks"]

        elif step_type == "conditional":
            if "truePath" in preserved_blocks:
                interpolated_config["truePath"] = preserved_blocks["truePath"]

            if "falsePath" in preserved_blocks:
                interpolated_config["falsePath"] = preserved_blocks["falsePath"]

        elif step_type == "loop":
            if "loopBlocks" in preserved_blocks:
                interpolated_config["loopBlocks"] = preserved_blocks["loopBlocks"]

        return interpolated_config
