hive-router 0.2.0

use crate::query_planner::ast::merge_path::{Condition, FieldPathSegment, MergePath, Segment};
use crate::query_planner::ast::selection_item::SelectionItem;
use crate::query_planner::ast::selection_set::{
    FieldSelection, InlineFragmentSelection, SelectionSet,
};
use crate::query_planner::ast::type_aware_selection::TypeAwareSelection;
use crate::query_planner::graph::edge::{
    Edge, FieldMove, InterfaceObjectTypeMove, PlannerOverrideContext, ReentryMove,
};
use crate::query_planner::graph::node::Node;
use crate::query_planner::graph::Graph;
use crate::query_planner::planner::fetch::fetch_step_data::{
    FetchStepData, FetchStepFlags, FetchStepKind,
};
use crate::query_planner::planner::fetch::selections::FetchStepSelections;
use crate::query_planner::planner::fetch::state::{MultiTypeFetchStep, SingleTypeFetchStep};
use crate::query_planner::planner::plan_nodes::{FetchNodePathSegment, FetchRewrite, ValueSetter};
use crate::query_planner::planner::tree::query_tree::QueryTree;
use crate::query_planner::planner::tree::query_tree_node::{MutationFieldPosition, QueryTreeNode};
use crate::query_planner::planner::walker::path::OperationPath;
use crate::query_planner::planner::walker::pathfinder::can_satisfy_edge;
use crate::query_planner::planner::QueryPlannerOptions;
use crate::query_planner::state::supergraph_state::{OperationKind, SubgraphName, SupergraphState};
use crate::query_planner::utils::cancellation::CancellationToken;
use petgraph::algo::has_path_connecting;
use petgraph::graph::EdgeReference;
use petgraph::stable_graph::{EdgeIndex, NodeIndex, NodeIndices, NodeReferences, StableDiGraph};
use petgraph::visit::EdgeRef;
use petgraph::visit::{Bfs, IntoNodeReferences};
use petgraph::Directed;
use petgraph::Direction;
use std::collections::{BTreeSet, VecDeque};
use std::fmt::{Debug, Display};
use tracing::{instrument, trace};

use super::error::FetchGraphError;

#[derive(Debug, Clone)]
pub struct FetchGraph<State> {
    pub(crate) graph: StableDiGraph<FetchStepData<State>, ()>,
    pub root_index: Option<NodeIndex>,
    pub(crate) operation_kind: OperationKind,
}

impl FetchGraph<SingleTypeFetchStep> {
    pub fn to_multi_type(self) -> FetchGraph<MultiTypeFetchStep> {
        let new_graph = self
            .graph
            .map(|_, w| w.clone().into_multi_type(), |_, _| ());

        FetchGraph {
            graph: new_graph,
            root_index: self.root_index,
            operation_kind: self.operation_kind,
        }
    }

    pub fn new(kind: OperationKind) -> Self {
        Self {
            graph: StableDiGraph::new(),
            root_index: None,
            operation_kind: kind,
        }
    }
}

impl<State> FetchGraph<State> {
    pub fn all_nodes(&self) -> NodeReferences<'_, FetchStepData<State>> {
        self.graph.node_references()
    }

    pub fn create_fetch_id(&self) -> i64 {
        (self.graph.node_count() + 1) as i64
    }

    pub fn get_step_data_mut(
        &mut self,
        index: NodeIndex,
    ) -> Result<&mut FetchStepData<State>, FetchGraphError> {
        self.graph
            .node_weight_mut(index)
            .ok_or(FetchGraphError::MissingStep(
                index.index(),
                String::from("when getting mutable step data"),
            ))
    }
}

impl<State> FetchGraph<State> {
    pub fn parents_of(&self, index: NodeIndex) -> petgraph::stable_graph::Edges<'_, (), Directed> {
        self.graph.edges_directed(index, Direction::Incoming)
    }

    pub fn children_of(&self, index: NodeIndex) -> petgraph::stable_graph::Edges<'_, (), Directed> {
        self.graph.edges_directed(index, Direction::Outgoing)
    }

    pub fn is_descendant_of(&self, descendant: NodeIndex, ancestor: NodeIndex) -> bool {
        has_path_connecting(&self.graph, ancestor, descendant, None)
    }

    /// Checks if one is ancestor of the other and vice versa
    pub fn is_ancestor_or_descendant(&self, a: NodeIndex, b: NodeIndex) -> bool {
        self.is_descendant_of(a, b) || self.is_descendant_of(b, a)
    }

    pub fn step_indices<'a>(&'a self) -> NodeIndices<'a, FetchStepData<State>> {
        self.graph.node_indices()
    }

    pub fn get_step_data(
        &self,
        index: NodeIndex,
    ) -> Result<&FetchStepData<State>, FetchGraphError> {
        self.graph
            .node_weight(index)
            .ok_or(FetchGraphError::MissingStep(
                index.index(),
                String::from("when getting step data"),
            ))
    }

    pub fn connect(&mut self, parent_index: NodeIndex, child_index: NodeIndex) -> EdgeIndex {
        self.graph.update_edge(parent_index, child_index, ())
    }

    pub fn remove_edge(&mut self, edge_index: EdgeIndex) -> bool {
        self.graph.remove_edge(edge_index).is_some_and(|_| true)
    }

    #[instrument(level = "trace", skip_all, fields(
      index = index.index(),
    ))]
    pub fn remove_step(&mut self, index: NodeIndex) -> bool {
        self.graph.remove_node(index).is_some_and(|_| true)
    }

    pub fn add_step(&mut self, data: FetchStepData<State>) -> NodeIndex {
        self.graph.add_node(data)
    }

    pub fn bfs<F>(&self, root_index: NodeIndex, mut visitor: F) -> Option<NodeIndex>
    where
        F: FnMut(&NodeIndex, &FetchStepData<State>) -> bool,
    {
        self.graph.node_weight(root_index)?;

        let mut bfs = Bfs::new(&self.graph, root_index);

        while let Some(step_index) = bfs.next(&self.graph) {
            // Get the data for the current step. bfs.next() should yield valid indices.
            let step_data = self
                .graph
                .node_weight(step_index)
                .expect("BFS returned invalid node index");

            if visitor(&step_index, step_data) {
                return Some(step_index);
            }
        }

        None
    }

    /// Checks whether the given step is an ancestor of a step that has the given condition
    pub fn is_ancestor_of_condition(&self, step_index: NodeIndex, condition: &Condition) -> bool {
        let mut bfs = Bfs::new(&self.graph, step_index);

        while let Some(current_index) = bfs.next(&self.graph) {
            let current_step = match self.get_step_data(current_index) {
                Ok(s) => s,
                Err(_) => continue,
            };

            if let Some(step_condition) = &current_step.condition {
                if step_condition == condition {
                    return true;
                }
            }
        }

        false
    }
}

