use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::iter::Zip;
use std::rc::Rc;
use std::str::FromStr as _;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::vec::{IntoIter, Vec};

use actix_http::StatusCode;
use actix_web::web::Data;
use index_scheduler::filter::{
    filters_into_index_filters, filters_into_index_filters_unchecked,
    retrieve_foreign_keys_settings, SourceIndexUid,
};
use index_scheduler::{IndexScheduler, RoFeatures};
use itertools::Itertools;
use meilisearch_types::dynamic_search_rules::DynamicSearchRules;
use meilisearch_types::error::{Code, ResponseError};
use meilisearch_types::milli::order_by_map::OrderByMap;
use meilisearch_types::milli::progress::Progress;
use meilisearch_types::milli::score_details::{ScoreDetails, WeightedScoreValue};
use meilisearch_types::milli::vector::Embedding;
use meilisearch_types::milli::{
    self, merge_positioned_hits_into_page, serialize_index_filter_to_filter_string, Deadline,
    DocumentId, FederatingResultsStep, IndexFilter, OrderBy, DEFAULT_VALUES_PER_FACET,
};
use meilisearch_types::network::{Network, Remote, RemoteAvailability};
use meilisearch_types::settings::DEFAULT_PAGINATION_MAX_TOTAL_HITS;
use meilisearch_types::Document;
use roaring::RoaringBitmap;
use tokio::task::JoinHandle;
use uuid::Uuid;

use super::super::ranking_rules::{self, RankingRules};
use super::super::{
    compute_facet_distribution_stats, prepare_search, resolve_pins, AttributesFormat,
    ComputedFacets, HitMaker, HitsInfo, RetrieveVectors, SearchHit, SearchKind, SearchMetadata,
    SearchQuery, SearchQueryWithIndex,
};
use super::proxy::{proxy_search, ProxySearchError, ProxySearchParams};
use super::types::{
    FederatedFacets, FederatedSearchResult, Federation, FederationOptions, MergeFacets, Weight,
    FEDERATION_HIT, FEDERATION_REMOTE, PINNED_POSITION, WEIGHTED_SCORE_VALUES,
};
use super::weighted_scores;
use crate::error::MeilisearchHttpError;
use crate::routes::indexes::search::search_kind;
use crate::search::federated::types::{
    FEDERATION_EXTRA_DOCUMENT, INDEX_UID, QUERIES_POSITION, WEIGHTED_RANKING_SCORE,
};
use crate::search::hydration::{FederatedHydrationFormatter, HydrationContext};
use crate::search::{parse_filter, NetworkableQuery as _, DEFAULT_SEARCH_LIMIT};

