package com.arms.egovframework.javaservice.esframework.custom.client;

/*
 * Copyright OpenSearch Contributors
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

import static org.springframework.util.CollectionUtils.*;

import java.time.Duration;
import java.util.EnumSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchType;
import org.opensearch.common.geo.GeoDistance;
import org.opensearch.common.unit.DistanceUnit;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.data.client.orhlc.*;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.search.aggregations.AbstractAggregationBuilder;
import org.opensearch.search.aggregations.PipelineAggregationBuilder;
import org.opensearch.search.builder.PointInTimeBuilder;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.subphase.FetchSourceContext;
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.opensearch.search.rescore.QueryRescoreMode;
import org.opensearch.search.rescore.QueryRescorerBuilder;
import org.opensearch.search.sort.FieldSortBuilder;
import org.opensearch.search.sort.GeoDistanceSortBuilder;
import org.opensearch.search.sort.ScoreSortBuilder;
import org.opensearch.search.sort.SortBuilder;
import org.opensearch.search.sort.SortBuilders;
import org.opensearch.search.sort.SortMode;
import org.opensearch.search.sort.SortOrder;
import org.springframework.data.domain.Sort;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.data.elasticsearch.core.query.RescorerQuery.ScoreMode;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

/**
 * Factory class to create OpenSearch request instances from Spring Data OpenSearch query objects.
 * @since 0.1
 */
class CustomRequestFactory {

    private final ElasticsearchConverter elasticsearchConverter;

    public CustomRequestFactory(ElasticsearchConverter elasticsearchConverter) {
        this.elasticsearchConverter = elasticsearchConverter;
    }

    // region search
    @Nullable
    public HighlightBuilder highlightBuilder(Query query) {
        HighlightBuilder highlightBuilder = query.getHighlightQuery()
                .map(highlightQuery -> new HighlightQueryBuilder(elasticsearchConverter.getMappingContext())
                        .getHighlightBuilder(highlightQuery.getHighlight(), highlightQuery.getType()))
                .orElse(null);

        if (highlightBuilder == null) {

            if (query instanceof NativeSearchQuery) {
                NativeSearchQuery searchQuery = (NativeSearchQuery) query;
                HighlightBuilder.Field[] highlightFields = searchQuery.getHighlightFields();

                if ((highlightFields != null && highlightFields.length > 0)
                        || searchQuery.getHighlightBuilder() != null) {
                    highlightBuilder = searchQuery.getHighlightBuilder();

                    if (highlightBuilder == null) {
                        highlightBuilder = new HighlightBuilder();
                    }

                    if (highlightFields != null) {
                        for (HighlightBuilder.Field highlightField : highlightFields) {
                            highlightBuilder.field(highlightField);
                        }
                    }
                }
            }
        }
        return highlightBuilder;
    }

    public SearchRequest searchRequest(
            Query query, @Nullable String routing, @Nullable Class<?> clazz, IndexCoordinates index) {

        elasticsearchConverter.updateQuery(query, clazz);
        SearchRequest searchRequest = prepareSearchRequest(query, routing, clazz, index);
        QueryBuilder opensearchQuery = getQuery(query);
        QueryBuilder opensearchFilter = getFilter(query);

        searchRequest.source().query(opensearchQuery);

        if (opensearchFilter != null) {
            searchRequest.source().postFilter(opensearchFilter);
        }

        return searchRequest;
    }