impl FetchGraph<MultiTypeFetchStep> {
    #[instrument(level = "trace", skip_all)]
    pub fn collect_variable_usages(&mut self) -> Result<(), FetchGraphError> {
        let nodes_idx = self.graph.node_indices().collect::<Vec<_>>();

        for node_idx in nodes_idx {
            let step_data = self.get_step_data_mut(node_idx)?;
            let usage = step_data.output.variable_usages();

            if !usage.is_empty() {
                step_data.variable_usages = Some(usage);
            }
        }

        Ok(())
    }

    pub fn get_pair_of_steps_mut(
        &mut self,
        index1: NodeIndex,
        index2: NodeIndex,
    ) -> Result<
        (
            &mut FetchStepData<MultiTypeFetchStep>,
            &mut FetchStepData<MultiTypeFetchStep>,
        ),
        FetchGraphError,
    > {
        // `index_twice_mut` panics when indexes are equal
        if index1 == index2 {
            return Err(FetchGraphError::SameNodeIndex(index1.index()));
        }

        // `index_twice_mut` panics when nodes do not exist
        if self.graph.node_weight(index1).is_none() {
            return Err(FetchGraphError::MissingStep(
                index1.index(),
                String::from("when checking existence"),
            ));
        }
        if self.graph.node_weight(index2).is_none() {
            return Err(FetchGraphError::MissingStep(
                index2.index(),
                String::from("when checking existence"),
            ));
        }

        Ok(self.graph.index_twice_mut(index1, index2))
    }
}

impl<State> Display for FetchGraph<State> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Nodes:")?;
        for node_index in self.graph.node_indices() {
            // ignore root node
            if node_index.index() == 0 {
                continue;
            }
            if let Some(fetch_step) = self.graph.node_weight(node_index) {
                fetch_step.pretty_write(f, node_index)?;
                writeln!(f)?;
            }
        }

        writeln!(f, "\nTree:")?;

        let mut stack: VecDeque<(NodeIndex, usize)> = VecDeque::new();

        let roots = find_graph_roots(self);

        if roots.is_empty() {
            writeln!(f, "Fetch step graph is empty or has no roots.")?;
            return Ok(());
        }

        if roots.len() > 1 {
            writeln!(f, "Fetch step graph has multiple roots:")?;
            return Ok(());
        }

        for root_index in roots {
            for child_index in self
                .graph
                .edges_directed(root_index, Direction::Outgoing)
                .map(|edge_ref| edge_ref.target())
            {
                stack.push_front((child_index, 0));
            }
        }

        while let Some((node_index, depth)) = stack.pop_back() {
            let indent = "  ".repeat(depth);
            writeln!(f, "{indent}[{}]", node_index.index())?;

            for child_index in self
                .graph
                .edges_directed(node_index, Direction::Outgoing)
                .map(|edge_ref| edge_ref.target())
            {
                stack.push_back((child_index, depth + 1));
            }
        }

        Ok(())
    }
}

fn create_noop_fetch_step(
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    created_from_requires: bool,
) -> NodeIndex {
    let flags = if created_from_requires {
        FetchStepFlags::USED_FOR_REQUIRES
    } else {
        FetchStepFlags::empty()
    };

    fetch_graph.add_step(FetchStepData {
        id: fetch_graph.create_fetch_id(),
        service_name: SubgraphName::any(),
        response_path: MergePath::default(),
        input: FetchStepSelections::new_empty(),
        output: FetchStepSelections::new_empty(),
        flags,
        condition: None,
        kind: FetchStepKind::Root,
        operation_kind: OperationKind::Query,
        input_rewrites: None,
        output_rewrites: None,
        variable_usages: None,
        variable_definitions: None,
        mutation_field_position: None,
        internal_aliases_locations: Vec::new(),
    })
}

