hive-router 0.3.2

use std::collections::{HashMap, HashSet, VecDeque};

use petgraph::{
    graph::NodeIndex,
    visit::{EdgeRef, NodeRef},
};
use tracing::{instrument, trace};

use crate::query_planner::{
    ast::{
        merge_path::{MergePath, Segment},
        selection_set::{find_arguments_conflicts, find_selection_set_by_path},
    },
    planner::fetch::{
        error::FetchGraphError,
        fetch_graph::FetchGraph,
        fetch_step_data::{FetchStepData, FetchStepKind},
        response_keys::read_as_written,
        selections::FetchStepSelections,
        state::MultiTypeFetchStep,
    },
};

/// Where merged-away steps went, kept for a whole pass. Passes queue steps and pick merge pairs
/// before merging, so a step they hold on to may be gone by the time they use it. A step merged
/// into one that got merged again later resolves to the last one, the one still in the graph.
#[derive(Default)]
pub(crate) struct MergedSteps(HashMap<NodeIndex, NodeIndex>);

impl MergedSteps {
    pub(crate) fn resolve(&self, mut step: NodeIndex) -> NodeIndex {
        while let Some(into) = self.0.get(&step) {
            step = *into;
        }
        step
    }

    pub(crate) fn record(&mut self, merged: NodeIndex, into: NodeIndex) {
        self.0.insert(merged, into);
    }
}

/// Handles the "target is non-entity, source has step-level condition" case.
/// When merging an entity fetch into a non-entity target, the condition must
/// stay attached to the source branch data, not to the whole target step.
fn merge_source_condition_into_non_entity_target(
    target: &FetchStepData<MultiTypeFetchStep>,
    source: &mut FetchStepData<MultiTypeFetchStep>,
) -> Result<bool, FetchGraphError> {
    let Some(condition) = source.condition.clone() else {
        return Ok(false);
    };

    if target.is_entity_call() {
        return Ok(false);
    }

    // The source fetch step is an entity call, so it has both input and output.
    // The target fetch step is a regular (non-entity) query fetch.
    //
    // We can safely migrate source.input into source.output before merging:
    //
    // Input             -> Output
    // { id __typename } -> { price }
    //
    // becomes
    // { products { ... on Product @skip(if: $bool) { __typename id price } } }
    //
    // We do this only for non-entity target merges. Entity-to-entity merges use
    // different path/type rules and are handled in a separate branch.
    // The input is written the way a representation reads it (`id: _internal_qp_alias_0`),
    // so it goes in the way the fetch writes it (`_internal_qp_alias_0: id`).
    let mut input = source.input.clone();
    for (_, selection_set) in input.iter_selections_mut() {
        read_as_written(selection_set);
    }
    source
        .output
        .migrate_from_another(&input, &MergePath::default())?;

    // Check if the condition is already enforced by the path
    let condition_redundant = matches!(
        source.response_path.last(),
        Some(Segment::TypeCondition(_, Some(c)) | Segment::Field(_, _, Some(c))) if c == &condition
    );

    if !condition_redundant {
        source.output.wrap_with_condition(condition);
    }

    Ok(true)
}

