package com.arms.egovframework.javaservice.esframework.factory.builder;

import com.arms.egovframework.javaservice.esframework.annotation.Recent;
import com.arms.egovframework.javaservice.esframework.esquery.EsQuery;
import com.arms.egovframework.javaservice.esframework.model.dto.esquery.CompositeGroupDTO;
import com.arms.egovframework.javaservice.esframework.model.entity.BaseEntity;
import org.opensearch.data.client.orhlc.NativeSearchQuery;
import org.opensearch.data.client.orhlc.NativeSearchQueryBuilder;
import org.opensearch.index.query.BoolQueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.search.aggregations.AggregationBuilders;
import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder;
import org.opensearch.search.aggregations.metrics.TopHitsAggregationBuilder;
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.SortOrder;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

import static com.arms.egovframework.javaservice.esframework.util.ReflectionUtil.fieldInfo;

public class CompositeAggregationQueryBuilder implements AbstractQueryBuilder<NativeSearchQuery> {

    private final NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder();
    private final BoolQueryBuilder boolQuery;

    protected CompositeAggregationQueryBuilder(CompositeGroupDTO compositeGroupDTO, EsQuery esQuery, CompositeAggregationBuilder compositeAggregationBuilder) {

        String topHitsSortField = compositeGroupDTO.getSortField();
        String[] fieldIncludes = compositeGroupDTO.getIncludeFields();
        List<FieldSortBuilder> fieldSortBuilders = esQuery.getSortBuilders();

        this.boolQuery = esQuery.getBoolQueryBuilder();
        this.nativeSearchQueryBuilder.withMaxResults(compositeGroupDTO.isContentView() ? compositeGroupDTO.getSize() : 0);
        this.nativeSearchQueryBuilder.withTrackTotalHits(false);

        Optional.ofNullable(compositeGroupDTO.getAfterKey())
                .ifPresent(compositeAggregationBuilder::aggregateAfter);

        this.nativeSearchQueryBuilder.withAggregations(compositeAggregationBuilder);


        Optional.ofNullable(boolQuery)
                .ifPresent(query -> this.nativeSearchQueryBuilder.withQuery(boolQuery));

        Optional.ofNullable(fieldSortBuilders)
                .ifPresent(sorts ->
                        nativeSearchQueryBuilder.withSorts(new ArrayList<>(fieldSortBuilders)));

        if (topHitsSortField != null && fieldIncludes != null && fieldIncludes.length > 0) {
            int topHitsSize = compositeGroupDTO.getHitSize();
            SortOrder topHitsSortOrder = compositeGroupDTO.isSortAscending() ? SortOrder.ASC : SortOrder.DESC;

            TopHitsAggregationBuilder topHits = AggregationBuilders.topHits("top_hit")
                    .size(topHitsSize)
                    .fetchSource(false)
                    .sort(topHitsSortField, topHitsSortOrder);

            for (String field : fieldIncludes) {
                topHits.docValueField(field);
            }

            compositeAggregationBuilder.subAggregation(topHits);
        }
    }

    @Override
    public NativeSearchQuery create() {
        return nativeSearchQueryBuilder.build();
    }

    @Override
    public NativeSearchQuery createForRecentTrue(Class<? extends BaseEntity> entityClass) {
        String recentFieldName = fieldInfo(entityClass, Recent.class).getName();
        boolQuery.filter(QueryBuilders.termQuery(recentFieldName, true));
        return nativeSearchQueryBuilder.build();
    }
}