fn create_fetch_step_for_entity_call(
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    subgraph_name: &SubgraphName,
    input_type_name: &str,
    output_type_name: &str,
    response_path: &MergePath,
    condition: Option<&Condition>,
    used_for_requires: bool,
) -> NodeIndex {
    let flags = if used_for_requires {
        FetchStepFlags::USED_FOR_REQUIRES
    } else {
        FetchStepFlags::empty()
    };
    let mut input = FetchStepSelections::new(input_type_name);
    input
        .add(&SelectionSet {
            items: vec![SelectionItem::Field(FieldSelection::new_typename())],
        })
        .unwrap();

    fetch_graph.add_step(FetchStepData {
        id: fetch_graph.create_fetch_id(),
        service_name: subgraph_name.clone(),
        response_path: response_path.clone(),
        input,
        output: FetchStepSelections::new(output_type_name),
        flags,
        condition: condition.cloned(),
        kind: FetchStepKind::Entity,
        operation_kind: OperationKind::Query,
        input_rewrites: None,
        output_rewrites: None,
        variable_usages: None,
        variable_definitions: None,
        mutation_field_position: None,
        internal_aliases_locations: Vec::new(),
    })
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
fn create_fetch_step_for_root_move(
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    root_step_index: NodeIndex,
    subgraph_name: &SubgraphName,
    type_name: &str,
    operation_kind: OperationKind,
    mutation_field_position: MutationFieldPosition,
    response_path: &MergePath,
    condition: Option<&Condition>,
) -> NodeIndex {
    let idx = fetch_graph.add_step(FetchStepData {
        id: fetch_graph.create_fetch_id(),
        service_name: subgraph_name.clone(),
        response_path: response_path.clone(),
        input: FetchStepSelections::new(type_name),
        output: FetchStepSelections::new(type_name),
        flags: FetchStepFlags::empty(),
        condition: condition.cloned(),
        kind: FetchStepKind::Root,
        operation_kind,
        variable_usages: None,
        variable_definitions: None,
        input_rewrites: None,
        output_rewrites: None,
        mutation_field_position,
        internal_aliases_locations: Vec::new(),
    });

    fetch_graph.connect(root_step_index, idx);

    idx
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
fn ensure_fetch_step_for_subgraph(
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    parent_fetch_step_index: NodeIndex,
    subgraph_name: &SubgraphName,
    input_type_name: &str,
    output_type_name: &str,
    response_path: &MergePath,
    key: Option<&TypeAwareSelection>,
    requires: Option<&TypeAwareSelection>,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<NodeIndex, FetchGraphError> {
    let matching_child_index = if requires.is_some() {
        None
    } else {
        fetch_graph
            .children_of(parent_fetch_step_index)
            .find_map(|to_child_edge_ref| {
                if let Ok(fetch_step) = fetch_graph.get_step_data(to_child_edge_ref.target()) {
                    if fetch_step.service_name != *subgraph_name {
                        return None;
                    }

                    if fetch_step.input.definition_name() != input_type_name {
                        return None;
                    }

                    if fetch_step.response_path != *response_path {
                        return None;
                    }

                    if fetch_step.condition.as_ref() != condition {
                        return None;
                    }

                    if let Some(key) = &key {
                        if fetch_step.input.definition_name() != key.type_name
                            || !fetch_step
                                .input
                                .selection_set()
                                .contains(&key.selection_set)
                        {
                            // requested key fields are not part of the input
                            return None;
                        }
                    }

                    // If there are requirements, then we do not re-use
                    // optimizations will try to re-use the existing step later, if possible.
                    if fetch_step.flags.contains(FetchStepFlags::USED_FOR_REQUIRES)
                        || requires.is_some()
                    {
                        return None;
                    }

                    return Some(to_child_edge_ref.target());
                }

                None
            })
    };

    match matching_child_index {
        Some(idx) => {
            trace!(
                "found existing fetch step [{}] for entity move requirement({}) key({}) in children of {}",
                idx.index(),
                requires.map(|r| r.to_string()).unwrap_or_default(),
                key.map(|r| r.to_string()).unwrap_or_default(),
                parent_fetch_step_index.index(),
            );
            Ok(idx)
        }
        None => {
            let step_index = create_fetch_step_for_entity_call(
                fetch_graph,
                subgraph_name,
                input_type_name,
                output_type_name,
                response_path,
                condition,
                created_from_requires || requires.is_some(),
            );
            if let Some(selection) = key {
                let step = fetch_graph.get_step_data_mut(step_index)?;
                step.input.add(&selection.selection_set)?
            }

            trace!(
                "created a new fetch step [{}] subgraph({}) type({}) requirement({}) key({}) in children of {}",
                step_index.index(),
                subgraph_name,
                input_type_name,
                requires.map(|r| r.to_string()).unwrap_or_default(),
                key.map(|r| r.to_string()).unwrap_or_default(),
                parent_fetch_step_index.index(),
            );

            Ok(step_index)
        }
    }
}

fn ensure_fetch_step_for_requirement(
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    parent_fetch_step_index: NodeIndex,
    subgraph_name: &SubgraphName,
    type_name: &String,
    response_path: &MergePath,
    condition: Option<&Condition>,
    requirement: &TypeAwareSelection,
) -> Result<NodeIndex, FetchGraphError> {
    let matching_child_index =
        fetch_graph
            .children_of(parent_fetch_step_index)
            .find_map(|to_child_edge_ref| {
                if let Ok(fetch_step) = fetch_graph.get_step_data(to_child_edge_ref.target()) {
                    if fetch_step.service_name != *subgraph_name {
                        return None;
                    }

                    if fetch_step.condition.as_ref() != condition {
                        return None;
                    }

                    if fetch_step.input.definition_name() != *type_name {
                        return None;
                    }

                    if fetch_step.response_path != *response_path {
                        return None;
                    }

                    if !fetch_step
                        .input
                        .contains(&requirement.type_name, &requirement.selection_set)
                    {
                        return None;
                    }

                    return Some(to_child_edge_ref.target());
                }

                None
            });

    match matching_child_index {
        Some(idx) => {
            trace!(
                "found existing fetch step [{}] children of {}",
                idx.index(),
                parent_fetch_step_index.index(),
            );
            Ok(idx)
        }
        None => {
            let step_index = create_fetch_step_for_entity_call(
                fetch_graph,
                subgraph_name,
                type_name,
                type_name,
                response_path,
                condition,
                true,
            );

            trace!(
                "created a new fetch step [{}] subgraph({}) type({}) requirement({}) in children of {}",
                step_index.index(),
                subgraph_name,
                type_name,
                requirement.to_string(),
                parent_fetch_step_index.index(),
            );

            Ok(step_index)
        }
    }
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace", skip_all, fields(
  count = query_node.children.len(),
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|s| s.index()),
  fetch_path = fetch_path.to_string(),
  response_path = response_path.to_string()
))]
fn process_children_for_fetch_steps(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    if query_node.children.is_empty() {
        return Ok(vec![parent_fetch_step_index]);
    }

    let mut leaf_fetch_step_indexes: Vec<NodeIndex> = vec![];
    for sub_step in query_node.children.iter() {
        leaf_fetch_step_indexes.extend(process_query_node(
            graph,
            fetch_graph,
            supergraph,
            override_context,
            sub_step,
            Some(parent_fetch_step_index),
            response_path,
            fetch_path,
            requiring_fetch_step_index,
            condition,
            created_from_requires,
        )?);
    }

    Ok(leaf_fetch_step_indexes)
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  count = query_node.requirements.len(),
  fetch_path = fetch_path.to_string(),
  response_path = response_path.to_string()
))]
fn process_requirements_for_fetch_steps(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    requiring_fetch_step_index: Option<NodeIndex>,
    condition: Option<&Condition>,
    response_path: &MergePath,
    fetch_path: &MergePath,
) -> Result<(), FetchGraphError> {
    if query_node.requirements.is_empty() {
        return Ok(());
    }

    for req_query_node in query_node.requirements.iter() {
        process_query_node(
            graph,
            fetch_graph,
            supergraph,
            override_context,
            req_query_node,
            Some(parent_fetch_step_index),
            response_path,
            fetch_path,
            requiring_fetch_step_index,
            condition,
            true,
        )?;
    }

    Ok(())
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
fn process_noop_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: Option<NodeIndex>,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    // We're at the root
    let fetch_step_index = parent_fetch_step_index
        .unwrap_or_else(|| create_noop_fetch_step(fetch_graph, created_from_requires));

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        fetch_step_index,
        response_path,
        fetch_path,
        requiring_fetch_step_index,
        condition,
        created_from_requires,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  edge = graph.pretty_print_edge(edge_index, false),
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
))]
fn process_entity_move_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    edge_index: EdgeIndex,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let edge = graph.edge(edge_index)?;
    let (requirement, is_interface) = match edge {
        Edge::EntityMove(em) => (
            TypeAwareSelection {
                selection_set: em.requirements.selection_set.clone(),
                type_name: em.requirements.type_name.clone(),
            },
            em.is_interface,
        ),
        _ => {
            return Err(FetchGraphError::UnexpectedEdgeMove(
                "EntityMove".to_string(),
            ))
        }
    };

    let head_node_index = graph.get_edge_head(&edge_index)?;
    let head_node = graph.node(head_node_index)?;
    let input_type_name = match head_node {
        Node::SubgraphType(t) => &t.name,
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    let tail_node_index = graph.get_edge_tail(&edge_index)?;
    let tail_node = graph.node(tail_node_index)?;
    let (output_type_name, subgraph_name) = match tail_node {
        Node::SubgraphType(t) => (&t.name, &t.subgraph),
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    let fetch_step_index = ensure_fetch_step_for_subgraph(
        fetch_graph,
        parent_fetch_step_index,
        subgraph_name,
        input_type_name,
        output_type_name,
        response_path,
        Some(&requirement),
        None,
        condition,
        created_from_requires,
    )?;

    let fetch_step = fetch_graph.get_step_data_mut(fetch_step_index)?;
    trace!(
        "adding input requirement '{}' to fetch step [{}]",
        requirement,
        fetch_step_index.index()
    );
    fetch_step.input.add(&requirement.selection_set)?;

    if is_interface {
        // We use `output_type_name` as there's no connection from `Interface` to `Object`,
        // it's always Object -> Interface.
        trace!(
            "adding input rewrite '... on {} {{ __typename }}' to '{}' of [{}]",
            output_type_name,
            output_type_name,
            fetch_step_index.index()
        );

        fetch_step.add_input_rewrite(FetchRewrite::ValueSetter(ValueSetter {
            path: vec![
                FetchNodePathSegment::typename_equals_from_type(output_type_name.to_string()),
                FetchNodePathSegment::Key("__typename".to_string()),
            ],
            set_value_to: output_type_name.clone(),
        }));

        fetch_step
            .flags
            .insert(FetchStepFlags::USED_FOR_TYPE_CONDITION);
    }

    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    parent_fetch_step
        .output
        .add_selection_typename(fetch_path)?;

    // Make the fetch step a child of the parent fetch step
    trace!(
        "connecting fetch step to parent [{}] -> [{}]",
        parent_fetch_step_index.index(),
        fetch_step_index.index()
    );
    fetch_graph.connect(parent_fetch_step_index, fetch_step_index);

    process_requirements_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        parent_fetch_step_index,
        Some(fetch_step_index),
        condition,
        response_path,
        fetch_path,
    )?;

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        fetch_step_index,
        response_path,
        &MergePath::default(),
        requiring_fetch_step_index,
        condition,
        created_from_requires,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  edge = graph.pretty_print_edge(edge_index, false),
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
))]
fn process_interface_object_type_move_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    edge_index: EdgeIndex,
    object_type_name: &str,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    if fetch_graph.parents_of(parent_fetch_step_index).count() != 1 {
        return Err(FetchGraphError::NonSingleParent(
            parent_fetch_step_index.index(),
        ));
    }
    let edge = graph.edge(edge_index)?;
    let requirement = match edge {
        Edge::InterfaceObjectTypeMove(m) => TypeAwareSelection {
            selection_set: m.requirements.selection_set.clone(),
            type_name: m.requirements.type_name.clone(),
        },
        _ => {
            return Err(FetchGraphError::UnexpectedEdgeMove(
                "InterfaceObjectTypeMove".to_string(),
            ))
        }
    };

    let head_node_index = graph.get_edge_head(&edge_index)?;
    let head_node = graph.node(head_node_index)?;
    let head_subgraph_name = match head_node {
        Node::SubgraphType(t) => &t.subgraph,
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    let tail_node_index = graph.get_edge_tail(&edge_index)?;
    let tail_node = graph.node(tail_node_index)?;
    let (interface_type_name, subgraph_name) = match tail_node {
        Node::SubgraphType(t) => (&t.name, &t.subgraph),
        // todo: FetchGraphError::MissingSubgraphName(tail_node.clone())
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    let step_for_children_index = ensure_fetch_step_for_subgraph(
        fetch_graph,
        parent_fetch_step_index,
        subgraph_name,
        object_type_name,
        interface_type_name,
        response_path,
        Some(&requirement),
        None,
        condition,
        created_from_requires,
    )?;

    let step_for_children = fetch_graph.get_step_data_mut(step_for_children_index)?;
    trace!(
        "adding input requirement '{}' to fetch step [{}]",
        requirement,
        step_for_children_index.index()
    );
    step_for_children.input.add(&requirement.selection_set)?;
    let key_to_reenter_subgraph = find_satisfiable_key(
        graph,
        supergraph,
        override_context,
        query_node.requirements.first().unwrap(),
    )?;
    trace!(
        "adding key '{}' to fetch step [{}]",
        key_to_reenter_subgraph,
        step_for_children_index.index()
    );
    step_for_children
        .input
        .add(&key_to_reenter_subgraph.selection_set)?;

    trace!(
        "adding input rewrite '... on {} {{ __typename }}' to '{}' of [{}]",
        interface_type_name,
        interface_type_name,
        step_for_children_index.index()
    );
    step_for_children.add_input_rewrite(FetchRewrite::ValueSetter(ValueSetter {
        path: vec![
            FetchNodePathSegment::typename_equals_from_type(interface_type_name.to_string()),
            FetchNodePathSegment::Key("__typename".to_string()),
        ],
        set_value_to: interface_type_name.clone(),
    }));

    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    parent_fetch_step
        .output
        .add_selection_typename(fetch_path)?;

    // In all cases it's `__typename` that needs to be resolved by another subgraph.
    trace!("Creating a fetch step for requirement of @interfaceObject");
    let step_for_requirements_index = create_fetch_step_for_entity_call(
        fetch_graph,
        head_subgraph_name,
        object_type_name,
        interface_type_name,
        response_path,
        condition,
        false,
    );
    let step_for_requirements = fetch_graph.get_step_data_mut(step_for_requirements_index)?;
    trace!(
        "Adding {} to fetch([{}]).input",
        key_to_reenter_subgraph,
        step_for_requirements_index.index()
    );
    step_for_requirements
        .input
        .add(&key_to_reenter_subgraph.selection_set)?;

    //
    // Given `f0 { ... on User { f1 } }` where f1 is a field contributed by @interfaceObject,
    // the requirement to resolve `f1` with `User` type condition is a real value of `__typename` field.
    // - step_for_requirements    -> __typename
    // - step_for_children        -> f1
    // - parent                   -> f0
    //
    // parent -> step_for_requirements -> step_for_children
    //
    // f0 -> __typename -> f1
    //

    fetch_graph.connect(parent_fetch_step_index, step_for_requirements_index);

    trace!("Processing requirements");
    let leaf_fetch_step_indexes = process_query_node(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node.requirements.first().unwrap(),
        Some(step_for_requirements_index),
        response_path,
        &MergePath::default(),
        None,
        condition,
        true,
    )?;

    if leaf_fetch_step_indexes.is_empty() {
        fetch_graph.connect(step_for_requirements_index, step_for_children_index);
    } else {
        for idx in leaf_fetch_step_indexes {
            fetch_graph.connect(idx, step_for_children_index);
        }
    }

    trace!("Processing children");
    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        step_for_children_index,
        response_path,
        &MergePath::default(),
        requiring_fetch_step_index,
        condition,
        created_from_requires,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  subgraph = subgraph_name.0,
  type_name = type_name,
  parent_fetch_step_index = parent_fetch_step_index.index(),
))]
fn process_subgraph_entrypoint_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    subgraph_name: &SubgraphName,
    type_name: &str,
    operation_kind: OperationKind,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let fetch_step_index = create_fetch_step_for_root_move(
        fetch_graph,
        parent_fetch_step_index,
        subgraph_name,
        type_name,
        operation_kind,
        query_node.mutation_field_position,
        &MergePath::default(),
        None,
    );

    fetch_graph.connect(parent_fetch_step_index, fetch_step_index);

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        fetch_step_index,
        &MergePath::default(),
        &MergePath::default(),
        None,
        None,
        created_from_requires,
    )
}