#[instrument(level = "trace", skip_all)]
pub(crate) fn perform_fetch_step_merge(
    target_index: NodeIndex,
    source_index: NodeIndex,
    fetch_graph: &mut FetchGraph<MultiTypeFetchStep>,
    force_merge_inputs: bool,
) -> Result<(), FetchGraphError> {
    let (target, source) = fetch_graph.get_pair_of_steps_mut(target_index, source_index)?;

    trace!(
        "merging fetch steps [{}] + [{}]",
        target_index.index(),
        source_index.index(),
    );

    let source_condition_merged = merge_source_condition_into_non_entity_target(target, source)?;
    if !source_condition_merged {
        target.scope_fetch_conditions_before_merge(source);
    }

    let source_fetch_path = source.response_path.slice_from(target.response_path.len());
    // At the same spot both steps write the same object. If the target fetches one type and
    // the source another, the source's fields would land on the target's type: `Cat` fields
    // on `Animal`. `can_merge` keeps such steps apart, so getting here is a bug.
    if source_fetch_path.is_empty() {
        if let Some(target_type) = target.output.try_as_single() {
            if let Some((source_type, _)) = source
                .output
                .iter_selections()
                .find(|(source_type, _)| source_type.as_str() != target_type)
            {
                return Err(FetchGraphError::Internal(format!(
                    "can't merge the `{}` fields of step [{}] into the `{}` fields of step [{}]",
                    source_type,
                    source_index.index(),
                    target_type,
                    target_index.index(),
                )));
            }
        }
    }
    target
        .output
        .migrate_from_another(&source.output, &source_fetch_path)?;

    if let Some(input_rewrites) = source.input_rewrites.take() {
        if !input_rewrites.is_empty() {
            for input_rewrite in input_rewrites {
                target.add_input_rewrite(input_rewrite);
            }
        }
    }

    if force_merge_inputs {
        target
            .input
            .migrate_from_another(&source.input, &MergePath::default())?;
    } else if target.input.selecting_same_types(&source.input) {
        // It's safe to not check if a condition was turned into an inline fragment,
        // because if a condition is present and "me" is a non-entity fetch step,
        // then the type_name values of the inputs are different.
        if target.response_path == source.response_path {
            target
                .input
                .migrate_from_another(&source.input, &MergePath::default())?;
        } else if !source.response_path.starts_with(&target.response_path) {
            return Err(FetchGraphError::MismatchedResponsePath);
        }
        // Otherwise we pulled in a step from deeper in the response. Our own output
        // already has its input - that is why we could pull it in - so there is
        // nothing to copy over. See `can_absorb_nested_entity_call`.
    }

    // Conditions may have been pushed down to keep the merge correct.
    // If the merged fetch is still guarded by one shared condition, lift it back to
    // step level.
    if target.is_entity_call() {
        target.lift_shared_output_condition_to_fetch();
    }

    let mut children_indexes: Vec<NodeIndex> = vec![];
    let mut parents_indexes: Vec<NodeIndex> = vec![];
    for edge_ref in fetch_graph.children_of(source_index) {
        children_indexes.push(edge_ref.target().id());
    }

    for edge_ref in fetch_graph.parents_of(source_index) {
        // We ignore self_index
        if edge_ref.source().id() != target_index {
            parents_indexes.push(edge_ref.source().id());
        }
    }

    // Replace parents:
    // 1. Add self -> child
    for child_index in children_indexes.iter() {
        fetch_graph.connect(target_index, *child_index);
    }
    // 2. Add parent -> self
    for parent_index in parents_indexes {
        fetch_graph.connect(parent_index, target_index);
    }
    // 3. Drop other -> child and parent -> other
    fetch_graph.remove_step(source_index);

    Ok(())
}

/// Checks if an ancestor node (`target_ancestor_index`) is reachable from a
/// child node (`child_index`) in a directed graph by following paths upwards
/// (traversing incoming edges), EXCLUDING any paths that start by traversing
/// the direct edge from the `target_ancestor_index` down to the `child_index`.
///
/// This is implemented as an iterative Breadth-First Search (BFS).
/// The search starts from all direct parents of `child_index` *except*
/// `target_ancestor_index`, and follows incoming edges from there.
pub fn is_reachable_via_alternative_upstream_path(
    graph: &FetchGraph<MultiTypeFetchStep>,
    child_index: NodeIndex,
    target_ancestor_index: NodeIndex,
) -> Result<bool, FetchGraphError> {
    let mut queue: VecDeque<NodeIndex> = VecDeque::new();
    let mut visited: HashSet<NodeIndex> = HashSet::new();

    // Start BFS queue with all parents of `child_index` except `target_ancestor_index`
    for edge_ref in graph.parents_of(child_index) {
        let parent_index = edge_ref.source();

        if parent_index != target_ancestor_index {
            queue.push_back(parent_index);
            visited.insert(parent_index);
        }
    }

    if queue.is_empty() {
        return Ok(false);
    }

    // Perform BFS upwards (following incoming edges)
    while let Some(current_index) = queue.pop_front() {
        // If we reached the target ancestor indirectly
        if current_index == target_ancestor_index {
            return Ok(true);
        }

        // Explore further up the graph via the parents of the current node
        for edge_ref in graph.parents_of(current_index) {
            let parent_of_current_index = edge_ref.source();

            if visited.insert(parent_of_current_index) {
                queue.push_back(parent_of_current_index);
            }
        }
    }

    // no indirect path exists
    Ok(false)
}