#[allow(clippy::too_many_arguments)]
pub async fn perform_federated_search(
    index_scheduler: Data<IndexScheduler>,
    mut queries: Vec<SearchQueryWithIndex>,
    federation: Federation,
    features: RoFeatures,
    is_proxy: bool,
    request_uid: Uuid,
    include_metadata: bool,
    progress: &Progress,
) -> Result<(FederatedSearchResult, Deadline), ResponseError> {
    if is_proxy {
        features.check_network("Performing a remote federated search")?;
    }
    let before_search = std::time::Instant::now();

    let timeout = std::env::var("MEILI_EXPERIMENTAL_REMOTE_SEARCH_TIMEOUT_SECONDS")
        .ok()
        .map(|p| p.parse().unwrap())
        .unwrap_or(25);

    let deadline = before_search + std::time::Duration::from_secs(timeout);

    let required_hit_count = match (federation.page, federation.hits_per_page) {
        // no pagination, use limit and offset
        (None, None) => federation.limit + federation.offset,
        // pagination, default page (1)
        (None, Some(hits_per_page)) => hits_per_page,
        // pagination, no hits
        (Some(0), _) => 0,
        // pagination, default hits per page
        (Some(page), None) => DEFAULT_SEARCH_LIMIT() * page,
        // pagination
        (Some(page), Some(hits_per_page)) => hits_per_page * page,
    };

    let retrieve_vectors = queries.iter().any(|q| q.retrieve_vectors);

    let network = index_scheduler.network();
    let dynamic_search_rules = index_scheduler.dynamic_search_rules();

    // Preconstruct metadata keeping the original queries order for later metadata building
    let precomputed_query_metadata: Option<Vec<_>> = include_metadata.then(|| {
        queries
            .iter()
            .map(|q| {
                (
                    q.index_uid.to_string(),
                    q.federation_options.as_ref().and_then(|o| o.remote.clone()),
                )
            })
            .collect()
    });

    // Document join: list of indexes in the order of the queries
    // only create the hydration cache if the foreign keys feature is enabled
    let filter_values = queries.iter_mut().map(|q| q.filter.take()).collect::<Vec<_>>();
    let (mut hydration_cache, precomputed_filters) = if features.runtime_features().foreign_keys
        && !is_proxy
    {
        let index_uids: Vec<_> =
            queries.iter().map(|q| SourceIndexUid(Rc::from(q.index_uid.as_str()))).collect();
        let foreign_keys_settings = retrieve_foreign_keys_settings(&index_scheduler, &index_uids)?;

        // parse each query filter and bind them to their respective index
        let filters = index_uids
            .iter()
            .zip(filter_values.iter())
            .map(|(index_uid, filter)| match filter {
                Some(filter) => {
                    let filter = parse_filter(filter, Code::InvalidSearchFilter, features)?;

                    Ok((index_uid.clone(), filter))
                }
                None => Ok((index_uid.clone(), None)),
            })
            .collect::<Result<_, ResponseError>>()?;

        // convert the filters to index filters by evaluating the foreign filters
        let filters: Vec<_> = filters_into_index_filters(
            filters,
            &foreign_keys_settings,
            &index_scheduler,
            progress,
        )?;

        let hydration_cache = HydrationContext::new(index_uids, foreign_keys_settings);
        (Some(hydration_cache), filters)
    } else {
        let filters = filter_values
            .iter()
            .map(|f| {
                f.as_ref()
                    .and_then(|f| parse_filter(f, Code::InvalidSearchFilter, features).transpose())
                    .transpose()
            })
            .collect::<Result<_, ResponseError>>()?;
        let filters: Vec<_> = filters_into_index_filters_unchecked(filters)?;

        (None, filters)
    };

    // this implementation partition the queries by index to guarantee an important property:
    // - all the queries to a particular index use the same read transaction.
    // This is an important property, otherwise we cannot guarantee the self-consistency of the results.

    // 1. partition queries by host and index
    progress.update_progress(FederatingResultsStep::PartitionQueries);
    let mut partitioned_queries = PartitionedQueries::new();

    let mut federation = federation;
    for (query_index, (federated_query, filter)) in
        queries.into_iter().zip(precomputed_filters.into_iter()).enumerate()
    {
        partitioned_queries.partition(
            &mut federation,
            federated_query,
            filter,
            query_index,
            &network,
            features,
            index_scheduler.remote_availability(),
        )?
    }
    let federation = federation;

    // 2. perform queries, merge and make hits index by index
    // 2.1. start remote queries
    progress.update_progress(FederatingResultsStep::StartRemoteSearch);
    let remote_search = RemoteSearch::start(
        partitioned_queries.remote_queries_by_host,
        &federation,
        deadline,
        include_metadata,
        index_scheduler.web_client(),
    );

    // 2.2. concurrently execute local queries
    progress.update_progress(FederatingResultsStep::ExecuteLocalSearch);
    let params = SearchByIndexParams {
        index_scheduler,
        features,
        is_proxy,
        network,
        has_remote: partitioned_queries.has_remote,
        is_exhaustive: federation.is_exhaustive(),
        required_hit_count,
        dynamic_search_rules,
    };
    let mut search_by_index = SearchByIndex::new(
        federation,
        partitioned_queries.local_queries_by_index.len(),
        params.has_remote,
    );

    let mut deadline = Deadline::never();

    let (search_by_index, params, deadline) = tokio::task::spawn_blocking({
        let progress = progress.clone();
        move || -> Result<_, ResponseError> {
            for (index_uid, queries) in partitioned_queries.local_queries_by_index {
                // note: this is the only place we open `index_uid`
                let index_deadline =
                    search_by_index.execute(index_uid, queries, &params, &progress)?;
                deadline = Deadline::earliest(deadline, index_deadline);
            }

            // bonus step, make sure to return an error if an index wants a non-faceted field, even if no query actually uses that index.
            search_by_index.check_unused_facets(&params.index_scheduler)?;

            Ok((search_by_index, params, deadline))
        }
    })
    .await??;

    let SearchByIndexParams { network, index_scheduler, .. } = params;

    let SearchByIndex {
        federation,
        mut semantic_hit_count,
        mut results_by_index,
        mut query_vectors,
        previous_query_data: _,
        facet_order,
    } = search_by_index;

    progress.update_progress(FederatingResultsStep::WaitForRemoteResults);
    let before_waiting_remote_results = std::time::Instant::now();

    // 2.3. Wait for proxy search requests to complete
    let (mut remote_results, remote_errors) = remote_search.finish().await;

    let after_waiting_remote_results = std::time::Instant::now();

    // 3. merge hits and metadata across indexes and hosts
    progress.update_progress(FederatingResultsStep::MergeResults);
    // 3.1. Build metadata in the same order as the original queries
    let query_metadata = precomputed_query_metadata.map(|precomputed_query_metadata| {
        // If a remote is present, set the local remote name
        let local_remote_name = network.local.clone().filter(|_| partitioned_queries.has_remote);

        build_query_metadata(
            precomputed_query_metadata,
            local_remote_name,
            &remote_results,
            &results_by_index,
        )
    });

    // 3.2. merge federation metadata
    let (mut hit_number, degraded, used_negative_operator, facets, max_remote_duration) =
        merge_metadata(&mut results_by_index, &remote_results);

    let (skip, take) = match (federation.page, federation.hits_per_page) {
        // no pagination
        (None, None) => (federation.offset, federation.limit),
        // pagination: default page (1)
        (None, Some(hits_per_page)) => (0, hits_per_page),
        // special page 0: no hits
        (Some(0), _) => (0, 0),
        // pagination: default hits per page
        (Some(page), None) => ((page - 1) * DEFAULT_SEARCH_LIMIT(), DEFAULT_SEARCH_LIMIT()),
        (Some(page), Some(hits_per_page)) => ((page - 1) * hits_per_page, hits_per_page),
    };

    // 3.3. merge hits
    // Pinned documents are extracted before the score-based merge so they do
    // not compete with organic results. We keep an organic prefix and then
    // merge pins into the requested page in linear time.
    let mut pins = Vec::new();
    for result_by_index in &mut results_by_index {
        let prev_hits = std::mem::take(&mut result_by_index.hits);
        for hit in prev_hits {
            if let Some(ScoreDetails::Pin { position }) = hit.score.first() {
                pins.push((*position, hit.query_index, hit.hit));
            } else {
                result_by_index.hits.push(hit);
            }
        }
    }
    extract_remote_pin_hits(&mut remote_results, &mut pins);
    pins.sort_by_key(|&(pos, _, _)| pos);

    // store remote rejected hits to fixup the facet distributions
    // ideally could fixup in the iterator,
    // but we cannot double borrows of `remote_results` (from `hit` and `facets_by_index`).
    let mut rejected_hits: BTreeMap<String, Vec<SearchHit>> = Default::default();

    let mut distinct_values = HashSet::new();

    // When pins are present we need the organic prefix up to the end of the requested page,
    // then we inject pins into that prefix before applying the final slice. This mirrors
    // milli's bucket_sort behavior and naturally pumps late pins forward when organic
    // results run out.
    let (ranked_skip, ranked_take) =
        if pins.is_empty() { (skip, take) } else { (0, skip.saturating_add(take)) };
    let mut hit_it = merge_index_global_results(&mut results_by_index, &mut remote_results)
        .filter_map(|hit| {
            if let Some(distinct) = federation.distinct.as_deref() {
                let mut facet_values = Vec::new();
                hit.as_hit().facet_values(distinct, |value| facet_values.push(value));
                let is_rejected =
                    facet_values.iter().any(|facet_value| distinct_values.contains(facet_value));

                if is_rejected {
                    hit_number = hit_number.saturating_sub(1);

                    let index_uid = match hit.index_uid() {
                        Ok(index_uid) => index_uid,
                        Err(err) => {
                            tracing::warn!("skipping remote hit due to error: {err}");
                            return None;
                        }
                    };
                    rejected_hits.entry(index_uid.to_string()).or_default().push(hit.into_hit());

                    return None;
                }

                distinct_values.extend(facet_values.into_iter());
            }
            Some(hit)
        })
        .skip(ranked_skip);
    let mut merged_hits: Vec<_> = (&mut hit_it)
        .take(ranked_take)
        .inspect(|hit| {
            if let Some(semantic_hit_count) = &mut semantic_hit_count {
                if hit.to_score().0.any(|score| matches!(&score, WeightedScoreValue::VectorSort(_)))
                {
                    *semantic_hit_count += 1;
                }
            }
        })
        .map(|hit| (hit.query_index(), hit.into_hit()))
        .collect();

    if federation.distinct.is_some() {
        // when using distinct, exhaust any remaining hit from the merged list
        // so that the list of rejected hits is complete and the facet distribution is
        // fully adjusted by removing the facet values of all rejected hits.
        let _ = hit_it.count();
    } else {
        // since this variable is Drop and borrows local variables,
        // it needs to be dropped explicitly before going to the next step that use these
        // same variable.
        drop(hit_it);
    }

    merged_hits = merge_pinned_hits_into_page(pins, skip, take, merged_hits);

    // 3.3.1. hydrate documents based on the hydration points
    progress.update_progress(FederatingResultsStep::HydrateDocuments);
    if let Some(hydration_cache) = hydration_cache.as_mut() {
        for (query_index, hit) in &merged_hits {
            hydration_cache.register_foreign_docids(hit, *query_index);
        }
    }
    if let Some(hydration_cache) = hydration_cache {
        let hydration_formatter =
            FederatedHydrationFormatter::new(hydration_cache, &index_scheduler)?;
        hydration_formatter.hydrate_documents(&mut merged_hits)?;
    }

    let mut merged_hits = merged_hits.into_iter().map(|(_, hit)| hit).collect::<Vec<_>>();

    merged_hits.truncate(take);

    // 3.4. merge query vectors
    let query_vectors = if retrieve_vectors {
        for remote_results in remote_results.iter_mut() {
            if let Some(remote_vectors) = remote_results.query_vectors.take() {
                for (key, value) in remote_vectors.into_iter() {
                    debug_assert!(
                        !query_vectors.contains_key(&key),
                        "Query vector for query {key} already exists"
                    );
                    query_vectors.insert(key, value);
                }
            }
        }

        Some(query_vectors)
    } else {
        None
    };

    // 3.5. merge facets
    progress.update_progress(FederatingResultsStep::MergeFacets);
    let (facet_distribution, facet_stats, facets_by_index) =
        facet_order.merge(federation.merge_facets, remote_results, facets, rejected_hits);

    let after_merge = std::time::Instant::now();

    let local_duration = (before_waiting_remote_results - before_search)
        + (after_merge - after_waiting_remote_results);
    let max_duration = Duration::max(local_duration, max_remote_duration);

    let hits_info = match (federation.page, federation.hits_per_page) {
        // no pagination
        (None, None) => HitsInfo::OffsetLimit {
            limit: federation.limit,
            offset: federation.offset,
            estimated_total_hits: hit_number,
        },
        // pagination: default page number (1)
        (page, hits_per_page) => {
            let page = page.unwrap_or(1);
            let hits_per_page = hits_per_page.unwrap_or_else(DEFAULT_SEARCH_LIMIT);
            // If hit_per_page is 0, then pages can't be computed and so we respond 0.
            let total_pages = (hit_number + hits_per_page.saturating_sub(1))
                .checked_div(hits_per_page)
                .unwrap_or(0);
            HitsInfo::Pagination { hits_per_page, page, total_pages, total_hits: hit_number }
        }
    };

    let remote_errors = partitioned_queries.has_remote.then_some(remote_errors);
    let performance_details =
        federation.show_performance_details.then(|| progress.accumulated_durations());

    if !network.shards.is_empty() {
        for (remote_name, error) in remote_errors.iter().flatten() {
            if error.code.is_server_error() {
                index_scheduler.mark_remote_unavailable(remote_name.clone())?;
            }
        }
    }

    Ok((
        FederatedSearchResult {
            hits: merged_hits,
            processing_time_ms: max_duration.as_millis(),
            hits_info,
            query_vectors,
            semantic_hit_count,
            degraded,
            used_negative_operator,
            facet_distribution,
            facet_stats,
            facets_by_index,
            remote_errors,
            request_uid: Some(request_uid),
            metadata: query_metadata,
            performance_details,
        },
        deadline,
    ))
}

