/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.atlas.discovery;

import org.apache.atlas.SortOrder;
import org.apache.atlas.model.discovery.SearchParameters;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.repository.graphdb.AtlasEdge;
import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.type.AtlasRelationshipType;
import org.apache.atlas.type.AtlasStructType;
import org.apache.atlas.util.SearchPredicateUtil;
import org.apache.atlas.utils.AtlasPerfTracer;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.Predicate;
import org.apache.commons.collections.PredicateUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static org.apache.atlas.repository.Constants.RELATIONSHIP_TYPE_PROPERTY_KEY;
import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.ASC;
import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.DESC;

public class RelationshipSearchProcessor extends SearchProcessor {
    private static final Logger LOG      = LoggerFactory.getLogger(RelationshipSearchProcessor.class);
    private static final Logger PERF_LOG = AtlasPerfTracer.getPerfLogger("RelationshipSearchProcessor");

    private final AtlasIndexQuery indexQuery;
    private final AtlasGraphQuery graphQuery;

    RelationshipSearchProcessor(SearchContext context, Set<String> indexedKeys) {
        super(context);

        context.setEdgeIndexKeys(indexedKeys);

        final Set<AtlasRelationshipType>      types                 = context.getRelationshipTypes();
        final SearchParameters.FilterCriteria filterCriteria        = context.getSearchParameters().getRelationshipFilters();
        final Set<String>                     indexAttributes       = new HashSet<>();
        final Set<String>                     graphAttributes       = new HashSet<>();
        final Set<String>                     allAttributes         = new HashSet<>();
        final Set<String>                     typeNames             = CollectionUtils.isNotEmpty(types) ? types.stream().map(AtlasRelationshipType::getTypeName).collect(Collectors.toSet()) : null;
        final String                          typeAndSubTypesQryStr = AtlasStructType.AtlasAttribute.escapeIndexQueryValue(typeNames, true);
        final Predicate                       typeNamePredicate     = SearchPredicateUtil.generateIsRelationshipEdgePredicate(context.getTypeRegistry());
        final String                          sortBy                = context.getSearchParameters().getSortBy();
        final SortOrder                       sortOrder             = context.getSearchParameters().getSortOrder();

        inMemoryPredicate = typeNamePredicate;

        processSearchAttributes(types, filterCriteria, indexAttributes, graphAttributes, allAttributes);

        final boolean typeSearchByIndex = typeAndSubTypesQryStr.length() <= MAX_QUERY_STR_LENGTH_TYPES;
        final boolean attrSearchByIndex = CollectionUtils.isNotEmpty(indexAttributes) && canApplyIndexFilter(types, filterCriteria, false);

        //index query
        StringBuilder indexQuery = new StringBuilder();

        if (typeSearchByIndex) {
            graphIndexQueryBuilder.addTypeAndSubTypesQueryFilterEdges(indexQuery, typeAndSubTypesQryStr);
        }

        if (attrSearchByIndex) {
            constructFilterQuery(indexQuery, types, filterCriteria, indexAttributes);

            Predicate attributePredicate = constructInMemoryPredicate(types, filterCriteria, indexAttributes);

            if (attributePredicate != null) {
                inMemoryPredicate = PredicateUtils.andPredicate(inMemoryPredicate, attributePredicate);
            }
        } else {
            graphAttributes.addAll(indexAttributes);
        }

        if (indexQuery.length() > 0) {
            String indexQueryString = STRAY_AND_PATTERN.matcher(indexQuery).replaceAll(")");

            indexQueryString = STRAY_OR_PATTERN.matcher(indexQueryString).replaceAll(")");
            indexQueryString = STRAY_ELIPSIS_PATTERN.matcher(indexQueryString).replaceAll("");

            this.indexQuery = context.getGraph().indexQuery(Constants.EDGE_INDEX, indexQueryString);
        } else {
            this.indexQuery = null;
        }

        //graph query
        if (CollectionUtils.isNotEmpty(graphAttributes) || !typeSearchByIndex) {
            AtlasGraphQuery query = context.getGraph().query();

            if (!typeSearchByIndex) {
                query.in(RELATIONSHIP_TYPE_PROPERTY_KEY, types);
            }

            graphQuery = toGraphFilterQuery(types, filterCriteria, graphAttributes, query);

            Predicate attributePredicate = constructInMemoryPredicate(types, filterCriteria, graphAttributes);

            if (attributePredicate != null) {
                inMemoryPredicate = PredicateUtils.andPredicate(inMemoryPredicate, attributePredicate);
            }

            if (StringUtils.isNotEmpty(sortBy)) {
                final AtlasRelationshipType    relationshipType = types.iterator().next();
                AtlasStructType.AtlasAttribute sortByAttribute  = relationshipType.getAttribute(sortBy);

                if (sortByAttribute != null && StringUtils.isNotEmpty(sortByAttribute.getVertexPropertyName())) {
                    AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;

                    graphQuery.orderBy(sortByAttribute.getVertexPropertyName(), qrySortOrder);
                }
            }
        } else {
            graphQuery = null;
        }
    }