impl FetchStepData<MultiTypeFetchStep> {
    pub fn can_merge(
        &self,
        self_index: NodeIndex,
        other_index: NodeIndex,
        other: &Self,
        fetch_graph: &FetchGraph<MultiTypeFetchStep>,
    ) -> bool {
        if self_index == other_index {
            return false;
        }

        if self.service_name != other.service_name {
            return false;
        }

        // Is `this` FetchStep the only one `other` waits for?
        let is_only_parent = fetch_graph.parents_of(other_index).count() == 1
            && fetch_graph
                .parents_of(other_index)
                .all(|edge| edge.source() == self_index);

        // A step-level condition skips the whole request, so entity calls under different
        // conditions stay apart. A root fetch takes a conditional entity call by putting its
        // fields under `... on T @include(if: $x)`.
        //
        // An unconditional entity call can do the same with a nested entity call only it feeds:
        // `listings { ... on Listing @include(if: $x) { pet } }`.
        //
        // See `issue_1308_conditional_rank_keeps_concrete_requires_types` test
        //
        // Two entity calls for the same objects can share a request when one of them has no
        // condition, as that request goes out anyway: `... on User @include(if: $x) { name }`.
        //
        // See `client_alias_named_like_a_key_field_under_a_condition` test
        if self.is_entity_call() && other.is_entity_call() && self.condition != other.condition {
            let absorbs_conditional_child = self.condition.is_none()
                && is_only_parent
                && self.can_absorb_nested_entity_call(other);
            let one_call_for_same_objects = self.response_path == other.response_path
                && (self.condition.is_none() || other.condition.is_none());
            if !absorbs_conditional_child && !one_call_for_same_objects {
                return false;
            }
        }

        // If both are entities, their response_paths should match,
        // as we can't merge entity calls resolving different entities.
        // The one exception is a nested entity call that we feed ourselves,
        // see `can_absorb_nested_entity_call`.
        if matches!(self.kind, FetchStepKind::Entity) && self.kind == other.kind {
            if self.response_path.eq(&other.response_path) {
                // Same objects, but their fields sit under other types: an `Animal` call next to
                // a `Cat | Dog` one. Merged, the `Cat` fields would end up on `Animal`.
                // Outputs, not inputs: an `@interfaceObject` call reads `User`s and writes
                // `NodeWithName`, like the plain `NodeWithName` call next to it.
                if !self.output.selecting_same_types(&other.output) {
                    return false;
                }
            } else if !(is_only_parent && self.can_absorb_nested_entity_call(other)) {
                return false;
            }
        } else {
            // otherwise we can merge
            if !other.response_path.starts_with(&self.response_path) {
                return false;
            }
        }

        if self.has_arguments_conflicts_with(other) {
            return false;
        }

        if is_only_parent {
            return true;
        }

        // if they do not share parents, they can't be merged
        if !fetch_graph.parents_of(self_index).all(|self_edge| {
            fetch_graph
                .parents_of(other_index)
                .any(|other_edge| other_edge.source() == self_edge.source())
        }) {
            return false;
        }

        true
    }