#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  parent_fetch_step_index = parent_fetch_step_index.index(),
))]
fn process_subgraph_reentry(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    edge_index: EdgeIndex,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    created_from_requires: bool,
    reentry_move: &ReentryMove,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    condition: Option<&Condition>,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let node_index = graph.get_edge_tail(&edge_index).unwrap();
    let node = graph.node(node_index).unwrap();

    let (type_name, subgraph_name) = match node {
        Node::SubgraphType(subgraph_type) => (&subgraph_type.name, &subgraph_type.subgraph),
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    // The re-entry field is part of the parent's response shape, but its children resolve in a
    // separate (re-entered) fetch. Record the field in the parent output with a `__typename`
    // so the parent operation stays valid.
    //
    let condition_in_path = if let Some(condition) = condition {
        fetch_path.inner.iter().any(|segment| {
            matches!(
                segment,
                Segment::TypeCondition(_, Some(c)) | Segment::Field(_, _, Some(c)) if c == condition
            )
        })
    } else {
        false
    };

    let ancestor_of_condition = match condition {
        Some(c) => fetch_graph.is_ancestor_of_condition(parent_fetch_step_index, c),
        None => false,
    };

    let should_strip_condition = condition_in_path || ancestor_of_condition;
    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    parent_fetch_step.output.add_at_path(
        fetch_path,
        SelectionSet {
            items: vec![SelectionItem::Field(FieldSelection {
                name: reentry_move.name.to_string(),
                alias: query_node.selection_alias().map(|a| a.to_string()),
                selections: SelectionSet {
                    items: vec![SelectionItem::Field(FieldSelection::new_typename())],
                },
                arguments: query_node.selection_arguments().cloned(),
                skip_if: condition.and_then(|c| {
                    if should_strip_condition {
                        None
                    } else {
                        c.to_skip_if()
                    }
                }),
                include_if: condition.and_then(|c| {
                    if should_strip_condition {
                        None
                    } else {
                        c.to_include_if()
                    }
                }),
                omit_from_response: false,
            })],
        },
    )?;

    let segment_field = FieldPathSegment::new(
        reentry_move.name.to_string(),
        query_node.selection_alias().map(|a| a.to_string()),
    );
    let mut child_response_path = response_path.push(Segment::Field(segment_field, 0, None));
    if reentry_move.is_list {
        child_response_path = child_response_path.push(Segment::List);
    }

    let fetch_step_index = create_fetch_step_for_root_move(
        fetch_graph,
        parent_fetch_step_index,
        subgraph_name,
        type_name,
        OperationKind::Query,
        query_node.mutation_field_position,
        &child_response_path,
        condition,
    );

    fetch_graph.connect(parent_fetch_step_index, fetch_step_index);

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        fetch_step_index,
        &child_response_path,
        &MergePath::default(),
        requiring_fetch_step_index,
        None,
        created_from_requires,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
  type_name = target_type_name,
  response_path = response_path.to_string(),
  fetch_path = fetch_path.to_string(),
))]
fn process_selfie_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    requiring_fetch_step_index: Option<NodeIndex>,
    response_path: &MergePath,
    fetch_path: &MergePath,
    target_type_name: &String,
    condition: Option<&Condition>,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let is_ancestor_of_condition = match condition {
        Some(c) => fetch_graph.is_ancestor_of_condition(parent_fetch_step_index, c),
        None => false,
    };

    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    trace!(
        "starting an inline fragment for type '{}' to fetch step [{}]",
        parent_fetch_step_index.index(),
        target_type_name,
    );
    parent_fetch_step.output.add_at_path(
        fetch_path,
        SelectionSet {
            items: vec![SelectionItem::InlineFragment(InlineFragmentSelection {
                type_condition: target_type_name.clone(),
                selections: SelectionSet::default(),
                skip_if: condition.and_then(|c| {
                    if is_ancestor_of_condition {
                        None
                    } else {
                        c.to_skip_if()
                    }
                }),
                include_if: condition.and_then(|c| {
                    if is_ancestor_of_condition {
                        None
                    } else {
                        c.to_include_if()
                    }
                }),
            })],
        },
    )?;

    let segment_condition = if is_ancestor_of_condition {
        None
    } else {
        condition.cloned()
    };
    let child_response_path = response_path.push(Segment::TypeCondition(
        BTreeSet::from([target_type_name.clone()]),
        segment_condition.clone(),
    ));
    let child_fetch_path = fetch_path.push(Segment::TypeCondition(
        BTreeSet::from([target_type_name.clone()]),
        segment_condition,
    ));

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        parent_fetch_step_index,
        &child_response_path,
        &child_fetch_path,
        requiring_fetch_step_index,
        condition,
        false,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
  type_name = target_type_name,
  response_path = response_path.to_string(),
  fetch_path = fetch_path.to_string()
))]
fn process_abstract_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    requiring_fetch_step_index: Option<NodeIndex>,
    response_path: &MergePath,
    fetch_path: &MergePath,
    target_type_name: &String,
    condition: Option<&Condition>,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    trace!(
        "adding output field '__typename' and starting an inline fragment for type '{}' to fetch step [{}]",
        parent_fetch_step_index.index(),
        target_type_name,
    );

    parent_fetch_step.output.add_at_path(
        fetch_path,
        SelectionSet {
            items: vec![
                SelectionItem::Field(FieldSelection::new_typename()),
                SelectionItem::InlineFragment(InlineFragmentSelection {
                    type_condition: target_type_name.clone(),
                    selections: SelectionSet::default(),
                    skip_if: None,
                    include_if: None,
                }),
            ],
        },
    )?;

    let child_response_path = response_path.push(Segment::TypeCondition(
        BTreeSet::from([target_type_name.clone()]),
        None,
    ));
    let child_fetch_path = fetch_path.push(Segment::TypeCondition(
        BTreeSet::from([target_type_name.clone()]),
        None,
    ));

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        parent_fetch_step_index,
        &child_response_path,
        &child_fetch_path,
        requiring_fetch_step_index,
        condition,
        false,
    )
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
  type_name = field_move.type_name,
  field = field_move.name,
  alias = query_node.selection_attributes.as_ref().and_then(|v| v.alias.as_ref()),
  arguments = query_node.selection_attributes.as_ref().and_then(|v| v.arguments.as_ref()).map(|v| format!("{}", v)),
  leaf = field_move.is_leaf,
  list = field_move.is_list,
  response_path = response_path.to_string(),
  fetch_path = fetch_path.to_string(),
))]
fn process_plain_field_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    requiring_fetch_step_index: Option<NodeIndex>,
    response_path: &MergePath,
    fetch_path: &MergePath,
    field_move: &FieldMove,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    if let Some(requiring_fetch_step_index) = requiring_fetch_step_index {
        trace!(
            "connecting parent fetch step [{}] to requiring fetch step [{}]",
            parent_fetch_step_index.index(),
            requiring_fetch_step_index.index()
        );
        fetch_graph.connect(parent_fetch_step_index, requiring_fetch_step_index);
    }

    // If we are inserting into "... | [Product] @include(if: $bool)", we don't need the condition on the field.
    let condition_in_path = if let Some(condition) = condition {
        fetch_path.inner.iter().any(|segment| {
            matches!(
                segment,
                Segment::TypeCondition(_, Some(c)) | Segment::Field(_, _, Some(c)) if c == condition
            )
        })
    } else {
        false
    };

    let ancestor_of_condition = match condition {
        Some(c) => fetch_graph.is_ancestor_of_condition(parent_fetch_step_index, c),
        None => false,
    };

    let should_strip_condition = condition_in_path || ancestor_of_condition;

    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
    trace!(
        "adding output field '{}' to fetch step [{}]",
        field_move.name,
        parent_fetch_step_index.index()
    );

    parent_fetch_step.output.add_at_path(
        fetch_path,
        SelectionSet {
            items: vec![SelectionItem::Field(FieldSelection {
                name: field_move.name.to_string(),
                alias: query_node.selection_alias().map(|a| a.to_string()),
                selections: SelectionSet::default(),
                arguments: query_node.selection_arguments().cloned(),
                skip_if: condition.and_then(|c| {
                    if should_strip_condition {
                        None
                    } else {
                        c.to_skip_if()
                    }
                }),
                include_if: condition.and_then(|c| {
                    if should_strip_condition {
                        None
                    } else {
                        c.to_include_if()
                    }
                }),
                omit_from_response: false,
            })],
        },
    )?;

    let segment_args_hash = query_node
        .selection_arguments()
        .map(|a| a.hash_u64())
        .unwrap_or(0);
    let segment_field = FieldPathSegment::new(
        field_move.name.to_string(),
        query_node.selection_alias().map(|a| a.to_string()),
    );
    let segment_condition = if should_strip_condition {
        None
    } else {
        condition.cloned()
    };
    let mut child_response_path = response_path.push(Segment::Field(
        segment_field.clone(),
        segment_args_hash,
        segment_condition.clone(),
    ));
    let mut child_fetch_path = fetch_path.push(Segment::Field(
        segment_field,
        segment_args_hash,
        segment_condition,
    ));

    if field_move.is_list {
        child_response_path = child_response_path.push(Segment::List);
        child_fetch_path = child_fetch_path.push(Segment::List);
    }

    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        parent_fetch_step_index,
        &child_response_path,
        &child_fetch_path,
        requiring_fetch_step_index,
        condition,
        created_from_requires,
    )
}