    private SearchRequest prepareSearchRequest(
            Query query, @Nullable String routing, @Nullable Class<?> clazz, IndexCoordinates indexCoordinates) {
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        SearchRequest request;
        // Point in time requests do not allow index definition as PIT id references specific indices
        if(query.getPointInTime()== null) {
            String[] indexNames = indexCoordinates.getIndexNames();
            Assert.notNull(indexNames, "No index defined for Query");
            Assert.notEmpty(indexNames, "No index defined for Query");
            request = new SearchRequest(indexNames);
        } else {
            request = new SearchRequest();
            PointInTimeBuilder pointInTimeBuilder = new PointInTimeBuilder(query.getPointInTime().id());
            pointInTimeBuilder.setKeepAlive(TimeValue.timeValueMillis(query.getPointInTime().keepAlive().toMillis()));
            sourceBuilder.pointInTimeBuilder(pointInTimeBuilder);
        }

        sourceBuilder.version(true);
        sourceBuilder.trackScores(query.getTrackScores());
        if (hasSeqNoPrimaryTermProperty(clazz)) {
            sourceBuilder.seqNoAndPrimaryTerm(true);
        }

        sourceBuilder
                .from((int) (query.getPageable().isPaged() ? query.getPageable().getOffset() : 0))
                .size(query.getRequestSize());

        SourceFilter sourceFilter = query.getSourceFilter();
        if (sourceFilter != null) {
            sourceBuilder.fetchSource(getFetchSourceContext(query));
            sourceBuilder.fetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
        }

        if(!isEmpty(query.getDocValueFields())) {
            List<DocValueField> docValueFields = query.getDocValueFields();
            docValueFields.forEach(field -> sourceBuilder.docValueField(field.field()));
        }

        if (!query.getFields().isEmpty()) {
            query.getFields().forEach(sourceBuilder::fetchField);
        }

        if (!isEmpty(query.getStoredFields())) {
            sourceBuilder.storedFields(query.getStoredFields());
        }

        if (query.getIndicesOptions() != null) {
            request.indicesOptions(toOpenSearchIndicesOptions(query.getIndicesOptions()));
        }

        if (query.isLimiting()) {
            // noinspection ConstantConditions
            sourceBuilder.size(query.getMaxResults());
        }

        if (query.getMinScore() > 0) {
            sourceBuilder.minScore(query.getMinScore());
        }

        if (query.getPreference() != null) {
            request.preference(query.getPreference());
        }

        request.searchType(SearchType.fromString(query.getSearchType().name().toLowerCase()));

        prepareSort(query, sourceBuilder, getPersistentEntity(clazz));

        HighlightBuilder highlightBuilder = highlightBuilder(query);

        if (highlightBuilder != null) {
            sourceBuilder.highlighter(highlightBuilder);
        }

        if (query instanceof NativeSearchQuery) {
            prepareNativeSearch((NativeSearchQuery) query, sourceBuilder);
        }

        if (query.getTrackTotalHits() != null) {
            sourceBuilder.trackTotalHits(query.getTrackTotalHits());
        } else if (query.getTrackTotalHitsUpTo() != null) {
            sourceBuilder.trackTotalHitsUpTo(query.getTrackTotalHitsUpTo());
        }

        if (StringUtils.hasLength(query.getRoute())) {
            request.routing(query.getRoute());
        } else if (StringUtils.hasText(routing)) {
            request.routing(routing);
        }

        Duration timeout = query.getTimeout();
        if (timeout != null) {
            sourceBuilder.timeout(new TimeValue(timeout.toMillis()));
        }

        sourceBuilder.explain(query.getExplain());

        List<Object> searchAfter = query.getSearchAfter();
        if (searchAfter != null) {
            sourceBuilder.searchAfter(searchAfter.toArray());
        }

        query.getRescorerQueries().forEach(rescorer -> sourceBuilder.addRescorer(getQueryRescorerBuilder(rescorer)));

        if (query.getRequestCache() != null) {
            request.requestCache(query.getRequestCache());
        }

        Duration scrollTime = query.getScrollTime();
        if (scrollTime != null) {
            request.scroll(TimeValue.timeValueMillis(scrollTime.toMillis()));
        }

        request.source(sourceBuilder);
        return request;
    }


