/**
 * 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.graph;

import org.apache.atlas.AtlasException;
import org.apache.atlas.listener.ChangedTypeDefs;
import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasGraphIndexClient;
import org.apache.atlas.type.AtlasBusinessMetadataType;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasStructType;
import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.util.AtlasRepositoryConfiguration;
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.HashMap;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.DEFAULT_SEARCHWEIGHT;
import static org.apache.atlas.repository.Constants.CLASSIFICATION_TEXT_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.TYPE_NAME_PROPERTY_KEY;

/**
 * This is a component that will go through all entity type definitions and create free text index
 * request handler with SOLR. This is a no op class in non-solr index based deployments.
 * This component needs to be initialized after type definitions are completely fixed with the needed patches (Ordder 3 initialization).
 */
public class SolrIndexHelper implements IndexChangeListener {
    public static final int DEFAULT_SEARCHWEIGHT_FOR_STRINGS = 3;
    public static final int SEARCHWEIGHT_FOR_CLASSIFICATIONS = 10;
    public static final int SEARCHWEIGHT_FOR_LABELS          = 10;
    public static final int SEARCHWEIGHT_FOR_CUSTOM_ATTRS    = 3;
    public static final int SEARCHWEIGHT_FOR_TYPENAME        = 1;
    private static final Logger LOG = LoggerFactory.getLogger(SolrIndexHelper.class);
    private static final int MIN_SEARCH_WEIGHT_FOR_SUGGESTIONS = 8;

    private final AtlasTypeRegistry typeRegistry;

    private boolean initializationCompleted;

    public SolrIndexHelper(AtlasTypeRegistry typeRegistry) {
        this.typeRegistry = typeRegistry;
    }

    @Override
    public void onChange(ChangedTypeDefs changedTypeDefs) {
        if (!AtlasRepositoryConfiguration.isFreeTextSearchEnabled()) {
            return;
        }

        if (changedTypeDefs == null || !(changedTypeDefs.hasEntityDef() || changedTypeDefs.hasBusinessMetadataDef())) {
            LOG.info("SolrIndexHelper.onChange(): no change in entity/business-metadata types. No updates needed.");

            return;
        }

        if (initializationCompleted) {
            try {
                AtlasGraph            graph                          = AtlasGraphProvider.getGraphInstance();
                AtlasGraphIndexClient graphIndexClient               = graph.getGraphIndexClient();
                Map<String, Integer>  indexFieldName2SearchWeightMap = geIndexFieldNamesWithSearchWeights();

                graphIndexClient.applySearchWeight(Constants.VERTEX_INDEX, indexFieldName2SearchWeightMap);
                graphIndexClient.applySuggestionFields(Constants.VERTEX_INDEX, getIndexFieldNamesForSuggestions(indexFieldName2SearchWeightMap));
            } catch (AtlasException e) {
                LOG.error("Error encountered in handling type system change notification.", e);

                throw new RuntimeException("Error encountered in handling type system change notification.", e);
            }
        }
    }

    @Override
    public void onInitStart() {
        LOG.info("SolrIndexHelper Initialization started.");

        initializationCompleted = false;
    }

    @Override
    public void onInitCompletion(ChangedTypeDefs changedTypeDefs) {
        LOG.info("SolrIndexHelper Initialization completed.");

        initializationCompleted = true;

        onChange(changedTypeDefs);
    }

    private List<String> getIndexFieldNamesForSuggestions(Map<String, Integer> indexFieldName2SearchWeightMap) {
        List<String> ret = new ArrayList<>();

        for (Map.Entry<String, Integer> entry : indexFieldName2SearchWeightMap.entrySet()) {
            if (entry.getValue() >= MIN_SEARCH_WEIGHT_FOR_SUGGESTIONS) {
                String indexFieldName = entry.getKey();

                LOG.debug("Adding indexFieldName {} for suggestions.", indexFieldName);

                ret.add(indexFieldName);
            }
        }

        return ret;
    }

    private Map<String, Integer> geIndexFieldNamesWithSearchWeights() {
        Map<String, Integer> ret = new HashMap<>();

        //the following properties are specially added manually.
        //as, they don't come in the entity definitions as attributes.
        ret.put(typeRegistry.getIndexFieldName(CLASSIFICATION_TEXT_KEY), SEARCHWEIGHT_FOR_CLASSIFICATIONS);
        ret.put(typeRegistry.getIndexFieldName(LABELS_PROPERTY_KEY), SEARCHWEIGHT_FOR_LABELS);
        ret.put(typeRegistry.getIndexFieldName(CUSTOM_ATTRIBUTES_PROPERTY_KEY), SEARCHWEIGHT_FOR_CUSTOM_ATTRS);
        ret.put(typeRegistry.getIndexFieldName(TYPE_NAME_PROPERTY_KEY), SEARCHWEIGHT_FOR_TYPENAME);

        for (AtlasEntityType entityType : typeRegistry.getAllEntityTypes()) {
            if (entityType.isInternalType()) {
                continue;
            }

            processType(ret, entityType);
        }

        for (AtlasBusinessMetadataType businessMetadataType : typeRegistry.getAllBusinessMetadataTypes()) {
            processType(ret, businessMetadataType);
        }

        return ret;
    }

    private void processType(Map<String, Integer> indexFieldNameWithSearchWeights, AtlasStructType structType) {
        List<AtlasAttributeDef> attributes = structType.getStructDef().getAttributeDefs();

        if (CollectionUtils.isNotEmpty(attributes)) {
            for (AtlasAttributeDef attribute : attributes) {
                processAttribute(indexFieldNameWithSearchWeights, structType.getAttribute(attribute.getName()));
            }
        } else {
            LOG.debug("No attributes are defined for type {}", structType.getTypeName());
        }
    }

    private void processAttribute(Map<String, Integer> indexFieldNameWithSearchWeights, AtlasAttribute attribute) {
        if (attribute != null && GraphBackedSearchIndexer.isStringAttribute(attribute) && StringUtils.isNotEmpty(attribute.getIndexFieldName())) {
            int searchWeight = attribute.getSearchWeight();

            if (searchWeight == DEFAULT_SEARCHWEIGHT) {
                //We will use default search weight of 3 for string attributes.
                //this will make the string data searchable like in FullTextIndex Searcher using Free Text searcher.
                searchWeight = DEFAULT_SEARCHWEIGHT_FOR_STRINGS;
            } else if (!GraphBackedSearchIndexer.isValidSearchWeight(searchWeight)) { //validate the value provided in the model.
                LOG.warn("Invalid search weight {} for attribute {}. Will use default {}", searchWeight, attribute.getQualifiedName(), DEFAULT_SEARCHWEIGHT_FOR_STRINGS);

                searchWeight = DEFAULT_SEARCHWEIGHT_FOR_STRINGS;
            }

            LOG.debug("Applying search weight {} for attribute={}: indexFieldName={}", searchWeight, attribute.getQualifiedName(), attribute.getIndexFieldName());

            indexFieldNameWithSearchWeights.put(attribute.getIndexFieldName(), searchWeight);
        }
    }
}