#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",skip_all, fields(
  parent_fetch_step_index = parent_fetch_step_index.index(),
  requiring_fetch_step_index = requiring_fetch_step_index.map(|s| s.index()),
))]
fn process_requires_field_edge(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: NodeIndex,
    response_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    field_move: &FieldMove,
    edge_index: EdgeIndex,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    if fetch_graph.parents_of(parent_fetch_step_index).count() != 1 {
        return Err(FetchGraphError::NonSingleParent(
            parent_fetch_step_index.index(),
        ));
    }

    let parent_parent_index = fetch_graph
        .parents_of(parent_fetch_step_index)
        .next()
        .map(|edge| edge.source())
        .unwrap();

    let requires = field_move
        .requirements
        .as_ref()
        .ok_or(FetchGraphError::MissingRequires)?;

    let head_node_index = graph.get_edge_head(&edge_index)?;
    let head_node = graph.node(head_node_index)?;
    let (head_type_name, head_subgraph_name, is_interface_object) = match head_node {
        Node::SubgraphType(t) => (&t.name, &t.subgraph, &t.is_interface_object),
        _ => return Err(FetchGraphError::ExpectedSubgraphType),
    };

    let key_to_reenter_subgraph = find_satisfiable_key(
        graph,
        supergraph,
        override_context,
        query_node.requirements.first().unwrap(),
    )?;

    let parent_fetch_step = fetch_graph.get_step_data(parent_fetch_step_index)?;
    // In case of a field with `@requires`, the parent will be the current subgraph we're in.
    let real_parent_fetch_step_index = match !parent_fetch_step
        .output
        .is_selecting_definition(head_type_name)
    {
        // If the parent's output resolves a different type, then it's a root type.
        // We can use that as a parent.
        true => parent_fetch_step_index,
        // If the parent's output resolves the same type, it means we're in an entity call.
        // We need to move up, as the entity call was created to fetch regular fields of the type
        // (those without @requires).
        //
        // Example: Fetch Step was created to get `baz`
        // {
        //   foo
        //   bar @requires(fields: "foo")
        //   baz
        // }
        //
        // We need to stick to the parent of the parent.
        false => parent_parent_index,
    };

    // When a field (foo) is annotated with `@requires(fields: "bar")`
    // We want to create new FetchStep (entity move) for that field (foo)
    // or reuse an existing one if the requirement matches
    // The new FetchStep will have `foo` as the output and `bar` as the input.
    // The `bar` should be fetched by one of the parents.
    trace!("Creating a fetch step for children of @requires");
    let step_for_children_index = ensure_fetch_step_for_requirement(
        fetch_graph,
        real_parent_fetch_step_index,
        head_subgraph_name,
        head_type_name,
        response_path,
        condition,
        requires,
    )?;

    let step_for_children = fetch_graph.get_step_data_mut(step_for_children_index)?;

    step_for_children.output.add_at_path(
        &MergePath::default(),
        SelectionSet {
            items: vec![SelectionItem::Field(FieldSelection {
                name: field_move.name.to_string(),
                alias: query_node.selection_alias().map(|a| a.to_string()),
                selections: SelectionSet::default(),
                arguments: query_node.selection_arguments().cloned(),
                skip_if: None,
                include_if: None,
                omit_from_response: false,
            })],
        },
    )?;

    if *is_interface_object {
        trace!(
            "adding input rewrite '... on {} {{ __typename }}' to '{}' of [{}]",
            head_type_name,
            head_type_name,
            step_for_children_index.index(),
        );
        step_for_children.add_input_rewrite(FetchRewrite::ValueSetter(ValueSetter {
            path: vec![
                FetchNodePathSegment::typename_equals_from_type(head_type_name.to_string()),
                FetchNodePathSegment::Key("__typename".to_string()),
            ],
            set_value_to: head_type_name.clone(),
        }));
    }

    trace!(
        "Adding {} to fetch([{}]).input (requires)",
        requires,
        step_for_children_index.index()
    );
    step_for_children.input.add(&requires.selection_set)?;
    trace!(
        "Adding {} to fetch([{}]).input (key re-enter)",
        key_to_reenter_subgraph,
        step_for_children_index.index()
    );
    step_for_children
        .input
        .add(&key_to_reenter_subgraph.selection_set)?;

    trace!("Creating a fetch step for requirement of @requires");
    let step_for_requirements_index = create_fetch_step_for_entity_call(
        fetch_graph,
        head_subgraph_name,
        head_type_name,
        head_type_name,
        response_path,
        condition,
        true,
    );
    let step_for_requirements = fetch_graph.get_step_data_mut(step_for_requirements_index)?;
    trace!(
        "Adding {} to fetch([{}]).input",
        key_to_reenter_subgraph,
        step_for_requirements_index.index()
    );
    step_for_requirements
        .input
        .add(&key_to_reenter_subgraph.selection_set)?;

    let real_parent_fetch_step = fetch_graph.get_step_data_mut(real_parent_fetch_step_index)?;

    let key_to_reenter_at = if real_parent_fetch_step.response_path.len() > response_path.len() {
        return Err(FetchGraphError::Internal(
            "Response path is longer than expected".to_string(),
        ));
    } else {
        &response_path.slice_from(real_parent_fetch_step.response_path.len())
    };

    trace!(
        "Adding {} to fetch([{}]).output at path {}",
        key_to_reenter_subgraph,
        real_parent_fetch_step_index.index(),
        key_to_reenter_at
    );

    real_parent_fetch_step.output.add_at_path(
        key_to_reenter_at,
        key_to_reenter_subgraph.clone().selection_set,
    )?;

    real_parent_fetch_step
        .output
        .add_selection_typename(key_to_reenter_at)?;

    fetch_graph.connect(real_parent_fetch_step_index, step_for_requirements_index);

    let segment_args_hash = query_node
        .selection_arguments()
        .map(|a| a.hash_u64())
        .unwrap_or(0);
    let segment_field = FieldPathSegment::new(
        field_move.name.to_string(),
        query_node.selection_alias().map(|a| a.to_string()),
    );
    let mut child_response_path = response_path.push(Segment::Field(
        segment_field.clone(),
        segment_args_hash,
        None,
    ));
    let mut child_fetch_path =
        MergePath::default().push(Segment::Field(segment_field, segment_args_hash, None));

    if field_move.is_list {
        child_response_path = child_response_path.push(Segment::List);
        child_fetch_path = child_fetch_path.push(Segment::List);
    }

    trace!("Processing requirements");
    let leaf_fetch_step_indexes = process_query_node(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node.requirements.first().unwrap(),
        Some(step_for_requirements_index),
        response_path,
        &MergePath::default(),
        None,
        condition,
        true,
    )?;

    //
    // Given `f0 { f1 @requires(fields: f2) }`
    // - step_for_requirements    -> f2
    // - step_for_children        -> f1
    // - parent                   -> f0
    //
    // parent -> step_for_requirements -> step_for_children
    //
    // f0 -> f2 -> f1
    //
    // in case of `f2 @requires(fields: f3)`:
    // f0 -> f3 -> f2 -> f1
    //
    // and so on.
    //
    // Basically any leaf becomes a parent of current `step_for_children`.
    // This way we wait for the entire chain of fetches to be resolved before we move to resolve a field with `@requires`.

    if leaf_fetch_step_indexes.is_empty() {
        trace!("Connecting fetch that pulls requirements with fetch that resolves the field with @requires");
        fetch_graph.connect(step_for_requirements_index, step_for_children_index);
    } else {
        trace!("Connecting leaf fetches of requirements to fetch with @requires");
        for idx in leaf_fetch_step_indexes {
            fetch_graph.connect(idx, step_for_children_index);
        }
    }

    trace!("Processing children");
    process_children_for_fetch_steps(
        graph,
        fetch_graph,
        supergraph,
        override_context,
        query_node,
        step_for_children_index,
        &child_response_path,
        &child_fetch_path,
        requiring_fetch_step_index,
        condition,
        created_from_requires,
    )
}