    private void prepareNativeSearch(NativeSearchQuery query, SearchSourceBuilder sourceBuilder) {

        if (!query.getScriptFields().isEmpty()) {
            for (ScriptField scriptedField : query.getScriptFields()) {
                sourceBuilder.scriptField(scriptedField.fieldName(), scriptedField.script());
            }
        }

        if (query.getCollapseBuilder() != null) {
            sourceBuilder.collapse(query.getCollapseBuilder());
        }

        if (!isEmpty(query.getIndicesBoost())) {
            for (IndexBoost indexBoost : query.getIndicesBoost()) {
                sourceBuilder.indexBoost(indexBoost.getIndexName(), indexBoost.getBoost());
            }
        }

        List<AbstractAggregationBuilder<?>> aggregations = query.getAggregations();
        if (!isEmpty(aggregations)) {
            aggregations.forEach(sourceBuilder::aggregation);
        }

        List<PipelineAggregationBuilder> pipelineAggregations = query.getPipelineAggregations();
        if (!isEmpty(pipelineAggregations)) {
            pipelineAggregations.forEach(sourceBuilder::aggregation);
        }

        if (query.getSuggestBuilder() != null) {
            sourceBuilder.suggest(query.getSuggestBuilder());
        }

        if (!isEmpty(query.getSearchExtBuilders())) {
            sourceBuilder.ext(query.getSearchExtBuilders());
        }
    }

