use std::future::Future;
use std::sync::Arc;
use std::time::Duration;

use ahash::AHashMap;
use api::rest::{BaseGroupRequest, SearchGroupsRequestInternal, SearchRequestInternal};
use common::counter::hardware_accumulator::HwMeasurementAcc;
use segment::json_path::JsonPath;
use segment::types::WithVector;
use shard::grouping::{GroupByDriver, RequestBudget};

use super::types::QueryGroupRequest;
use crate::collection::Collection;
use crate::common::fetch_vectors;
use crate::common::fetch_vectors::build_vector_resolver_query;
use crate::lookup::WithLookup;
use crate::operations::consistency_params::ReadConsistency;
use crate::operations::routing::RoutingToken;
use crate::operations::shard_selector_internal::ShardSelectorInternal;
use crate::operations::types::{
    CollectionResult, PointGroup, RecommendGroupsRequestInternal, RecommendRequestInternal,
};
use crate::operations::universal_query::collection_query::{
    CollectionQueryGroupsRequest, CollectionQueryRequest,
};
use crate::operations::universal_query::shard_query::{self, ShardQueryRequest};
use crate::recommendations::recommend_into_core_search;

#[derive(Clone, Debug, PartialEq)]
pub enum SourceRequest {
    Search(SearchRequestInternal),
    Recommend(RecommendRequestInternal),
    Query(CollectionQueryRequest),
}

#[derive(Clone, Debug, PartialEq)]
pub struct GroupRequest {
    /// Request to use (search or recommend)
    pub source: SourceRequest,

    /// Path to the field to group by
    pub group_by: JsonPath,

    /// Limit of points to return per group
    pub group_size: usize,

    /// Limit of groups to return
    pub limit: usize,

    /// Options for specifying how to use the group id to lookup points in another collection
    pub with_lookup: Option<WithLookup>,
}