struct QueryByIndex {
    query: SearchQuery,
    filter: Option<IndexFilter<'static>>,
    weight: Weight,
    query_index: usize,
}

struct SearchResultByQuery<'a> {
    documents_ids: Vec<DocumentId>,
    document_scores: Vec<Vec<ScoreDetails>>,
    weight: Weight,
    hit_maker: HitMaker<'a>,
    query_index: usize,
}

struct SearchResultByQueryIter<'a> {
    it: Zip<IntoIter<DocumentId>, IntoIter<Vec<ScoreDetails>>>,
    weight: Weight,
    hit_maker: Rc<HitMaker<'a>>,
    query_index: usize,
}

impl<'a> SearchResultByQueryIter<'a> {
    fn new(
        SearchResultByQuery {
        documents_ids,
        document_scores,
        weight,
        hit_maker,
        query_index,
    }: SearchResultByQuery<'a>,
    ) -> Self {
        let it = documents_ids.into_iter().zip(document_scores);
        Self { it, weight, hit_maker: Rc::new(hit_maker), query_index }
    }
}

struct SearchResultByQueryIterItem<'a> {
    docid: DocumentId,
    score: Vec<ScoreDetails>,
    weight: Weight,
    hit_maker: Rc<HitMaker<'a>>,
    query_index: usize,
}

fn merge_index_local_results(
    results_by_query: Vec<SearchResultByQuery<'_>>,
) -> impl Iterator<Item = SearchResultByQueryIterItem<'_>> + '_ {
    itertools::kmerge_by(
        results_by_query.into_iter().map(SearchResultByQueryIter::new),
        |left: &SearchResultByQueryIterItem, right: &SearchResultByQueryIterItem| {
            match weighted_scores::compare(
                ScoreDetails::weighted_score_values(left.score.iter(), *left.weight),
                ScoreDetails::global_score(left.score.iter()) * *left.weight,
                ScoreDetails::weighted_score_values(right.score.iter(), *right.weight),
                ScoreDetails::global_score(right.score.iter()) * *right.weight,
            ) {
                // the biggest score goes first
                Ordering::Greater => true,
                // break ties using query index
                Ordering::Equal => left.query_index < right.query_index,
                Ordering::Less => false,
            }
        },
    )
}

// NOTE: Pinned documents (ScoreDetails::Pin) are extracted by the caller before invoking this
// function, so they never reach the score-based comparator. The caller re-injects pins at their
// target positions after this merge completes.
fn merge_index_global_results<'a>(
    results_by_index: &'a mut [SearchResultByIndex],
    remote_results: &'a mut [FederatedSearchResult],
) -> impl Iterator<Item = MergedSearchHit<'a>> + 'a {
    itertools::kmerge_by(
        // local results
        results_by_index
            .iter_mut()
            .map(|result_by_index| {
                either::Either::Left(std::mem::take(&mut result_by_index.hits).into_iter().map(
                    |hit| MergedSearchHit::Local { hit, index: result_by_index.index.as_str() },
                ))
            })
            // remote results
            .chain(remote_results.iter_mut().map(|x| either::Either::Right(iter_remote_hits(x)))),
        |left: &MergedSearchHit, right: &MergedSearchHit| {
            let (left_it, left_weighted_global_score, left_query_index) = left.to_score();
            let (right_it, right_weighted_global_score, right_query_index) = right.to_score();

            match weighted_scores::compare(
                left_it,
                left_weighted_global_score,
                right_it,
                right_weighted_global_score,
            ) {
                // the biggest score goes first
                Ordering::Greater => true,
                // break ties using query index
                Ordering::Equal => left_query_index < right_query_index,
                Ordering::Less => false,
            }
        },
    )
}

enum MergedSearchHit<'a> {
    Local { hit: SearchHitByIndex, index: &'a str },
    Remote(RemoteSearchHit),
}

struct RemoteSearchHit {
    hit: SearchHit,
    score: Vec<WeightedScoreValue>,
    global_weighted_score: f64,
    query_index: usize,
}

impl RemoteSearchHit {
    fn remote_index_uid(&self) -> Result<&str, ProxySearchError> {
        let hit = &self.hit;
        let federation = hit
            .document
            .get(FEDERATION_HIT)
            .ok_or(ProxySearchError::MissingPathInResponse("._federation"))?;
        let federation = match federation.as_object() {
            Some(federation) => federation,
            None => {
                return Err(ProxySearchError::UnexpectedValueInPath {
                    path: "._federation",
                    expected_type: "map",
                    received_value: federation.to_string(),
                });
            }
        };
        let index = federation
            .get(INDEX_UID)
            .ok_or(ProxySearchError::MissingPathInResponse("._federation.indexUid"))?;
        let index = index.as_str().ok_or_else(|| ProxySearchError::UnexpectedValueInPath {
            path: "._federation.indexUid",
            expected_type: "string",
            received_value: index.to_string(),
        })?;

        Ok(index)
    }
}

impl MergedSearchHit<'_> {
    fn remote(mut hit: SearchHit) -> Result<Self, ProxySearchError> {
        let federation = hit
            .document
            .get_mut(FEDERATION_HIT)
            .ok_or(ProxySearchError::MissingPathInResponse("._federation"))?;
        let federation = match federation.as_object_mut() {
            Some(federation) => federation,
            None => {
                return Err(ProxySearchError::UnexpectedValueInPath {
                    path: "._federation",
                    expected_type: "map",
                    received_value: federation.to_string(),
                });
            }
        };

        let global_weighted_score = federation
            .get(WEIGHTED_RANKING_SCORE)
            .ok_or(ProxySearchError::MissingPathInResponse("._federation.weightedRankingScore"))?;
        let global_weighted_score = global_weighted_score.as_f64().ok_or_else(|| {
            ProxySearchError::UnexpectedValueInPath {
                path: "._federation.weightedRankingScore",
                expected_type: "number",
                received_value: global_weighted_score.to_string(),
            }
        })?;

        let score: Vec<WeightedScoreValue> =
            serde_json::from_value(federation.remove(WEIGHTED_SCORE_VALUES).ok_or(
                ProxySearchError::MissingPathInResponse("._federation.weightedScoreValues"),
            )?)
            .map_err(ProxySearchError::CouldNotParseWeightedScoreValues)?;

        let query_index = federation
            .get(QUERIES_POSITION)
            .ok_or(ProxySearchError::MissingPathInResponse("._federation.queriesPosition"))?;
        let query_index =
            query_index.as_u64().ok_or_else(|| ProxySearchError::UnexpectedValueInPath {
                path: "._federation.queriesPosition",
                expected_type: "integer",
                received_value: query_index.to_string(),
            })? as usize;

        // Mount `extra_document` to the search hit, see documentation for `Hit::extra_document`
        let extra_document =
            federation.remove(FEDERATION_EXTRA_DOCUMENT).and_then(|extra_document| {
                if let serde_json::Value::Object(extra_document) = extra_document {
                    Some(extra_document)
                } else {
                    None
                }
            });

        hit.extra_document = extra_document.unwrap_or_default();

        Ok(Self::Remote(RemoteSearchHit { hit, score, global_weighted_score, query_index }))
    }

    fn into_hit(self) -> SearchHit {
        match self {
            MergedSearchHit::Local { hit: search_hit_by_index, .. } => search_hit_by_index.hit,
            MergedSearchHit::Remote(RemoteSearchHit { hit, .. }) => hit,
        }
    }

    fn as_hit(&self) -> &SearchHit {
        match self {
            MergedSearchHit::Local { hit: search_hit_by_index, .. } => &search_hit_by_index.hit,
            MergedSearchHit::Remote(RemoteSearchHit { hit, .. }) => hit,
        }
    }

    fn to_score(&self) -> (impl Iterator<Item = WeightedScoreValue> + '_, f64, usize) {
        match self {
            MergedSearchHit::Local { hit: search_hit_by_index, .. } => (
                either::Left(ScoreDetails::weighted_score_values(
                    search_hit_by_index.score.iter(),
                    *search_hit_by_index.weight,
                )),
                ScoreDetails::global_score(search_hit_by_index.score.iter())
                    * *search_hit_by_index.weight,
                search_hit_by_index.query_index,
            ),
            MergedSearchHit::Remote(RemoteSearchHit {
                score,
                global_weighted_score,
                query_index,
                hit: _,
            }) => {
                let global_weighted_score = *global_weighted_score;
                let query_index = *query_index;
                (either::Right(score.iter().cloned()), global_weighted_score, query_index)
            }
        }
    }

    fn index_uid(&self) -> Result<&str, ProxySearchError> {
        match self {
            MergedSearchHit::Local { hit: _, index } => Ok(index),
            MergedSearchHit::Remote(remote_search_hit) => remote_search_hit.remote_index_uid(),
        }
    }

    fn query_index(&self) -> usize {
        match self {
            MergedSearchHit::Local { hit: search_hit_by_index, .. } => {
                search_hit_by_index.query_index
            }
            MergedSearchHit::Remote(RemoteSearchHit { query_index, .. }) => *query_index,
        }
    }
}