    /// Only call this when we are the only step `other` waits for.
    /// If something else feeds it, the keys it sends may come from there,
    /// and pointing its other parents at us could change the order of the graph,
    /// or add a cycle.
    fn can_absorb_nested_entity_call(&self, other: &Self) -> bool {
        if !other.response_path.starts_with(&self.response_path) {
            return false;
        }

        let Some(input_type) = other.input.try_as_single() else {
            return false;
        };
        let Some(input_selections) = other.input.selections_for_definition(input_type) else {
            return false;
        };

        // The merge only knows where to put `other`'s fields when we fetch a single type.
        let Some(output_type) = self.output.try_as_single() else {
            return false;
        };

        // We must already have `other`'s input at that path.
        // If we do not, the fields we move here would have no object to sit in.
        let path = other.response_path.slice_from(self.response_path.len());
        self.output
            .selections_for_definition(output_type)
            .and_then(|output| find_selection_set_by_path(output, &path))
            .is_some_and(|at_path| at_path.contains(input_selections))
    }

    pub fn has_arguments_conflicts_with(&self, other: &Self) -> bool {
        let input_conflicts = FetchStepSelections::<MultiTypeFetchStep>::iter_matching_types(
            &self.input,
            &other.input,
            |_, self_selections, other_selections| {
                find_arguments_conflicts(self_selections, other_selections)
            },
        );

        input_conflicts
            .iter()
            .any(|(_, conflicts)| !conflicts.is_empty())
    }
}

#[cfg(test)]
mod tests {
    use graphql_tools::parser::query::{Definition, OperationDefinition};

    use crate::query_planner::{
        ast::{
            merge_path::{FieldPathSegment, MergePath, Segment},
            selection_set::SelectionSet,
        },
        planner::fetch::{
            error::FetchGraphError,
            fetch_graph::FetchGraph,
            fetch_step_data::{FetchStepData, FetchStepFlags, FetchStepKind},
            selections::FetchStepSelections,
            state::{MultiTypeFetchStep, SingleTypeFetchStep},
        },
        state::supergraph_state::{OperationKind, SubgraphName},
        utils::parsing::parse_operation,
    };

    use petgraph::graph::NodeIndex;

    use super::{perform_fetch_step_merge, MergedSteps};

    /// Selections for one or more types, e.g. `&[("User", "{ id }"), ("Admin", "{ id }")]`.
    fn selections(types: &[(&str, &str)]) -> FetchStepSelections<MultiTypeFetchStep> {
        let parse = |query: &str| -> SelectionSet {
            match parse_operation(query).definitions.first() {
                Some(Definition::Operation(OperationDefinition::SelectionSet(s))) => {
                    s.clone().into()
                }
                _ => panic!("expected a selection set"),
            }
        };

        let mut result = FetchStepSelections::<SingleTypeFetchStep>::new_empty().into_multi_type();
        for (type_name, query) in types {
            let mut single = FetchStepSelections::<SingleTypeFetchStep>::new(type_name);
            single.add(&parse(query)).unwrap();
            result.declare_known_type(type_name);
            result
                .migrate_from_another(&single.into_multi_type(), &MergePath::default())
                .unwrap();
        }
        result
    }

    fn path(fields: &[&str]) -> MergePath {
        MergePath::new(
            fields
                .iter()
                .map(|name| match *name {
                    "@" => Segment::List,
                    name => Segment::Field(FieldPathSegment::named(name.to_string()), 0, None),
                })
                .collect(),
        )
    }

    fn entity_step(
        service: &str,
        response_path: MergePath,
        input: FetchStepSelections<MultiTypeFetchStep>,
        output: FetchStepSelections<MultiTypeFetchStep>,
    ) -> FetchStepData<MultiTypeFetchStep> {
        FetchStepData {
            id: 0,
            service_name: SubgraphName(service.to_string()),
            response_path,
            input,
            output,
            kind: FetchStepKind::Entity,
            operation_kind: OperationKind::Query,
            flags: FetchStepFlags::empty(),
            condition: None,
            variable_usages: None,
            variable_definitions: None,
            mutation_field_position: None,
            input_rewrites: None,
            output_rewrites: None,
        }
    }