/// A field marked with `@requires` tells us that in order to resolve it,
/// the subgraph needs data from another subgraph.
///
/// This interaction ALWAYS involves an entity call,
/// but the trigger for resolving the field differs:
///
/// 1.  Cross-Subgraph: If the field annotated with `@requires` is being
///     fetched after an `EntityMove` brought us to the current subgraph from another,
///     the `EntityMove` itself already handled fetching the necessary `@key` fields.
///     We don't need to think about it, we just need to add the field set of `@requires(fields:)`
///     to the `FetchStep` input.
///
/// 2.  Root: If the field annotated with `@requires` is being fetched locally,
///     because the query started from the root type,
///     the `EntityMove` is needed, but the data to satisfy the key fields is all local.
///     The subgraph effectively makes an "internal" entity call to itself.
///     - We first fetch the fields needed to satisfy some resolvable `@key` of the
///       entity type within this subgraph.
///     - Then, we use that to resolve the fields specified in the `@requires(fields:)`.
///     - This effectively splits the resolution within the subgraph: part comes from the
///       main query path, part comes via the internal entity resolution triggered by `@requires`.
///
/// The key fields need to be added to the output of the parent,
/// and input of the entity move.
#[instrument(level = "trace",skip_all, fields(
  node = graph.node(query_node.node_index).unwrap().to_string()
))]
fn find_satisfiable_key<'a>(
    graph: &'a Graph,
    supergraph: &'a SupergraphState,
    override_context: &'a PlannerOverrideContext,
    query_node: &QueryTreeNode,
) -> Result<&'a TypeAwareSelection, FetchGraphError> {
    // This could be improved...
    // We added a flag to `can_satisfy_edge` and increased the complexity.

    let mut entity_moves_edges_to_self: Vec<
        EdgeReference<crate::query_planner::graph::edge::Edge>,
    > = graph
        .edges_from(query_node.node_index)
        .filter(|edge_reference| {
            let edge = graph.edge(edge_reference.id()).unwrap();
            matches!(edge, Edge::EntityMove(_))
        })
        .collect();
    entity_moves_edges_to_self.sort_by_key(|edge| std::cmp::Reverse(edge.weight().cost()));

    for edge_ref in entity_moves_edges_to_self {
        if can_satisfy_edge(
            graph,
            supergraph,
            override_context,
            &edge_ref,
            &OperationPath {
                root_node: query_node.node_index,
                last_segment: None,
                cost: 0,
                union_context: None,
            },
            &Default::default(),
            true,
            // It's safe to use a noop CancellationToken here,
            // as the result of this function is guaranteed to be successful,
            // and fast.
            &CancellationToken::new(),
        )?
        .is_some()
        {
            return edge_ref
                .weight()
                .requirements()
                .ok_or(FetchGraphError::Internal(String::from(
                    "Resolved empty Satisfiable Key",
                )));
        }
    }

    Err(FetchGraphError::Internal(String::from(
        "Failed to find Satisfiable Key",
    )))
}