fn iter_remote_hits(
    results_by_host: &mut FederatedSearchResult,
) -> impl Iterator<Item = MergedSearchHit<'_>> + '_ {
    // have a per node registry of failed hits
    results_by_host.hits.drain(..).filter_map(move |hit| match MergedSearchHit::remote(hit) {
        Ok(hit) => Some(hit),
        Err(err) => {
            tracing::warn!("skipping remote hit due to error: {err}");
            None
        }
    })
}

fn merge_pinned_hits_into_page<T>(
    pins: Vec<(u32, usize, T)>,
    skip: usize,
    take: usize,
    organic_hits: Vec<(usize, T)>,
) -> Vec<(usize, T)> {
    merge_positioned_hits_into_page(
        pins,
        skip,
        take,
        organic_hits,
        |&(position, _, _)| position,
        |(_, query_index, hit)| (query_index, hit),
    )
}

fn extract_remote_pin_hits(
    remote_results: &mut [FederatedSearchResult],
    pins: &mut Vec<(u32, usize, SearchHit)>,
) {
    fn parse_pin_pos_and_query_idx(
        federation: &mut serde_json::Map<String, serde_json::Value>,
    ) -> Option<(u32, usize)> {
        let pin_pos: u32 = federation.remove(PINNED_POSITION)?.as_u64()?.try_into().ok()?;
        let _ = federation.remove(WEIGHTED_SCORE_VALUES);
        let _ = federation.remove(FEDERATION_EXTRA_DOCUMENT);
        let query_idx: usize = federation.get(QUERIES_POSITION)?.as_u64()?.try_into().ok()?;

        Some((pin_pos, query_idx))
    }

    for remote_result in remote_results {
        let previous_hits = std::mem::take(&mut remote_result.hits);
        for mut hit in previous_hits {
            let pair = hit
                .document
                .get_mut(FEDERATION_HIT)
                .and_then(|federation| federation.as_object_mut())
                .and_then(parse_pin_pos_and_query_idx);

            if let Some((position, query_index)) = pair {
                pins.push((position, query_index, hit));
            } else {
                remote_result.hits.push(hit);
            }
        }
    }
}

fn build_federation_hit(
    params: &SearchByIndexParams,
    index_uid: &str,
    query_index: usize,
    score: &[ScoreDetails],
    weight: Weight,
    extra_document: &mut Document,
) -> serde_json::Value {
    let weighted_score = ScoreDetails::global_score(score.iter()) * *weight;
    let mut federation = serde_json::json!({
        INDEX_UID: index_uid,
        QUERIES_POSITION: query_index,
        WEIGHTED_RANKING_SCORE: weighted_score,
    });

    if params.has_remote && !params.is_proxy {
        federation
            .as_object_mut()
            .unwrap()
            .insert(FEDERATION_REMOTE.to_string(), params.network.local.clone().into());
    }

    if params.is_proxy {
        let federation = federation.as_object_mut().unwrap();
        federation.insert(
            WEIGHTED_SCORE_VALUES.to_string(),
            serde_json::json!(
                ScoreDetails::weighted_score_values(score.iter(), *weight).collect_vec()
            ),
        );
        federation.insert(
            FEDERATION_EXTRA_DOCUMENT.to_string(),
            serde_json::Value::Object(std::mem::take(extra_document)),
        );

        if let Some(ScoreDetails::Pin { position }) = score.first() {
            federation.insert(PINNED_POSITION.to_string(), serde_json::json!(position));
        }
    }

    federation
}

impl<'a> Iterator for SearchResultByQueryIter<'a> {
    type Item = SearchResultByQueryIterItem<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        let (docid, score) = self.it.next()?;
        Some(SearchResultByQueryIterItem {
            docid,
            score,
            weight: self.weight,
            hit_maker: Rc::clone(&self.hit_maker),
            query_index: self.query_index,
        })
    }
}

struct SearchHitByIndex {
    hit: SearchHit,
    score: Vec<ScoreDetails>,
    weight: Weight,
    query_index: usize,
}

struct SearchResultByIndex {
    index: String,
    primary_key: Option<String>,
    hits: Vec<SearchHitByIndex>,
    estimated_total_hits: usize,
    degraded: bool,
    used_negative_operator: bool,
    facets: Option<ComputedFacets>,
}

/// Builds query metadata for federated search results.
///
/// This function creates metadata for each query in the same order as the original queries,
/// combining information from both local and remote search results. It handles the mapping
/// of primary keys to their respective indexes and remotes to prevent collisions when
/// multiple remotes have the same index_uid but different primary keys.
fn build_query_metadata(
    precomputed_query_metadata: Vec<(String, Option<String>)>,
    local_remote_name: Option<String>,
    remote_results: &[FederatedSearchResult],
    results_by_index: &[SearchResultByIndex],
) -> Vec<SearchMetadata> {
    // Create a map of (remote, index_uid) -> primary_key for quick lookup
    // This prevents collisions when multiple remotes have the same index_uid but different primary keys
    let mut primary_key_per_index = std::collections::HashMap::new();

    // Build metadata for remote results
    for remote_result in remote_results {
        if let Some(remote_metadata) = &remote_result.metadata {
            for remote_meta in remote_metadata {
                if let SearchMetadata {
                    remote: Some(remote_name),
                    index_uid,
                    primary_key: Some(primary_key),
                    ..
                } = remote_meta
                {
                    let key = (Some(remote_name), index_uid);
                    primary_key_per_index.insert(key, primary_key);
                }
            }
        }
    }

    // Build metadata for local results
    for local_meta in results_by_index {
        if let SearchResultByIndex { index, primary_key: Some(primary_key), .. } = local_meta {
            let key = (None, index);
            primary_key_per_index.insert(key, primary_key);
        }
    }

    // Build metadata in the same order as the original queries
    let mut query_metadata = Vec::new();
    for (index_uid, remote) in precomputed_query_metadata {
        let primary_key =
            primary_key_per_index.get(&(remote.as_ref(), &index_uid)).map(|pk| pk.to_string());
        let query_uid = Uuid::now_v7();
        // if the remote is not set, use the local remote name
        let remote = remote.or_else(|| local_remote_name.clone());
        query_metadata.push(SearchMetadata { query_uid, primary_key, index_uid, remote });
    }
    query_metadata
}

