/**
 * 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.repository.graphdb.janus;

import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.discovery.SearchParameters.FilterCriteria;
import org.apache.atlas.model.discovery.SearchParameters.Operator;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
import org.apache.atlas.model.typedef.AtlasStructDef;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
import org.apache.commons.collections.CollectionUtils;
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.List;
import java.util.Map;
import java.util.Set;

import static org.apache.atlas.repository.Constants.CLASSIFICATION_NAMES_KEY;
import static org.apache.atlas.repository.Constants.CUSTOM_ATTRIBUTES_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.LABELS_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.PROPAGATED_CLASSIFICATION_NAMES_KEY;

public class AtlasSolrQueryBuilder {
    private static final Logger LOG = LoggerFactory.getLogger(AtlasSolrQueryBuilder.class);

    public static final char   CUSTOM_ATTR_SEPARATOR     = '=';
    public static final String CUSTOM_ATTR_SEARCH_FORMAT = "\"\\\"%s\\\":\\\"%s\\\"\"";

    private Set<AtlasEntityType> entityTypes;
    private String               queryString;
    private FilterCriteria       criteria;
    private boolean              excludeDeletedEntities;
    private boolean              includeSubtypes;
    private Map<String, String>  indexFieldNameCache;

    public AtlasSolrQueryBuilder() {
    }

    public AtlasSolrQueryBuilder withEntityTypes(Set<AtlasEntityType> searchForEntityTypes) {
        this.entityTypes = searchForEntityTypes;

        return this;
    }

    public AtlasSolrQueryBuilder withQueryString(String queryString) {
        this.queryString = queryString;

        return this;
    }

    public AtlasSolrQueryBuilder withCriteria(FilterCriteria criteria) {
        this.criteria = criteria;

        return this;
    }

    public AtlasSolrQueryBuilder withExcludedDeletedEntities(boolean excludeDeletedEntities) {
        this.excludeDeletedEntities = excludeDeletedEntities;

        return this;
    }

    public AtlasSolrQueryBuilder withIncludeSubTypes(boolean includeSubTypes) {
        this.includeSubtypes = includeSubTypes;

        return this;
    }

    public AtlasSolrQueryBuilder withCommonIndexFieldNames(Map<String, String> indexFieldNameCache) {
        this.indexFieldNameCache = indexFieldNameCache;

        return this;
    }

    public String build() throws AtlasBaseException {
        StringBuilder queryBuilder = new StringBuilder();
        boolean       isAndNeeded  = false;

        if (StringUtils.isNotEmpty(queryString)) {
            LOG.debug("Initial query string is {}.", queryString);

            queryBuilder.append("+").append(queryString.trim()).append(" ");

            isAndNeeded = true;
        }

        if (excludeDeletedEntities) {
            if (isAndNeeded) {
                queryBuilder.append(" AND ");
            }

            dropDeletedEntities(queryBuilder);

            isAndNeeded = true;
        }

        if (CollectionUtils.isNotEmpty(entityTypes)) {
            if (isAndNeeded) {
                queryBuilder.append(" AND ");
            }

            buildForEntityType(queryBuilder);

            isAndNeeded = true;
        }

        if (criteria != null) {
            StringBuilder attrFilterQueryBuilder = new StringBuilder();

            withCriteria(attrFilterQueryBuilder, criteria);

            if (attrFilterQueryBuilder.length() != 0) {
                if (isAndNeeded) {
                    queryBuilder.append(" AND ");
                }

                queryBuilder.append(" ").append(attrFilterQueryBuilder);
            }
        }

        return queryBuilder.toString();
    }

    private void buildForEntityType(StringBuilder queryBuilder) {
        String typeIndexFieldName = indexFieldNameCache.get(Constants.ENTITY_TYPE_PROPERTY_KEY);

        queryBuilder.append(" +")
                .append(typeIndexFieldName)
                .append(":(");

        Set<String> typesToSearch = new HashSet<>();

        for (AtlasEntityType type : entityTypes) {
            if (includeSubtypes) {
                typesToSearch.addAll(type.getTypeAndAllSubTypes());
            } else {
                typesToSearch.add(type.getTypeName());
            }
        }

        queryBuilder.append(StringUtils.join(typesToSearch, " ")).append(" ) ");
    }

    private void dropDeletedEntities(StringBuilder queryBuilder) throws AtlasBaseException {
        LOG.debug("excluding the deleted entities.");

        String indexFieldName = indexFieldNameCache.get(Constants.STATE_PROPERTY_KEY);

        if (indexFieldName == null) {
            String msg = String.format("There is no index field name defined for attribute '%s'", Constants.STATE_PROPERTY_KEY);

            LOG.error(msg);

            throw new AtlasBaseException(msg);
        }

        queryBuilder.append(" -").append(indexFieldName).append(":").append(AtlasEntity.Status.DELETED.name());
    }

    private AtlasSolrQueryBuilder withCriteria(StringBuilder queryBuilder, FilterCriteria criteria) throws AtlasBaseException {
        List<FilterCriteria> criterion       = criteria.getCriterion();
        Set<String>          indexAttributes = new HashSet<>();

        if (StringUtils.isNotEmpty(criteria.getAttributeName()) && CollectionUtils.isEmpty(criterion)) { // no child criterion
            String   attributeName  = criteria.getAttributeName();
            String   attributeValue = criteria.getAttributeValue();
            Operator operator       = criteria.getOperator();

            ArrayList<StringBuilder> orExpQuery = new ArrayList<>();

            for (AtlasEntityType type : entityTypes) {
                String indexAttributeName = getIndexAttributeName(type, attributeName);

                //check to remove duplicate attribute query (for eg. name)
                if (!indexAttributes.contains(indexAttributeName)) {
                    StringBuilder sb = new StringBuilder();

                    if (attributeName.equals(CUSTOM_ATTRIBUTES_PROPERTY_KEY)) {
                        // CustomAttributes stores key value pairs in String format, so ideally it should be 'contains' operator to search for one pair,
                        // for use-case, E1 having key1=value1 and E2 having key1=value2, searching key1=value1 results both E1,E2
                        // surrounding inverted commas to attributeValue works
                        if (operator.equals(Operator.CONTAINS)) {
                            operator = Operator.EQ;
                        } else if (operator.equals(Operator.NOT_CONTAINS)) {
                            operator = Operator.NEQ;
                        }

                        attributeValue = getIndexQueryAttributeValue(attributeValue);
                    }

                    if (attributeValue != null) {
                        attributeValue = attributeValue.trim();
                    }

                    //when wildcard search -> escape special Char, don't quote
                    //      when  tokenized characters + index field Type TEXT -> remove wildcard '*' from query

                    //'CONTAINS and NOT_CONTAINS' operator with tokenize char in attributeValue doesn't guarantee correct results
                    // for aggregationMetrics
                    boolean                          replaceWildcardChar = false;
                    AtlasStructDef.AtlasAttributeDef def                 = type.getAttributeDef(attributeName);

                    if (!isPipeSeparatedSystemAttribute(attributeName) && isWildCardOperator(operator) && def.getTypeName().equalsIgnoreCase(AtlasBaseTypeDef.ATLAS_TYPE_STRING)) {
                        if (def.getIndexType() == null && AtlasAttribute.hastokenizeChar(attributeValue)) {
                            replaceWildcardChar = true;
                        }

                        attributeValue = AtlasAttribute.escapeIndexQueryValue(attributeValue, false, false);
                    } else {
                        attributeValue = AtlasAttribute.escapeIndexQueryValue(attributeValue);
                    }

                    withPropertyCondition(sb, indexAttributeName, operator, attributeValue, replaceWildcardChar);

                    indexAttributes.add(indexAttributeName);
                    orExpQuery.add(sb);
                }
            }

            if (CollectionUtils.isNotEmpty(orExpQuery)) {
                if (orExpQuery.size() > 1) {
                    String orExpStr = StringUtils.join(orExpQuery, FilterCriteria.Condition.OR.name());

                    queryBuilder.append(" ( ").append(orExpStr).append(" ) ");
                } else {
                    queryBuilder.append(orExpQuery.iterator().next());
                }
            }
        } else if (CollectionUtils.isNotEmpty(criterion)) {
            beginCriteria(queryBuilder);

            for (Iterator<FilterCriteria> iterator = criterion.iterator(); iterator.hasNext(); ) {
                FilterCriteria childCriteria = iterator.next();

                withCriteria(queryBuilder, childCriteria);

                if (iterator.hasNext()) {
                    withCondition(queryBuilder, criteria.getCondition().name());
                }
            }

            endCriteria(queryBuilder);
        }

        return this;
    }

    private void withPropertyCondition(StringBuilder queryBuilder, String indexFieldName, Operator operator, String attributeValue, boolean replaceWildCard) throws AtlasBaseException {
        if (StringUtils.isNotEmpty(indexFieldName) && operator != null) {
            beginCriteria(queryBuilder);

            switch (operator) {
                case EQ:
                    withEqual(queryBuilder, indexFieldName, attributeValue);
                    break;
                case NEQ:
                    withNotEqual(queryBuilder, indexFieldName, attributeValue);
                    break;
                case STARTS_WITH:
                    withStartsWith(queryBuilder, indexFieldName, attributeValue, replaceWildCard);
                    break;
                case ENDS_WITH:
                    withEndsWith(queryBuilder, indexFieldName, attributeValue, replaceWildCard);
                    break;
                case CONTAINS:
                    withContains(queryBuilder, indexFieldName, attributeValue, replaceWildCard);
                    break;
                case NOT_CONTAINS:
                    withNotContains(queryBuilder, indexFieldName, attributeValue, replaceWildCard);
                    break;
                case IS_NULL:
                    withIsNull(queryBuilder, indexFieldName);
                    break;
                case NOT_NULL:
                    withIsNotNull(queryBuilder, indexFieldName);
                    break;
                case LT:
                    withLessthan(queryBuilder, indexFieldName, attributeValue);
                    break;
                case GT:
                    withGreaterThan(queryBuilder, indexFieldName, attributeValue);
                    break;
                case LTE:
                    withLessthanOrEqual(queryBuilder, indexFieldName, attributeValue);
                    break;
                case GTE:
                    withGreaterThanOrEqual(queryBuilder, indexFieldName, attributeValue);
                    break;
                case IN:
                case LIKE:
                case CONTAINS_ANY:
                case CONTAINS_ALL:
                default:
                    String msg = String.format("%s is not supported operation.", operator.getSymbol());
                    LOG.error(msg);
                    throw new AtlasBaseException(msg);
            }

            endCriteria(queryBuilder);
        }
    }

    private boolean isPipeSeparatedSystemAttribute(String attrName) {
        return StringUtils.equals(attrName, CLASSIFICATION_NAMES_KEY) ||
                StringUtils.equals(attrName, PROPAGATED_CLASSIFICATION_NAMES_KEY) ||
                StringUtils.equals(attrName, LABELS_PROPERTY_KEY) ||
                StringUtils.equals(attrName, CUSTOM_ATTRIBUTES_PROPERTY_KEY);
    }

    private boolean isWildCardOperator(Operator operator) {
        return (operator == Operator.CONTAINS ||
                operator == Operator.STARTS_WITH ||
                operator == Operator.ENDS_WITH ||
                operator == Operator.NOT_CONTAINS);
    }

    private String getIndexQueryAttributeValue(String attributeValue) {
        if (StringUtils.isNotEmpty(attributeValue)) {
            int    separatorIdx = attributeValue.indexOf(CUSTOM_ATTR_SEPARATOR);
            String key          = separatorIdx != -1 ? attributeValue.substring(0, separatorIdx) : null;

            if (key != null) {
                String value = attributeValue.substring(separatorIdx + 1);

                return String.format(CUSTOM_ATTR_SEARCH_FORMAT, key, value);
            }
        }

        return attributeValue;
    }

    private String getIndexAttributeName(AtlasEntityType type, String attrName) throws AtlasBaseException {
        AtlasAttribute ret = type.getAttribute(attrName);

        if (ret == null) {
            String msg = String.format("Received unknown attribute '%s' for type '%s'.", attrName, type.getTypeName());

            LOG.error(msg);

            throw new AtlasBaseException(msg);
        }

        String indexFieldName = ret.getIndexFieldName();

        if (indexFieldName == null) {
            String msg = String.format("Received non-index attribute %s for type %s.", attrName, type.getTypeName());

            LOG.error(msg);

            throw new AtlasBaseException(msg);
        }

        return indexFieldName;
    }

    private void beginCriteria(StringBuilder queryBuilder) {
        queryBuilder.append("( ");
    }

    private void endCriteria(StringBuilder queryBuilder) {
        queryBuilder.append(" )");
    }

    private void withEndsWith(StringBuilder queryBuilder, String indexFieldName, String attributeValue, boolean replaceWildCard) {
        String attrValuePrefix = replaceWildCard ? ":" : ":*";

        queryBuilder.append("+").append(indexFieldName).append(attrValuePrefix).append(attributeValue).append(" ");
    }

    private void withStartsWith(StringBuilder queryBuilder, String indexFieldName, String attributeValue, boolean replaceWildCard) {
        String attrValuePostfix = replaceWildCard ? " " : "* ";

        queryBuilder.append("+").append(indexFieldName).append(":").append(attributeValue).append(attrValuePostfix);
    }

    private void withNotEqual(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        queryBuilder.append("*:* -").append(indexFieldName).append(":").append(attributeValue).append(" ");
    }

    private void withEqual(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        queryBuilder.append("+").append(indexFieldName).append(":").append(attributeValue).append(" ");
    }

    private void withGreaterThan(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        //{ == exclusive
        //] == inclusive
        //+__timestamp_l:{<attributeValue> TO *]
        queryBuilder.append("+").append(indexFieldName).append(":{ ").append(attributeValue).append(" TO * ] ");
    }

    private void withGreaterThanOrEqual(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        //[ == inclusive
        //] == inclusive
        //+__timestamp_l:[<attributeValue> TO *]
        queryBuilder.append("+").append(indexFieldName).append(":[ ").append(attributeValue).append(" TO * ] ");
    }

    private void withLessthan(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        //[ == inclusive
        //} == exclusive
        //+__timestamp_l:[* TO <attributeValue>}
        queryBuilder.append("+").append(indexFieldName).append(":[ * TO ").append(attributeValue).append("} ");
    }

    private void withLessthanOrEqual(StringBuilder queryBuilder, String indexFieldName, String attributeValue) {
        //[ == inclusive
        //[ == inclusive
        //+__timestamp_l:[* TO <attributeValue>]
        queryBuilder.append("+").append(indexFieldName).append(":[ * TO ").append(attributeValue).append(" ] ");
    }

    private void withContains(StringBuilder queryBuilder, String indexFieldName, String attributeValue, boolean replaceWildCard) {
        String attrValuePrefix  = replaceWildCard ? ":" : ":*";
        String attrValuePostfix = replaceWildCard ? " " : "* ";

        queryBuilder.append("+").append(indexFieldName).append(attrValuePrefix).append(attributeValue).append(attrValuePostfix);
    }

    private void withNotContains(StringBuilder queryBuilder, String indexFieldName, String attributeValue, boolean replaceWildCard) {
        String attrValuePrefix  = replaceWildCard ? ":" : ":*";
        String attrValuePostfix = replaceWildCard ? " " : "* ";

        queryBuilder.append("*:* -").append(indexFieldName).append(attrValuePrefix).append(attributeValue).append(attrValuePostfix);
    }

    private void withIsNull(StringBuilder queryBuilder, String indexFieldName) {
        queryBuilder.append("-").append(indexFieldName).append(":*").append(" ");
    }

    private void withIsNotNull(StringBuilder queryBuilder, String indexFieldName) {
        queryBuilder.append("+").append(indexFieldName).append(":*").append(" ");
    }

    private void withCondition(StringBuilder queryBuilder, String condition) {
        queryBuilder.append(" ").append(condition).append(" ");
    }
}