// TODO: simplfy args
#[allow(clippy::too_many_arguments)]
fn process_query_node(
    graph: &Graph,
    fetch_graph: &mut FetchGraph<SingleTypeFetchStep>,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_node: &QueryTreeNode,
    parent_fetch_step_index: Option<NodeIndex>,
    response_path: &MergePath,
    fetch_path: &MergePath,
    requiring_fetch_step_index: Option<NodeIndex>,
    condition: Option<&Condition>,
    created_from_requires: bool,
) -> Result<Vec<NodeIndex>, FetchGraphError> {
    let condition = query_node.condition.as_ref().or(condition);
    if let Some(edge_index) = query_node.edge_from_parent {
        let parent_fetch_step_index = parent_fetch_step_index.ok_or(FetchGraphError::IndexNone)?;
        let edge = graph.edge(edge_index)?;

        match edge {
            Edge::SubgraphEntrypoint {
                name,
                operation_kind,
            } => {
                let tail_node_index = graph.get_edge_tail(&edge_index)?;
                let tail_node = graph.node(tail_node_index)?;
                let type_name = match tail_node {
                    Node::QueryRoot(t) => t,
                    Node::MutationRoot(t) => t,
                    Node::SubscriptionRoot(t) => t,
                    Node::SubgraphType(t) => &t.name,
                };

                process_subgraph_entrypoint_edge(
                    graph,
                    fetch_graph,
                    supergraph,
                    override_context,
                    query_node,
                    parent_fetch_step_index,
                    name,
                    type_name,
                    operation_kind.clone(),
                    created_from_requires,
                )
            }
            Edge::EntityMove(_) => process_entity_move_edge(
                graph,
                fetch_graph,
                supergraph,
                override_context,
                query_node,
                parent_fetch_step_index,
                response_path,
                fetch_path,
                requiring_fetch_step_index,
                edge_index,
                condition,
                created_from_requires,
            ),
            Edge::ReentryMove(field) => process_subgraph_reentry(
                graph,
                fetch_graph,
                supergraph,
                edge_index,
                override_context,
                query_node,
                parent_fetch_step_index,
                created_from_requires,
                field,
                response_path,
                fetch_path,
                requiring_fetch_step_index,
                condition,
            ),
            Edge::FieldMove(field) => match field.requirements.is_some() {
                true => process_requires_field_edge(
                    graph,
                    fetch_graph,
                    supergraph,
                    override_context,
                    query_node,
                    parent_fetch_step_index,
                    response_path,
                    requiring_fetch_step_index,
                    field,
                    edge_index,
                    condition,
                    created_from_requires,
                ),
                false => process_plain_field_edge(
                    graph,
                    fetch_graph,
                    supergraph,
                    override_context,
                    query_node,
                    parent_fetch_step_index,
                    requiring_fetch_step_index,
                    response_path,
                    fetch_path,
                    field,
                    condition,
                    created_from_requires,
                ),
            },
            Edge::Selfie(type_name) => process_selfie_edge(
                graph,
                fetch_graph,
                supergraph,
                override_context,
                query_node,
                parent_fetch_step_index,
                requiring_fetch_step_index,
                response_path,
                fetch_path,
                type_name,
                condition,
            ),
            Edge::AbstractMove(type_name) => process_abstract_edge(
                graph,
                fetch_graph,
                supergraph,
                override_context,
                query_node,
                parent_fetch_step_index,
                requiring_fetch_step_index,
                response_path,
                fetch_path,
                type_name,
                condition,
            ),
            Edge::InterfaceObjectTypeMove(InterfaceObjectTypeMove {
                object_type_name, ..
            }) => process_interface_object_type_move_edge(
                graph,
                fetch_graph,
                supergraph,
                override_context,
                query_node,
                parent_fetch_step_index,
                response_path,
                fetch_path,
                requiring_fetch_step_index,
                edge_index,
                object_type_name,
                condition,
                created_from_requires,
            ),
        }
    } else {
        process_noop_edge(
            graph,
            fetch_graph,
            supergraph,
            override_context,
            query_node,
            parent_fetch_step_index,
            response_path,
            fetch_path,
            requiring_fetch_step_index,
            condition,
            created_from_requires,
        )
    }
}