fn merge_metadata(
    results_by_index: &mut Vec<SearchResultByIndex>,
    remote_results: &Vec<FederatedSearchResult>,
) -> (usize, bool, bool, FederatedFacets, Duration) {
    let mut estimated_total_hits = 0;
    let mut degraded = false;
    let mut used_negative_operator = false;
    let mut facets: FederatedFacets = FederatedFacets::default();
    let mut max_remote_duration = Duration::ZERO;
    for SearchResultByIndex {
        index,
        primary_key: _,
        hits: _,
        estimated_total_hits: estimated_total_hits_by_index,
        facets: facets_by_index,
        degraded: degraded_by_index,
        used_negative_operator: used_negative_operator_by_index,
    } in results_by_index
    {
        estimated_total_hits += *estimated_total_hits_by_index;
        degraded |= *degraded_by_index;
        used_negative_operator |= *used_negative_operator_by_index;

        let facets_by_index = std::mem::take(facets_by_index);

        facets.insert(index.clone(), facets_by_index);
    }
    for FederatedSearchResult {
        hits: _,
        processing_time_ms,
        hits_info,
        query_vectors: _,
        semantic_hit_count: _,
        facet_distribution: _,
        facet_stats: _,
        facets_by_index: _,
        degraded: degraded_for_host,
        used_negative_operator: host_used_negative_operator,
        remote_errors: _,
        metadata: _,
        request_uid: _,
        performance_details: _,
    } in remote_results
    {
        let this_remote_duration = Duration::from_millis(*processing_time_ms as u64);
        max_remote_duration = Duration::max(this_remote_duration, max_remote_duration);
        estimated_total_hits += match hits_info {
            HitsInfo::Pagination { total_hits: estimated_total_hits, .. }
            | HitsInfo::OffsetLimit { estimated_total_hits, .. } => estimated_total_hits,
        };
        // note that because `degraded` and `used_negative_operator` are #[serde(skip)],
        // `degraded_for_host` and `host_used_negative_operator` will always be false.
        degraded |= degraded_for_host;
        used_negative_operator |= host_used_negative_operator;
    }
    (estimated_total_hits, degraded, used_negative_operator, facets, max_remote_duration)
}

type LocalQueriesByIndex = BTreeMap<String, Vec<QueryByIndex>>;
type RemoteQueriesByHost = BTreeMap<String, (Remote, Vec<SearchQueryWithIndex>)>;

struct PartitionedQueries {
    local_queries_by_index: LocalQueriesByIndex,
    remote_queries_by_host: RemoteQueriesByHost,
    has_remote: bool,
}

impl PartitionedQueries {
    fn new() -> PartitionedQueries {
        PartitionedQueries {
            local_queries_by_index: Default::default(),
            remote_queries_by_host: Default::default(),
            has_remote: false,
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn partition(
        &mut self,
        federation: &mut Federation,
        mut federated_query: SearchQueryWithIndex,
        precomputed_filter: Option<IndexFilter>,
        query_index: usize,
        network: &Network,
        features: RoFeatures,
        remote_availability: &RemoteAvailability,
    ) -> Result<(), ResponseError> {
        if let Some(pagination_field) = federated_query.has_pagination() {
            return Err(MeilisearchHttpError::PaginationInFederatedQuery(
                query_index,
                pagination_field,
            )
            .into());
        }

        if let Some(facets) = federated_query.has_facets() {
            let facets = facets.to_owned();
            return Err(MeilisearchHttpError::FacetsInFederatedQuery(
                query_index,
                federated_query.index_uid.into_inner(),
                facets,
            )
            .into());
        }

        if federated_query.has_personalize() {
            return Err(MeilisearchHttpError::PersonalizationInFederatedQuery(query_index).into());
        }

        if federated_query.has_remote_and_use_network() {
            return Err(MeilisearchHttpError::RemoteAndUseNetwork(query_index).into());
        }

        if federated_query.has_show_performance_details() {
            return Err(
                MeilisearchHttpError::ShowPerformanceDetailsInFederatedQuery(query_index).into()
            );
        }

        if federated_query.has_distinct() && federation.distinct.is_some() {
            return Err(
                MeilisearchHttpError::DistinctInFederatedQueryAndFederation(query_index).into()
            );
        }

        let mut partition = None;

        let (index_uid, query, federation_options);
        let queries = if federated_query.must_use_network(network, &features)? {
            let partition = partition
                .get_or_insert_with(|| super::Partition::new(network.clone(), remote_availability));
            (index_uid, query, federation_options) = federated_query.into_index_query_federation();

            either::Left(partition.to_query_partition(
                federation,
                &query,
                federation_options,
                &index_uid,
            )?)
        } else {
            either::Right(std::iter::once(federated_query))
        };

        for federated_query in queries {
            let (index_uid, mut query, federation_options) =
                federated_query.into_index_query_federation();

            let federation_options = federation_options.unwrap_or_default();

            // local or remote node?
            'local_query: {
                let queries_by_index = match federation_options.remote {
                    None => self.local_queries_by_index.entry(index_uid.into_inner()).or_default(),
                    Some(remote_name) => {
                        self.has_remote = true;
                        features.check_network("Performing a remote federated search")?;

                        match &network.local {
                            Some(local) if local == &remote_name => self
                                .local_queries_by_index
                                .entry(index_uid.into_inner())
                                .or_default(),
                            _ => {
                                // node from the network
                                let Some(remote) = network.remotes.get(&remote_name) else {
                                    return Err(ResponseError::from_msg(
                                        format!(
                                            "Invalid `queries[{query_index}].federation_options.remote`: remote `{remote_name}` is not registered"
                                        ),
                                        meilisearch_types::error::Code::InvalidMultiSearchRemote,
                                    ));
                                };

                                // Insert back the filter into the query as a string before sending it to the remote
                                query.filter = precomputed_filter.as_ref().map(|f| {
                                    serde_json::Value::String(
                                        serialize_index_filter_to_filter_string(f).unwrap(),
                                    )
                                });

                                let query = SearchQueryWithIndex::from_index_query_federation(
                                    index_uid,
                                    query,
                                    Some(FederationOptions {
                                        weight: federation_options.weight,
                                        // do not pass the `remote` to not require the remote instance to have itself has a local node
                                        remote: None,
                                        // pass an explicit query index
                                        query_position: Some(query_index),
                                    }),
                                );

                                self.remote_queries_by_host
                                    .entry(remote_name)
                                    .or_insert_with(|| (remote.clone(), Default::default()))
                                    .1
                                    .push(query);
                                break 'local_query;
                            }
                        }
                    }
                };

                queries_by_index.push(QueryByIndex {
                    query,
                    filter: precomputed_filter.as_ref().map(|f| f.clone().into_owned()),
                    weight: federation_options.weight,
                    // override query index here with the one in federation.
                    // this will fix-up error messages to refer to the global query index of the original request.
                    query_index: if let Some(query_index) = federation_options.query_position {
                        features.check_network("Using `federationOptions.queryPosition`")?;
                        query_index
                    } else {
                        query_index
                    },
                })
            }
        }
        Ok(())
    }
}

struct RemoteSearch {
    in_flight_remote_queries:
        BTreeMap<String, JoinHandle<Result<FederatedSearchResult, ProxySearchError>>>,
}

impl RemoteSearch {
    fn start(
        queries: RemoteQueriesByHost,
        federation: &Federation,
        deadline: Instant,
        include_metadata: bool,
        client: &http_client::reqwest::Client,
    ) -> Self {
        let mut in_flight_remote_queries = BTreeMap::new();

        if queries.is_empty() {
            return Self { in_flight_remote_queries };
        }

        let params =
            ProxySearchParams { deadline: Some(deadline), try_count: 3, client: client.clone() };
        for (node_name, (node, queries)) in queries {
            // spawn one task per host
            in_flight_remote_queries.insert(
                node_name,
                tokio::spawn({
                    let mut proxy_federation = federation.clone();
                    // fixup limit and offset to not apply them twice
                    match (federation.page, federation.hits_per_page) {
                        (None, None) => {
                            proxy_federation.limit = federation.limit + federation.offset;
                            proxy_federation.offset = 0;
                        }
                        (Some(0), _) => {
                            // leave `proxy_federation.page` at Some(0)
                        }
                        (page, hits_per_page) => {
                            let page = page.unwrap_or(1);
                            let hits_per_page = hits_per_page.unwrap_or_else(DEFAULT_SEARCH_LIMIT);
                            proxy_federation.page = Some(1);
                            proxy_federation.hits_per_page = Some(page * hits_per_page);
                        }
                    }

                    // never merge distant facets
                    proxy_federation.merge_facets = None;
                    let params = params.clone();
                    async move {
                        proxy_search(&node, queries, proxy_federation, &params, include_metadata)
                            .await
                    }
                }),
            );
        }
        Self { in_flight_remote_queries }
    }

