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

import com.google.common.annotations.VisibleForTesting;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.impexp.AtlasExportRequest;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasStructType;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.util.AtlasGremlinQueryProvider;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.ScriptEngine;
import javax.script.ScriptException;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_CONTAINS;
import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_ENDS_WITH;
import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_FOR_TYPE;
import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_MATCHES;
import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_STARTS_WITH;

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

    static final String DEFAULT_MATCH               = "*";
    static final String BINDING_PARAMETER_TYPENAME  = "typeName";
    static final String BINDING_PARAMETER_ATTR_NAME = "attrName";
    static final String BINDING_PARAMTER_ATTR_VALUE = "attrValue";

    private final AtlasGraph          atlasGraph;
    private final AtlasTypeRegistry   typeRegistry;
    private final Map<String, String> matchTypeQuery;

    public StartEntityFetchByExportRequest(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, AtlasGremlinQueryProvider gremlinQueryProvider) {
        this.typeRegistry   = typeRegistry;
        this.atlasGraph     = atlasGraph;
        this.matchTypeQuery = initMatchTypeQueryMap(gremlinQueryProvider);
    }

    public List<AtlasObjectId> get(AtlasExportRequest exportRequest) {
        List<AtlasObjectId> list = new ArrayList<>();

        for (AtlasObjectId objectId : exportRequest.getItemsToExport()) {
            List<String> guids = get(exportRequest, objectId);

            if (guids.isEmpty()) {
                continue;
            }

            objectId.setGuid(guids.get(0));
            list.add(objectId);
        }

        return list;
    }

    public List<String> get(AtlasExportRequest exportRequest, AtlasObjectId item) {
        List<String> ret       = new ArrayList<>();
        String       matchType = exportRequest.getMatchTypeOptionValue();

        try {
            if (StringUtils.isNotEmpty(item.getGuid())) {
                ret.add(item.getGuid());

                return ret;
            }

            if (StringUtils.equalsIgnoreCase(matchType, MATCH_TYPE_FOR_TYPE) && StringUtils.isNotEmpty(item.getTypeName())) {
                ret = getEntitiesForMatchTypeType(item, matchType);

                return ret;
            }

            if (StringUtils.isNotEmpty(item.getTypeName()) && MapUtils.isNotEmpty(item.getUniqueAttributes())) {
                ret = getEntitiesForMatchTypeUsingUniqueAttributes(item, matchType);

                return ret;
            }
        } catch (AtlasBaseException ex) {
            LOG.error("Error fetching starting entity for: {}", item, ex);
        } finally {
            LOG.info("export(item={}; matchType={}, fetchType={}): found {} entities: options: {}", item,
                    exportRequest.getMatchTypeOptionValue(), exportRequest.getFetchTypeOptionValue(), ret.size(), AtlasType.toJson(exportRequest));
        }

        return ret;
    }

    public ScriptEngine getScriptEngine() {
        try {
            return atlasGraph.getGremlinScriptEngine();
        } catch (AtlasBaseException e) {
            LOG.error("Error initializing script engine.", e);
        }

        return null;
    }

    @VisibleForTesting
    String getQueryTemplateForMatchType(String matchType) {
        return matchTypeQuery.containsKey(matchType) ? matchTypeQuery.get(matchType) : matchTypeQuery.get(DEFAULT_MATCH);
    }

    @VisibleForTesting
    List<String> executeGremlinQuery(String query, Map<String, Object> bindings) {
        try {
            return (List<String>) atlasGraph.executeGremlinScript(getScriptEngine(), bindings, query, false);
        } catch (ScriptException e) {
            LOG.error("Script execution failed for query: {}, {} ", query, e);
            return null;
        }
    }

    private List<String> getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId item, String matchType) throws AtlasBaseException {
        final String          queryTemplate = getQueryTemplateForMatchType(matchType);
        final String          typeName      = item.getTypeName();
        final AtlasEntityType entityType    = typeRegistry.getEntityTypeByName(typeName);

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

        if (entityType == null) {
            throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, typeName);
        }

        for (Map.Entry<String, Object> e : item.getUniqueAttributes().entrySet()) {
            String attrName  = e.getKey();
            Object attrValue = e.getValue();

            AtlasStructType.AtlasAttribute attribute = entityType.getAttribute(attrName);

            if (attribute == null || attrValue == null) {
                continue;
            }

            List<String> guids = executeGremlinQuery(queryTemplate, getBindingsForObjectId(typeName, attribute.getQualifiedName(), e.getValue()));

            if (!CollectionUtils.isNotEmpty(guids)) {
                continue;
            }

            ret.addAll(guids);
        }

        return new ArrayList<>(ret);
    }

    private List<String> getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) {
        return executeGremlinQuery(getQueryTemplateForMatchType(matchType), getBindingsForTypeName(item.getTypeName()));
    }

    private HashMap<String, Object> getBindingsForTypeName(String typeName) {
        HashMap<String, Object> ret = new HashMap<>();

        ret.put(BINDING_PARAMETER_TYPENAME, new HashSet<>(Arrays.asList(StringUtils.split(typeName, ","))));

        return ret;
    }

    private HashMap<String, Object> getBindingsForObjectId(String typeName, String attrName, Object attrValue) {
        HashMap<String, Object> ret = new HashMap<>();

        ret.put(BINDING_PARAMETER_TYPENAME, typeName);
        ret.put(BINDING_PARAMETER_ATTR_NAME, attrName);
        ret.put(BINDING_PARAMTER_ATTR_VALUE, attrValue);

        return ret;
    }

    private static Map<String, String> initMatchTypeQueryMap(AtlasGremlinQueryProvider gremlinQueryProvider) {
        Map<String, String> ret = new HashMap<>();

        ret.put(MATCH_TYPE_STARTS_WITH, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_STARTS_WITH));
        ret.put(MATCH_TYPE_ENDS_WITH, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_ENDS_WITH));
        ret.put(MATCH_TYPE_CONTAINS, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_CONTAINS));
        ret.put(MATCH_TYPE_MATCHES, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_MATCHES));
        ret.put(MATCH_TYPE_FOR_TYPE, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_ALL_FOR_TYPE));
        ret.put(DEFAULT_MATCH, gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_TYPE_DEFAULT));

        return ret;
    }
}