pub fn find_graph_roots<State>(graph: &FetchGraph<State>) -> Vec<NodeIndex> {
    let mut roots = Vec::new();

    // Iterate over all nodes in the graph
    for node_idx in graph.step_indices() {
        // Check if the node has any incoming edges.
        // The `next().is_none()` checks if the iterator is empty - no incoming edges.
        if graph.parents_of(node_idx).next().is_none() {
            roots.push(node_idx);
        }
    }

    roots
}

#[instrument(level = "trace", skip(graph, query_tree, supergraph, override_context), fields(
    requirements_count = query_tree.root.requirements.len(),
    children_count = query_tree.root.children.len(),
))]
pub fn build_fetch_graph_from_query_tree(
    graph: &Graph,
    supergraph: &SupergraphState,
    override_context: &PlannerOverrideContext,
    query_tree: QueryTree,
    operation_kind: OperationKind,
    options: &QueryPlannerOptions,
    cancellation_token: &CancellationToken,
) -> Result<FetchGraph<MultiTypeFetchStep>, FetchGraphError> {
    let mut fetch_graph = FetchGraph::new(operation_kind);

    process_query_node(
        graph,
        &mut fetch_graph,
        supergraph,
        override_context,
        &query_tree.root,
        None,
        &MergePath::default(),
        &MergePath::default(),
        None,
        None,
        false,
    )?;

    trace!("Done");

    let root_indexes = find_graph_roots(&fetch_graph);

    trace!("found roots");

    if root_indexes.is_empty() {
        return Err(FetchGraphError::NonSingleRootStep(0));
    }

    if root_indexes.len() > 1 {
        return Err(FetchGraphError::NonSingleRootStep(root_indexes.len()));
    }

    trace!("fetch graph before optimizations:");
    trace!("{}", fetch_graph);

    // fine to unwrap as we have already checked the length
    fetch_graph.root_index = Some(*root_indexes.first().unwrap());
    let mut fetch_graph = fetch_graph.to_multi_type();
    fetch_graph.optimize(supergraph, options, cancellation_token)?;
    fetch_graph.collect_variable_usages()?;

    trace!("fetch graph after optimizations:");
    trace!("{}", fetch_graph);

    Ok(fetch_graph)
}