    async fn finish(self) -> (Vec<FederatedSearchResult>, BTreeMap<String, ResponseError>) {
        let mut remote_results = Vec::with_capacity(self.in_flight_remote_queries.len());
        let mut remote_errors: BTreeMap<String, ResponseError> = BTreeMap::new();
        'remote_queries: for (node_name, handle) in self.in_flight_remote_queries {
            match handle.await {
                Ok(Ok(mut res)) => {
                    for hit in &mut res.hits {
                        let Some(federation) = hit.document.get_mut(FEDERATION_HIT) else {
                            let error = ProxySearchError::MissingPathInResponse("._federation");
                            remote_errors.insert(node_name, error.as_response_error());
                            continue 'remote_queries;
                        };
                        let Some(federation) = federation.as_object_mut() else {
                            let error = ProxySearchError::UnexpectedValueInPath {
                                path: "._federation",
                                expected_type: "map",
                                received_value: federation.to_string(),
                            };
                            remote_errors.insert(node_name, error.as_response_error());
                            continue 'remote_queries;
                        };
                        if !federation.contains_key(WEIGHTED_SCORE_VALUES) {
                            let error = ProxySearchError::MissingPathInResponse(
                                "._federation.weightedScoreValues",
                            );
                            remote_errors.insert(node_name, error.as_response_error());
                            continue 'remote_queries;
                        }

                        if !federation.contains_key(WEIGHTED_RANKING_SCORE) {
                            let error = ProxySearchError::MissingPathInResponse(
                                "._federation.weightedRankingScore",
                            );
                            remote_errors.insert(node_name, error.as_response_error());
                            continue 'remote_queries;
                        }

                        // Add remote name to metadata
                        if let Some(metadata) = res.metadata.as_mut() {
                            for meta in metadata {
                                meta.remote = Some(node_name.clone());
                            }
                        }

                        federation.insert(
                            FEDERATION_REMOTE.to_string(),
                            serde_json::Value::String(node_name.clone()),
                        );
                    }

                    remote_results.push(res);
                }
                Ok(Err(error)) => {
                    remote_errors.insert(node_name, error.as_response_error());
                }
                Err(panic) => match panic.try_into_panic() {
                    Ok(panic) => {
                        let msg = match panic.downcast_ref::<&'static str>() {
                            Some(s) => *s,
                            None => match panic.downcast_ref::<String>() {
                                Some(s) => &s[..],
                                None => "Box<dyn Any>",
                            },
                        };
                        remote_errors.insert(
                            node_name,
                            ResponseError::from_msg(
                                msg.to_string(),
                                meilisearch_types::error::Code::Internal,
                            ),
                        );
                    }
                    Err(_) => tracing::error!("proxy search task was unexpectedly cancelled"),
                },
            }
        }
        (remote_results, remote_errors)
    }
}

struct SearchByIndexParams {
    index_scheduler: Data<IndexScheduler>,
    required_hit_count: usize,
    is_exhaustive: bool,
    features: RoFeatures,
    is_proxy: bool,
    has_remote: bool,
    network: Network,
    dynamic_search_rules: Arc<DynamicSearchRules>,
}

struct SearchByIndex {
    federation: Federation,
    // During search by index, semantic_hit_count will be set to Some(0) if any search kind uses semantic
    // Then when merging, we'll update its value if there is any semantic hit
    semantic_hit_count: Option<u32>,
    results_by_index: Vec<SearchResultByIndex>,
    query_vectors: BTreeMap<usize, Embedding>,
    previous_query_data: Option<(RankingRules, usize, String)>,
    // remember the order and name of first index for each facet when merging with index settings
    // to detect if the order is inconsistent for a facet.
    facet_order: FacetOrder,
}

impl SearchByIndex {
    fn new(federation: Federation, index_count: usize, has_remote: bool) -> Self {
        SearchByIndex {
            facet_order: match (federation.merge_facets, has_remote) {
                (None, true) => FacetOrder::ByIndex(Default::default()),
                (None, false) => FacetOrder::None,
                (Some(_), _) => FacetOrder::ByFacet(Default::default()),
            },
            federation,
            semantic_hit_count: None,
            results_by_index: Vec::with_capacity(index_count),
            query_vectors: BTreeMap::new(),
            previous_query_data: None,
        }
    }

