qdrant-edge 0.7.1

use std::cmp;

use crate::common::counter::hardware_accumulator::HwMeasurementAcc;
use crate::segment::common::operation_error::OperationResult;
use crate::segment::data_types::modifier::Modifier;
use crate::segment::data_types::vectors::QueryVector;
use crate::segment::types::{DEFAULT_FULL_SCAN_THRESHOLD, ScoredPoint, WithPayload};
use crate::shard::common::stopping_guard::StoppingGuard;
use crate::shard::query::query_context::{fill_query_context, init_query_context};
use crate::shard::search::CoreSearchRequest;
use crate::shard::search_result_aggregator::BatchResultAggregator;

use crate::edge::{DEFAULT_EDGE_TIMEOUT, EdgeShard};

impl EdgeShard {
    /// This method is DEPRECATED and should be replaced with query.
    pub fn search(&self, search: CoreSearchRequest) -> OperationResult<Vec<ScoredPoint>> {
        let is_stopped_guard = StoppingGuard::new();
        let searches = [search];
        let query_context = init_query_context(
            &searches,
            DEFAULT_FULL_SCAN_THRESHOLD,
            &is_stopped_guard,
            HwMeasurementAcc::disposable_edge(),
            |vector_name| {
                self.config
                    .read()
                    .sparse_vectors
                    .get(vector_name)
                    .is_some_and(|v| v.modifier == Some(Modifier::Idf))
            },
        );
        let [search] = searches;
        let Some(context) = fill_query_context(
            query_context,
            self.segments.clone(),
            DEFAULT_EDGE_TIMEOUT,
            &is_stopped_guard.get_is_stopped(),
        )?
        else {
            // No segments to search
            return Ok(vec![]);
        };

        let segments: Vec<_> = self
            .segments
            .read()
            .non_appendable_then_appendable_segments()
            .collect();

        let CoreSearchRequest {
            query,
            filter,
            params,
            limit,
            offset,
            with_payload,
            with_vector,
            score_threshold,
        } = search;

        let vector_name = query.get_vector_name().to_string();
        let query_vector = QueryVector::from(query);
        let with_payload = WithPayload::from(with_payload.unwrap_or_default());
        let with_vector = with_vector.unwrap_or_default();

        let mut points_by_segment = Vec::with_capacity(segments.len());

        for segment in segments {
            let batched_points = segment.get().read().search_batch(
                &vector_name,
                &[&query_vector],
                &with_payload,
                &with_vector,
                filter.as_ref(),
                offset + limit,
                params.as_ref(),
                &context.get_segment_query_context(),
            )?;

            debug_assert_eq!(batched_points.len(), 1);

            let [points] = batched_points
                .try_into()
                .expect("single batched search result");

            points_by_segment.push(points);
        }

        let mut aggregator = BatchResultAggregator::new([offset + limit]);
        aggregator.update_point_versions(points_by_segment.iter().flatten());

        for points in points_by_segment {
            aggregator.update_batch_results(0, points);
        }

        let [mut points] = aggregator
            .into_topk()
            .try_into()
            .expect("single batched search result");

        let distance = {
            let config = self.config.read();
            if let Some(dense) = config.vectors.get(&vector_name) {
                dense.distance
            } else if config.sparse_vectors.contains_key(&vector_name) {
                crate::segment::types::Distance::Dot
            } else {
                return Err(
                    crate::segment::common::operation_error::OperationError::service_error(format!(
                        "vector config for '{vector_name}' does not exist"
                    )),
                );
            }
        };

        match &query_vector {
            QueryVector::Nearest(_) => {
                for point in &mut points {
                    point.score = distance.postprocess_score(point.score);
                }
            }
            QueryVector::RecommendBestScore(_) => (),
            QueryVector::RecommendSumScores(_) => (),
            QueryVector::Discover(_) => (),
            QueryVector::Context(_) => (),
            QueryVector::FeedbackNaive(_) => (),
        }

        if let Some(score_threshold) = score_threshold {
            debug_assert!(
                points.is_sorted_by(|left, right| distance.is_ordered(left.score, right.score)),
            );

            let below_threshold = points
                .iter()
                .enumerate()
                .find(|(_, point)| !distance.check_threshold(point.score, score_threshold));

            if let Some((below_threshold_idx, _)) = below_threshold {
                points.truncate(below_threshold_idx);
            }
        }

        let _ = points.drain(..cmp::min(points.len(), offset));

        Ok(points)
    }
}