    /// A parent that fetches two types - what `batch_multi_type` builds - must not pull in a
    /// nested call. The merge only knows where to put the child's fields when the parent has a
    /// single type. Otherwise it looks for `Order` in the parent, and fails.
    #[test]
    fn multi_type_parent_does_not_absorb_nested_call() {
        let mut graph =
            FetchGraph::<SingleTypeFetchStep>::new(OperationKind::Query).to_multi_type();

        let parent = graph.add_step(entity_step(
            "orders",
            path(&["accounts", "@"]),
            selections(&[
                ("User", "{ __typename id }"),
                ("Admin", "{ __typename id }"),
            ]),
            selections(&[
                ("User", "{ orders { __typename id } }"),
                ("Admin", "{ orders { __typename id } }"),
            ]),
        ));
        let child = graph.add_step(entity_step(
            "orders",
            path(&["accounts", "@", "orders", "@"]),
            selections(&[("Order", "{ __typename id }")]),
            selections(&[("Order", "{ sku }")]),
        ));
        graph.connect(parent, child);

        let parent_data = graph.get_step_data(parent).unwrap();
        let child_data = graph.get_step_data(child).unwrap();
        assert!(!parent_data.can_merge(parent, child, child_data, &graph));
    }

    /// Two entity calls for the same type at unrelated paths can't be merged. The merged step
    /// would only fetch at `a.@`, and the entities at `b.@` would be lost.
    #[test]
    fn same_type_entity_calls_at_unrelated_paths_do_not_merge() {
        let mut graph =
            FetchGraph::<SingleTypeFetchStep>::new(OperationKind::Query).to_multi_type();

        let a = graph.add_step(entity_step(
            "catalog",
            path(&["a", "@"]),
            selections(&[("Order", "{ __typename id }")]),
            selections(&[("Order", "{ name }")]),
        ));
        let b = graph.add_step(entity_step(
            "catalog",
            path(&["b", "@"]),
            selections(&[("Order", "{ __typename id }")]),
            selections(&[("Order", "{ name }")]),
        ));

        let result = perform_fetch_step_merge(a, b, &mut graph, false);
        assert!(
            matches!(result, Err(FetchGraphError::MismatchedResponsePath)),
            "{result:?}"
        );
    }

    /// The `requires_requires` seed 4 order: 98 goes into 144, then 144 into 104. A pair that
    /// still names 98 has to land on 104, the one step left of the three.
    #[test]
    fn merged_steps_follow_the_whole_chain() {
        let [a, b, c, d] = [98, 144, 104, 3].map(NodeIndex::new);
        let mut merged = MergedSteps::default();
        merged.record(a, b);
        merged.record(b, c);

        assert_eq!(merged.resolve(a), c);
        assert_eq!(merged.resolve(b), c);
        assert_eq!(merged.resolve(c), c);
        assert_eq!(merged.resolve(d), d);
    }

    /// A child is a passthrough only when its input already has what it fetches for every
    /// type. Here it has the `User` fields but not the `Admin` ones.
    #[test]
    fn passthrough_needs_every_type_covered() {
        let mut graph =
            FetchGraph::<SingleTypeFetchStep>::new(OperationKind::Query).to_multi_type();

        let parent = graph.add_step(entity_step(
            "accounts",
            path(&["accounts", "@"]),
            selections(&[("User", "{ __typename id }")]),
            selections(&[("User", "{ id }")]),
        ));
        let child = graph.add_step(entity_step(
            "accounts",
            path(&["accounts", "@"]),
            selections(&[
                ("User", "{ __typename id }"),
                ("Admin", "{ __typename id }"),
            ]),
            selections(&[("User", "{ id }"), ("Admin", "{ name }")]),
        ));
        graph.connect(parent, child);

        let parent_data = graph.get_step_data(parent).unwrap();
        let child_data = graph.get_step_data(child).unwrap();
        assert!(!parent_data.can_merge_passthrough_child(parent, child, child_data, &graph));
    }
}