    fn execute(
        &mut self,
        index_uid: String,
        queries: Vec<QueryByIndex>,
        params: &SearchByIndexParams,
        progress: &Progress,
    ) -> Result<Deadline, ResponseError> {
        let first_query_index = queries.first().map(|query| query.query_index);
        let index = match params.index_scheduler.index(&index_uid) {
            Ok(index) => index,
            Err(err) => {
                let mut err = ResponseError::from(err);
                // Patch the HTTP status code to 400 as it defaults to 404 for `index_not_found`, but
                // here the resource not found is not part of the URL.
                err.code = StatusCode::BAD_REQUEST;
                if let Some(query_index) = first_query_index {
                    err.message = format!("Inside `.queries[{}]`: {}", query_index, err.message);
                }
                return Err(err);
            }
        };
        let rtxn = index.read_txn()?;
        let primary_key = index.primary_key(&rtxn)?.map(|pk| pk.to_string());
        let criteria = index.criteria(&rtxn)?;
        let dictionary = index.dictionary(&rtxn)?;
        let dictionary: Option<Vec<_>> =
            dictionary.as_ref().map(|x| x.iter().map(String::as_str).collect());
        let separators = index.allowed_separators(&rtxn)?;
        let separators: Option<Vec<_>> =
            separators.as_ref().map(|x| x.iter().map(String::as_str).collect());

        let max_total_hits = index
            .pagination_max_total_hits(&rtxn)
            .map_err(milli::Error::from)?
            .map(|x| x as usize)
            .unwrap_or(DEFAULT_PAGINATION_MAX_TOTAL_HITS);

        let required_hit_count = usize::min(params.required_hit_count, max_total_hits);

        let mut degraded = false;
        let mut used_negative_operator = false;
        let mut candidates = RoaringBitmap::new();
        let facets_by_index = self.federation.facets_by_index.remove(&index_uid).flatten();
        if let Err(mut error) =
            self.facet_order.check_facet_order(&index_uid, &facets_by_index, &index, &rtxn)
        {
            error.message = format!(
                "Inside `.federation.facetsByIndex.{index_uid}`: {error}{}",
                if let Some(query_index) = first_query_index {
                    format!("\n - Note: index `{index_uid}` used in `.queries[{query_index}]`")
                } else {
                    Default::default()
                }
            );
            return Err(error);
        }
        let mut results_by_query = Vec::with_capacity(queries.len());

        // all queries for an index share the same deadline
        let deadline = index.search_deadline(&rtxn)?;

        let extra_attributes_to_retrieve: BTreeSet<_> =
            if let Some(distinct) = self.federation.distinct.as_deref() {
                std::iter::once(distinct.to_string())
                    .chain(facets_by_index.iter().flatten().cloned())
                    .collect()
            } else {
                Default::default()
            };

        for QueryByIndex { query, weight, query_index, filter } in queries {
            // use an immediately invoked lambda to capture the result without returning from the function
            let res: Result<(), ResponseError> = (|| {
                let search_kind =
                    search_kind(&query, &params.index_scheduler, index_uid.to_string(), &index)?;

                let canonicalization_kind = match (&search_kind, &query.q) {
                    (SearchKind::SemanticOnly { .. }, _) => {
                        ranking_rules::CanonicalizationKind::Vector
                    }
                    (_, Some(q)) if !q.is_empty() => ranking_rules::CanonicalizationKind::Keyword,
                    _ => ranking_rules::CanonicalizationKind::Placeholder,
                };

                let sort = if let Some(sort) = &query.sort {
                    let sorts: Vec<_> =
                        match sort.iter().map(|s| milli::AscDesc::from_str(s)).collect() {
                            Ok(sorts) => sorts,
                            Err(asc_desc_error) => {
                                return Err(milli::SortError::from(asc_desc_error)
                                    .into_search_error()
                                    .into());
                            }
                        };
                    Some(sorts)
                } else {
                    None
                };

                let ranking_rules = ranking_rules::RankingRules::new(
                    criteria.clone(),
                    sort,
                    query.matching_strategy.into(),
                    canonicalization_kind,
                );

                if let Some((previous_ranking_rules, previous_query_index, previous_index_uid)) =
                    self.previous_query_data.take()
                {
                    if let Err(error) = ranking_rules.is_compatible_with(&previous_ranking_rules) {
                        return Err(error.to_response_error(
                            &ranking_rules,
                            &previous_ranking_rules,
                            query_index,
                            previous_query_index,
                            &index_uid,
                            &previous_index_uid,
                        ));
                    }
                    self.previous_query_data = if previous_ranking_rules.constraint_count()
                        > ranking_rules.constraint_count()
                    {
                        Some((previous_ranking_rules, previous_query_index, previous_index_uid))
                    } else {
                        Some((ranking_rules, query_index, index_uid.clone()))
                    };
                } else {
                    self.previous_query_data =
                        Some((ranking_rules, query_index, index_uid.clone()));
                }

                match search_kind {
                    SearchKind::KeywordOnly => {}
                    _ => self.semantic_hit_count = Some(0),
                }

                let retrieve_vectors = RetrieveVectors::new(query.retrieve_vectors);

                let (mut search, _is_finite_pagination, _max_total_hits, _offset) = prepare_search(
                    &index,
                    &rtxn,
                    &query,
                    filter,
                    &search_kind,
                    // clones of `Deadline` share the deadline rather than restart it
                    deadline.clone(),
                    params.features,
                    progress,
                )?;

                search.scoring_strategy(milli::score_details::ScoringStrategy::Detailed);

                search.offset(0);
                search.limit(required_hit_count);
                search.exhaustive_number_hits(params.is_exhaustive);
                let pins = if params.features.runtime_features().dynamic_search_rules {
                    resolve_pins(&params.dynamic_search_rules, &query, &index_uid, &index, &rtxn)?
                } else {
                    Vec::new()
                };
                if !pins.is_empty() {
                    search.pins(pins);
                }

                if let Some(distinct) = self.federation.distinct.as_deref() {
                    search.distinct(distinct.to_owned());
                }

                let (result, _semantic_hit_count) =
                    super::super::search_from_kind(index_uid.to_string(), search_kind, search)?;
                let format = AttributesFormat {
                    attributes_to_retrieve: query.attributes_to_retrieve,
                    extra_attributes_to_retrieve: extra_attributes_to_retrieve.clone(),
                    retrieve_vectors,
                    attributes_to_highlight: query.attributes_to_highlight,
                    attributes_to_crop: query.attributes_to_crop,
                    crop_length: query.crop_length,
                    crop_marker: query.crop_marker,
                    highlight_pre_tag: query.highlight_pre_tag,
                    highlight_post_tag: query.highlight_post_tag,
                    show_matches_position: query.show_matches_position,
                    sort: query.sort,
                    show_ranking_score: query.show_ranking_score,
                    show_ranking_score_details: query.show_ranking_score_details,
                    locales: query.locales.map(|l| l.iter().copied().map(Into::into).collect()),
                };

                let milli::SearchResult {
                    matching_words,
                    candidates: query_candidates,
                    documents_ids,
                    document_scores,
                    degraded: query_degraded,
                    used_negative_operator: query_used_negative_operator,
                    query_vector,
                } = result;

                if query.retrieve_vectors {
                    if let Some(query_vector) = query_vector {
                        debug_assert!(
                            !self.query_vectors.contains_key(&query_index),
                            "Query vector for query {query_index} already exists"
                        );
                        self.query_vectors.insert(query_index, query_vector);
                    }
                }

                candidates |= query_candidates;
                degraded |= query_degraded;
                used_negative_operator |= query_used_negative_operator;

                let tokenizer = HitMaker::tokenizer(dictionary.as_deref(), separators.as_deref());

                let formatter_builder = HitMaker::formatter_builder(matching_words, tokenizer);

                let hit_maker =
                    HitMaker::new(&index, &rtxn, format, formatter_builder).map_err(|e| {
                        MeilisearchHttpError::from_milli(e, Some(index_uid.to_string()))
                    })?;

                results_by_query.push(SearchResultByQuery {
                    weight,
                    hit_maker,
                    query_index,
                    documents_ids,
                    document_scores,
                });
                Ok(())
            })();

            if let Err(mut error) = res {
                error.message = format!("Inside `.queries[{query_index}]`: {}", error.message);
                return Err(error);
            }
        }
        let mut documents_seen = RoaringBitmap::new();
        let mut local_pinned_hits = Vec::new();
        for result_by_query in &mut results_by_query {
            let prev_documents_ids = std::mem::take(&mut result_by_query.documents_ids);
            let prev_scores = std::mem::take(&mut result_by_query.document_scores);

            for (doc_id, score) in prev_documents_ids.into_iter().zip(prev_scores.into_iter()) {
                if let Some(ScoreDetails::Pin { position }) = score.first() {
                    let mut hit = result_by_query.hit_maker.make_hit(doc_id, &score, progress)?;
                    let _federation = build_federation_hit(
                        params,
                        &index_uid,
                        result_by_query.query_index,
                        &score,
                        result_by_query.weight,
                        &mut hit.extra_document,
                    );

                    hit.document.insert(FEDERATION_HIT.to_string(), _federation);
                    local_pinned_hits.push((
                        *position,
                        result_by_query.query_index,
                        SearchHitByIndex {
                            hit,
                            score,
                            weight: result_by_query.weight,
                            query_index: result_by_query.query_index,
                        },
                    ));
                } else {
                    result_by_query.documents_ids.push(doc_id);
                    result_by_query.document_scores.push(score);
                }
            }
        }

        local_pinned_hits.sort_by_key(|&(pos, _, _)| pos);

        // A set of the seen values for the facet.
        // Whenever we consider a document, we check that its value for the distinct fid has not already been seen.
        // If it was seen, it is rejected, which shouldn't happen "too often" as the intermediate lists of results were
        // already dedup'd.
        // If it wasn't seen, it is accepted and we update the list of seen values accordingly.
        let mut distinct_values = HashSet::new();

        let organic_hits = merge_index_local_results(results_by_query)
            // skip documents we've already seen & mark that we saw the current document
            // 2.3 make hits
            .filter_map(
                |SearchResultByQueryIterItem { docid, hit_maker, score, weight, query_index }| {
                    let already_seen = !documents_seen.insert(docid);
                    if already_seen {
                        return None;
                    }

                    let hit = (|| {
                        let mut hit = hit_maker.make_hit(docid, &score, progress)?;

                        if let Some(distinct) = self.federation.distinct.as_deref() {
                            let mut facet_values = Vec::new();
                            hit.facet_values(distinct, |value| facet_values.push(value));
                            let is_rejected = facet_values
                                .iter()
                                .any(|facet_value| distinct_values.contains(facet_value));

                            if is_rejected {
                                candidates.remove(docid);

                                return Ok(None);
                            }

                            distinct_values.extend(facet_values);
                        }

                        let _federation = build_federation_hit(
                            params,
                            &index_uid,
                            query_index,
                            &score,
                            weight,
                            &mut hit.extra_document,
                        );

                        hit.document.insert(FEDERATION_HIT.to_string(), _federation);
                        Ok(Some(SearchHitByIndex { hit, score, weight, query_index }))
                    })();
                    hit.transpose()
                },
            )
            .take(required_hit_count)
            .map_ok(|hit| (hit.query_index, hit))
            .collect::<Result<Vec<_>, ResponseError>>()?;
        let merged_result =
            merge_pinned_hits_into_page(local_pinned_hits, 0, required_hit_count, organic_hits)
                .into_iter()
                .map(|(_, hit)| hit)
                .collect();
        let estimated_total_hits = candidates.len() as usize;
        let facets = facets_by_index
            .map(|facets_by_index| {
                compute_facet_distribution_stats(
                    &facets_by_index,
                    &index,
                    &rtxn,
                    candidates,
                    super::super::Route::MultiSearch,
                )
            })
            .transpose()
            .map_err(|mut error| {
                error.message = format!(
                    "Inside `.federation.facetsByIndex.{index_uid}`: {}{}",
                    error.message,
                    if let Some(query_index) = first_query_index {
                        format!("\n - Note: index `{index_uid}` used in `.queries[{query_index}]`")
                    } else {
                        Default::default()
                    }
                );
                error
            })?;

        self.results_by_index.push(SearchResultByIndex {
            index: index_uid,
            primary_key,
            hits: merged_result,
            estimated_total_hits,
            degraded,
            used_negative_operator,
            facets,
        });
        Ok(deadline)
    }