    @Override
    public List<AtlasVertex> execute() {
        return null;
    }

    @Override
    public long getResultCount() {
        if (indexQuery != null) {
            return indexQuery.edgeTotals();
        } else {
            return -1L;
        }
    }

    public List<AtlasEdge> executeEdges() {
        List<AtlasEdge> ret  = new ArrayList<>();
        AtlasPerfTracer perf = null;

        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "RelationshipSearchProcessor.execute(" + context + ")");
        }

        try {
            final int     limit    = context.getSearchParameters().getLimit();
            final Integer marker   = context.getMarker();
            final int     startIdx = marker != null ? marker : context.getSearchParameters().getOffset();

            // when subsequent filtering stages are involved, query should start at 0 even though startIdx can be higher
            // first 'startIdx' number of entries will be ignored
            // if marker is provided, start query with marker offset
            int qryOffset;

            if (marker != null) {
                qryOffset = marker;
            } else {
                qryOffset = (graphQuery != null && indexQuery != null) ? 0 : startIdx;
            }

            int                               resultIdx     = qryOffset;
            LinkedHashMap<Integer, AtlasEdge> offsetEdgeMap = new LinkedHashMap<>();

            for (; ret.size() < limit; qryOffset += limit) {
                offsetEdgeMap.clear();

                if (context.terminateSearch()) {
                    LOG.warn("query terminated: {}", context.getSearchParameters());

                    break;
                }

                if (indexQuery != null) {
                    Iterator<AtlasIndexQuery.Result> idxQueryResult = executeIndexQueryForEdge(context, indexQuery, qryOffset, limit);

                    offsetEdgeMap = getEdgesFromIndexQueryResult(idxQueryResult, offsetEdgeMap, qryOffset);
                } else {
                    Iterator<AtlasEdge> queryResult = graphQuery.edges(qryOffset, limit).iterator();

                    offsetEdgeMap = getEdges(queryResult, offsetEdgeMap, qryOffset);
                }

                final boolean isLastResultPage = offsetEdgeMap.size() < limit;

                // Do in-memory filtering
                offsetEdgeMap = offsetEdgeMap.entrySet()
                        .stream()
                        .filter(x -> inMemoryPredicate.evaluate(x.getValue()))
                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, (x, y) -> y, LinkedHashMap::new));

                resultIdx = collectResultEdges(ret, startIdx, limit, resultIdx, offsetEdgeMap, marker);

                if (isLastResultPage) {
                    resultIdx = SearchContext.MarkerUtil.MARKER_END - 1;
                    break;
                }
            }

            if (marker != null) {
                nextOffset = resultIdx + 1;
            }
        } finally {
            AtlasPerfTracer.log(perf);
        }

        LOG.debug("<== RelationshipSearchProcessor.execute({}): ret.size()={}", context, ret.size());

        return ret;
    }
}