    @SuppressWarnings("rawtypes")
    private void prepareSort(
        Query query, SearchSourceBuilder sourceBuilder, @Nullable ElasticsearchPersistentEntity<?> entity) {

        Sort sort = query.getSort();
        if (sort != null) {
            sort.forEach(order -> sourceBuilder.sort(getSortBuilder(order, entity)));
        }

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery nativeSearchQuery = (NativeSearchQuery) query;
            List<SortBuilder<?>> sorts = nativeSearchQuery.getOpenSearchSorts();
            if (sorts != null) {
                sorts.forEach(sourceBuilder::sort);
            }
        }
    }

    private SortBuilder<?> getSortBuilder(Sort.Order order, @Nullable ElasticsearchPersistentEntity<?> entity) {
        SortOrder sortOrder = order.getDirection().isDescending() ? SortOrder.DESC : SortOrder.ASC;

        Order.Mode mode = null;
        String unmappedType = null;

        if (order instanceof Order) {
            Order o = (Order) order;
            mode = o.getMode();
            unmappedType = o.getUnmappedType();
        }

        if (mode == null) {
            mode = Order.Mode.min;
        }

        if (ScoreSortBuilder.NAME.equals(order.getProperty())) {
            return SortBuilders //
                    .scoreSort() //
                    .order(sortOrder);
        } else {
            ElasticsearchPersistentProperty property = (entity != null) //
                    ? entity.getPersistentProperty(order.getProperty()) //
                    : null;
            String fieldName = property != null ? property.getFieldName() : order.getProperty();

            if (order instanceof GeoDistanceOrder) {
                GeoDistanceOrder geoDistanceOrder = (GeoDistanceOrder) order;

                GeoDistanceSortBuilder sort = SortBuilders.geoDistanceSort(
                        fieldName,
                        geoDistanceOrder.getGeoPoint().getLat(),
                        geoDistanceOrder.getGeoPoint().getLon());

                sort.geoDistance(GeoDistance.fromString(
                        geoDistanceOrder.getDistanceType().name()));
                sort.sortMode(SortMode.fromString(mode.name()));
                sort.unit(DistanceUnit.fromString(geoDistanceOrder.getUnit()));

                if (!GeoDistanceOrder.DEFAULT_IGNORE_UNMAPPED.equals(geoDistanceOrder.getIgnoreUnmapped())) {
                    sort.ignoreUnmapped(geoDistanceOrder.getIgnoreUnmapped());
                }

                return sort;
            } else {
                FieldSortBuilder sort = SortBuilders //
                        .fieldSort(fieldName) //
                        .order(sortOrder) //
                        .sortMode(SortMode.fromString(mode.name()));

                if (unmappedType != null) {
                    sort.unmappedType(unmappedType);
                }

                if (order.getNullHandling() == Sort.NullHandling.NULLS_FIRST) {
                    sort.missing("_first");
                } else if (order.getNullHandling() == Sort.NullHandling.NULLS_LAST) {
                    sort.missing("_last");
                }
                return sort;
            }
        }
    }

    private QueryRescorerBuilder getQueryRescorerBuilder(RescorerQuery rescorerQuery) {

        QueryBuilder queryBuilder = getQuery(rescorerQuery.getQuery());
        Assert.notNull("queryBuilder", "Could not build query for rescorerQuery");

        QueryRescorerBuilder builder = new QueryRescorerBuilder(queryBuilder);

        if (rescorerQuery.getScoreMode() != ScoreMode.Default) {
            builder.setScoreMode(
                    QueryRescoreMode.valueOf(rescorerQuery.getScoreMode().name()));
        }

        if (rescorerQuery.getQueryWeight() != null) {
            builder.setQueryWeight(rescorerQuery.getQueryWeight());
        }

        if (rescorerQuery.getRescoreQueryWeight() != null) {
            builder.setRescoreQueryWeight(rescorerQuery.getRescoreQueryWeight());
        }

        if (rescorerQuery.getWindowSize() != null) {
            builder.windowSize(rescorerQuery.getWindowSize());
        }

        return builder;
    }
    // endregion

    @Nullable
    private QueryBuilder getQuery(Query query) {
        QueryBuilder opensearchQuery;

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery searchQuery = (NativeSearchQuery) query;
            opensearchQuery = searchQuery.getQuery();
        }  else {
            throw new IllegalArgumentException(
                    "unhandled Query implementation " + query.getClass().getName());
        }

        return opensearchQuery;
    }

    @Nullable
    private QueryBuilder getFilter(Query query) {
        QueryBuilder opensearchFilter;

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery searchQuery = (NativeSearchQuery) query;
            opensearchFilter = searchQuery.getFilter();
        } else {
            throw new IllegalArgumentException(
                    "unhandled Query implementation " + query.getClass().getName());
        }

        return opensearchFilter;
    }

    @Nullable
    private FetchSourceContext getFetchSourceContext(Query searchQuery) {

        SourceFilter sourceFilter = searchQuery.getSourceFilter();

        if (sourceFilter != null) {
            Boolean fetchSource = sourceFilter.fetchSource();
            if (fetchSource != null && !fetchSource) {
                return new FetchSourceContext(false);
            }
            return new FetchSourceContext(true, sourceFilter.getIncludes(), sourceFilter.getExcludes());
        }

        return null;
    }

    public org.opensearch.action.support.IndicesOptions toOpenSearchIndicesOptions(IndicesOptions indicesOptions) {

        Assert.notNull(indicesOptions, "indicesOptions must not be null");

        Set<org.opensearch.action.support.IndicesOptions.Option> options = indicesOptions.getOptions().stream()
                .map(it -> org.opensearch.action.support.IndicesOptions.Option.valueOf(
                        it.name().toUpperCase()))
                .collect(Collectors.toSet());

        Set<org.opensearch.action.support.IndicesOptions.WildcardStates> wildcardStates =
                indicesOptions.getExpandWildcards().stream()
                        .map(it -> org.opensearch.action.support.IndicesOptions.WildcardStates.valueOf(
                                it.name().toUpperCase()))
                        .collect(Collectors.toSet());

        return new org.opensearch.action.support.IndicesOptions(
                options.isEmpty()
                        ? EnumSet.noneOf(org.opensearch.action.support.IndicesOptions.Option.class)
                        : EnumSet.copyOf(options),
                wildcardStates.isEmpty()
                        ? EnumSet.noneOf(org.opensearch.action.support.IndicesOptions.WildcardStates.class)
                        : EnumSet.copyOf(wildcardStates));
    }
    // endregion

    @Nullable
    private ElasticsearchPersistentEntity<?> getPersistentEntity(@Nullable Class<?> clazz) {
        return clazz != null ? elasticsearchConverter.getMappingContext().getPersistentEntity(clazz) : null;
    }
    // endregion

    private boolean hasSeqNoPrimaryTermProperty(@Nullable Class<?> entityClass) {

        if (entityClass == null) {
            return false;
        }

        if (!elasticsearchConverter.getMappingContext().hasPersistentEntityFor(entityClass)) {
            return false;
        }

        ElasticsearchPersistentEntity<?> entity =
                elasticsearchConverter.getMappingContext().getRequiredPersistentEntity(entityClass);
        return entity.hasSeqNoPrimaryTermProperty();
    }

    // endregion
}