    fn check_unused_facets(
        &mut self,
        index_scheduler: &IndexScheduler,
    ) -> Result<(), ResponseError> {
        for (index_uid, facets) in std::mem::take(&mut self.federation.facets_by_index) {
            let index = match index_scheduler.index(&index_uid) {
                Ok(index) => index,
                Err(err) => {
                    let mut err = ResponseError::from(err);
                    // Patch the HTTP status code to 400 as it defaults to 404 for `index_not_found`, but
                    // here the resource not found is not part of the URL.
                    err.code = StatusCode::BAD_REQUEST;
                    err.message = format!(
                        "Inside `.federation.facetsByIndex.{index_uid}`: {}\n - Note: index `{index_uid}` is not used in queries",
                        err.message
                    );
                    return Err(err);
                }
            };

            // Important: this is the only transaction we'll use for this index during this federated search
            let rtxn = index.read_txn()?;

            if let Err(mut error) =
                self.facet_order.check_facet_order(&index_uid, &facets, &index, &rtxn)
            {
                error.message = format!(
                    "Inside `.federation.facetsByIndex.{index_uid}`: {error}\n - Note: index `{index_uid}` is not used in queries",
                );
                return Err(error);
            }

            if let Some(facets) = facets {
                if let Err(mut error) = compute_facet_distribution_stats(
                    &facets,
                    &index,
                    &rtxn,
                    Default::default(),
                    super::super::Route::MultiSearch,
                ) {
                    error.message = format!(
                        "Inside `.federation.facetsByIndex.{index_uid}`: {}\n - Note: index `{index_uid}` is not used in queries",
                        error.message
                    );
                    return Err(error);
                }
            }
        }
        Ok(())
    }
}

enum FacetOrder {
    /// The order is stored by facet to be able to merge facets regardless of index of origin
    ///
    /// - key: facet name
    /// - value: (first_index_name, first_index_order)
    ///
    /// We store the name of the first index where the facet is present as well as its order,
    /// so that if encountering the same facet in a different index we can compare the order and send
    /// a readable error.
    ByFacet(BTreeMap<String, (String, OrderBy)>),
    /// The order is stored by index to be able to merge facets regardless of the remote of origin.
    ///
    /// This variant is only used when `is_remote = true`, and always used in that case.
    ///
    /// - key: index name
    /// - value: (order_by_map, max_values_per_facet)
    ///
    /// We store a map of the order per facet for that index, as well as the max values per facet.
    /// Both are retrieved from the settings of the local version of the index.
    ///
    /// It is not possible to have an index only existing in the remotes, because as of now all indexes that appear
    /// in `federation.facetsByIndex` must exist on all hosts.
    ByIndex(BTreeMap<String, (OrderByMap, usize)>),
    /// Do not merge facets. Used when `federation.mergeFacets = null` and `!has_remote`
    None,
}

type FacetDistributions = BTreeMap<String, indexmap::IndexMap<String, u64>>;
type FacetStats = BTreeMap<String, crate::search::FacetStats>;

impl FacetOrder {
    fn check_facet_order(
        &mut self,
        current_index: &str,
        facets_by_index: &Option<Vec<String>>,
        index: &milli::Index,
        rtxn: &milli::heed::RoTxn<'_>,
    ) -> Result<(), ResponseError> {
        match self {
            FacetOrder::ByFacet(facet_order) => {
                if let Some(facets_by_index) = facets_by_index {
                    let index_facet_order = index.sort_facet_values_by(rtxn)?;
                    for facet in facets_by_index {
                        let index_facet_order = index_facet_order.get(facet);
                        let (previous_index, previous_facet_order) = facet_order
                            .entry(facet.to_owned())
                            .or_insert_with(|| (current_index.to_owned(), index_facet_order));
                        if previous_facet_order != &index_facet_order {
                            return Err(MeilisearchHttpError::InconsistentFacetOrder {
                                facet: facet.clone(),
                                previous_facet_order: *previous_facet_order,
                                previous_uid: previous_index.clone(),
                                current_uid: current_index.to_owned(),
                                index_facet_order,
                            }
                            .into());
                        }
                    }
                }
            }
            FacetOrder::ByIndex(order_by_index) => {
                let max_values_per_facet = index
                    .max_values_per_facet(rtxn)?
                    .map(|x| x as usize)
                    .unwrap_or(DEFAULT_VALUES_PER_FACET);
                order_by_index.insert(
                    current_index.to_owned(),
                    (index.sort_facet_values_by(rtxn)?, max_values_per_facet),
                );
            }
            FacetOrder::None => {}
        }
        Ok(())
    }

    fn merge(
        self,
        merge_facets: Option<MergeFacets>,
        remote_results: Vec<FederatedSearchResult>,
        mut facets: FederatedFacets,
        rejected_hits: BTreeMap<String, Vec<SearchHit>>,
    ) -> (Option<FacetDistributions>, Option<FacetStats>, FederatedFacets) {
        let (facet_distribution, facet_stats, facets_by_index) = match (self, merge_facets) {
            (FacetOrder::ByFacet(facet_order), Some(merge_facets)) => {
                for remote_facets_by_index in
                    remote_results.into_iter().map(|result| result.facets_by_index)
                {
                    facets.append(remote_facets_by_index);
                }
                let mut facets = facets.merge(merge_facets, facet_order);

                if let Some(facets) = &mut facets {
                    let rejected_hits =
                        rejected_hits.into_values().fold(Vec::new(), |mut init, mut v| {
                            init.append(&mut v);
                            init
                        });
                    facets.remove_hits(&rejected_hits);
                }

                let (facet_distribution, facet_stats) = facets
                    .map(|ComputedFacets { distribution, stats }| (distribution, stats))
                    .unzip();

                (facet_distribution, facet_stats, FederatedFacets::default())
            }
            (FacetOrder::ByIndex(facet_order), _) => {
                for remote_facets_by_index in
                    remote_results.into_iter().map(|result| result.facets_by_index)
                {
                    facets.append(remote_facets_by_index);
                }
                facets.sort_and_truncate(facet_order);

                for (index, facets) in &mut facets.0 {
                    let Some(rejected_hits) = rejected_hits.get(index) else {
                        continue;
                    };
                    facets.remove_hits(rejected_hits);
                }
                (None, None, facets)
            }
            _ => {
                for (index, facets) in &mut facets.0 {
                    let Some(rejected_hits) = rejected_hits.get(index) else {
                        continue;
                    };
                    facets.remove_hits(rejected_hits);
                }
                (None, None, facets)
            }
        };
        (facet_distribution, facet_stats, facets_by_index)
    }
}