impl GroupRequest {
    pub fn with_limit_from_request(
        source: SourceRequest,
        group_by: JsonPath,
        group_size: usize,
    ) -> Self {
        let limit = match &source {
            SourceRequest::Search(request) => request.limit,
            SourceRequest::Recommend(request) => request.limit,
            SourceRequest::Query(request) => request.limit,
        };
        Self {
            source,
            group_by,
            group_size,
            limit,
            with_lookup: None,
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn into_query_group_request<F, Fut>(
        self,
        collection: &Collection,
        collection_by_name: F,
        read_consistency: Option<ReadConsistency>,
        routing_token: Option<RoutingToken>,
        shard_selection: ShardSelectorInternal,
        timeout: Option<Duration>,
        hw_measurement_acc: HwMeasurementAcc,
    ) -> CollectionResult<QueryGroupRequest>
    where
        F: Fn(String) -> Fut,
        Fut: Future<Output = Option<Arc<Collection>>>,
    {
        let query_search = match self.source {
            SourceRequest::Search(search_req) => ShardQueryRequest::from(search_req),
            SourceRequest::Recommend(recommend_req) => {
                let referenced_vectors = fetch_vectors::resolve_referenced_vectors_batch(
                    &[(recommend_req.clone(), shard_selection)],
                    collection,
                    collection_by_name,
                    read_consistency,
                    routing_token,
                    timeout,
                    hw_measurement_acc.clone(),
                )
                .await?;

                let core_search =
                    recommend_into_core_search(&collection.id, recommend_req, &referenced_vectors)?;
                ShardQueryRequest::from(core_search)
            }
            SourceRequest::Query(query_req) => {
                // Lift nested prefetches to root queries for vector resolution
                let resolver_requests = build_vector_resolver_query(&query_req, &shard_selection);

                let referenced_vectors = fetch_vectors::resolve_referenced_vectors_batch(
                    &resolver_requests,
                    collection,
                    collection_by_name,
                    read_consistency,
                    routing_token,
                    timeout,
                    hw_measurement_acc.clone(),
                )
                .await?;
                query_req.try_into_shard_request(&collection.id, &referenced_vectors)?
            }
        };

        Ok(QueryGroupRequest {
            source: query_search,
            group_by: self.group_by,
            group_size: self.group_size,
            groups: self.limit,
        })
    }
}

impl From<SearchGroupsRequestInternal> for GroupRequest {
    fn from(request: SearchGroupsRequestInternal) -> Self {
        let SearchGroupsRequestInternal {
            vector,
            filter,
            params,
            with_payload,
            with_vector,
            score_threshold,
            group_request:
                BaseGroupRequest {
                    group_by,
                    group_size,
                    limit,
                    with_lookup: with_lookup_interface,
                },
        } = request;

        let search = SearchRequestInternal {
            vector,
            filter,
            params,
            limit: 0,
            offset: Some(0),
            with_payload,
            with_vector,
            score_threshold,
        };

        GroupRequest {
            source: SourceRequest::Search(search),
            group_by,
            group_size: group_size as usize,
            limit: limit as usize,
            with_lookup: with_lookup_interface.map(Into::into),
        }
    }
}

impl From<RecommendGroupsRequestInternal> for GroupRequest {
    fn from(request: RecommendGroupsRequestInternal) -> Self {
        let RecommendGroupsRequestInternal {
            positive,
            negative,
            strategy,
            filter,
            params,
            with_payload,
            with_vector,
            score_threshold,
            using,
            lookup_from,
            group_request:
                BaseGroupRequest {
                    group_by,
                    group_size,
                    limit,
                    with_lookup: with_lookup_interface,
                },
        } = request;

        let recommend = RecommendRequestInternal {
            positive,
            negative,
            strategy,
            filter,
            params,
            limit: 0,
            offset: None,
            with_payload,
            with_vector,
            score_threshold,
            using,
            lookup_from,
        };

        GroupRequest {
            source: SourceRequest::Recommend(recommend),
            group_by,
            group_size: group_size as usize,
            limit: limit as usize,
            with_lookup: with_lookup_interface.map(Into::into),
        }
    }
}

impl From<CollectionQueryGroupsRequest> for GroupRequest {
    fn from(request: CollectionQueryGroupsRequest) -> Self {
        let CollectionQueryGroupsRequest {
            prefetch,
            query,
            using,
            filter,
            params,
            score_threshold,
            with_vector,
            with_payload,
            lookup_from,
            group_by,
            group_size,
            limit,
            with_lookup: with_lookup_interface,
        } = request;

        let collection_query_request = CollectionQueryRequest {
            prefetch: prefetch.into_iter().collect(),
            query,
            using,
            filter,
            score_threshold,
            limit,
            offset: 0,
            params,
            with_vector,
            with_payload,
            lookup_from,
        };

        GroupRequest {
            source: SourceRequest::Query(collection_query_request),
            group_by,
            group_size,
            limit,
            with_lookup: with_lookup_interface,
        }
    }
}

/// Uses the request to fill up groups of points.
pub async fn group_by(
    request: QueryGroupRequest,
    collection: &Collection,
    read_consistency: Option<ReadConsistency>,
    routing_token: Option<RoutingToken>,
    shard_selection: ShardSelectorInternal,
    timeout: Option<Duration>,
    hw_measurement_acc: HwMeasurementAcc,
) -> CollectionResult<Vec<PointGroup>> {
    let start = std::time::Instant::now();
    let collection_params = collection.collection_config.read().await.params.clone();
    let score_ordering =
        shard_query::query_result_order(request.source.query.as_ref(), &collection_params)?;

    let QueryGroupRequest {
        mut source,
        group_by,
        group_size,
        groups,
    } = request;

    // Groups are enriched with the user-requested payload and vectors at the end,
    // so candidates are fetched bare.
    let with_payload = source.with_payload.clone();
    let with_vector = source.with_vector.clone();
    source.with_vector = WithVector::Bool(false);

    let mut driver = GroupByDriver::new(
        source,
        group_by,
        groups,
        group_size,
        score_ordering,
        RequestBudget::default(),
    );

    while let Some(query) = driver.next_request() {
        // update timeout
        let timeout = timeout.map(|t| t.saturating_sub(start.elapsed()));

        let points = collection
            .query(
                query,
                read_consistency,
                routing_token,
                shard_selection.clone(),
                timeout,
                hw_measurement_acc.clone(),
            )
            .await?;

        driver.add_points(&points);
    }

    // extract best results
    let mut groups = driver.distill();

    // flatten results
    let bare_points = groups
        .iter()
        .cloned()
        .flat_map(|group| group.hits)
        .collect();

    // update timeout
    let timeout = timeout.map(|t| t.saturating_sub(start.elapsed()));

    // enrich with payload and vector
    let enriched_points: AHashMap<_, _> = collection
        .fill_search_result_with_payload(
            bare_points,
            Some(with_payload),
            with_vector,
            read_consistency,
            routing_token,
            &shard_selection,
            timeout,
            hw_measurement_acc.clone(),
        )
        .await?
        .into_iter()
        .map(|point| (point.id, point))
        .collect();

    // hydrate groups with enriched points
    groups
        .iter_mut()
        .for_each(|group| group.hydrate_from(&enriched_points));

    // turn into output form
    let groups = groups.into_iter().map(PointGroup::from).collect();

    Ok(groups)
}

#[cfg(test)]
mod tests {
    use ahash::AHashMap;
    use segment::data_types::groups::GroupId;
    use segment::payload_json;
    use segment::types::{Payload, ScoredPoint};
    use shard::grouping::Group;

    fn make_scored_point(id: u64, score: f32, payload: Option<Payload>) -> ScoredPoint {
        ScoredPoint {
            id: id.into(),
            version: 0,
            score,
            payload,
            vector: None,
            shard_key: None,
            order_value: None,
        }
    }

    #[test]
    fn test_hydrated_from() {
        // arrange
        let mut groups: Vec<Group> = Vec::new();
        [
            (
                "a",
                [
                    make_scored_point(1, 1.0, None),
                    make_scored_point(2, 1.0, None),
                ],
            ),
            (
                "b",
                [
                    make_scored_point(3, 1.0, None),
                    make_scored_point(4, 1.0, None),
                ],
            ),
        ]
        .into_iter()
        .for_each(|(key, points)| {
            let group = Group {
                key: GroupId::from(key),
                hits: points.into_iter().collect(),
            };
            groups.push(group);
        });

        let payload_a = payload_json! {"some_key": "some value a"};
        let payload_b = payload_json! {"some_key": "some value b"};

        let hydrated = vec![
            make_scored_point(1, 1.0, Some(payload_a.clone())),
            make_scored_point(2, 1.0, Some(payload_a.clone())),
            make_scored_point(3, 1.0, Some(payload_b.clone())),
            make_scored_point(4, 1.0, Some(payload_b.clone())),
        ];

        let set: AHashMap<_, _> = hydrated.into_iter().map(|p| (p.id, p)).collect();

        // act
        groups.iter_mut().for_each(|group| group.hydrate_from(&set));

        // assert
        assert_eq!(groups.len(), 2);
        assert_eq!(groups.first().unwrap().hits.len(), 2);
        assert_eq!(groups.get(1).unwrap().hits.len(), 2);

        let a = groups.first().unwrap();
        let b = groups.get(1).unwrap();

        assert!(
            a.hits
                .iter()
                .all(|x| x.payload.as_ref() == Some(&payload_a)),
        );
        assert!(
            b.hits
                .iter()
                .all(|x| x.payload.as_ref() == Some(&payload_b)),
        );
    }
}
