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

import com.google.common.annotations.VisibleForTesting;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.RequestContext;
import org.apache.atlas.annotation.GraphTransaction;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.impexp.AtlasImportResult;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.store.graph.AtlasEntityStore;
import org.apache.atlas.repository.store.graph.v2.AtlasEntityStreamForImport;
import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
import org.apache.atlas.repository.store.graph.v2.BulkImporterImpl;
import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
import org.apache.atlas.repository.store.graph.v2.EntityImportStream;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.v1.typesystem.types.utils.TypesUtil;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

import static org.apache.atlas.repository.Constants.HISTORICAL_GUID_PROPERTY_KEY;
import static org.apache.atlas.repository.store.graph.v2.BulkImporterImpl.updateImportProgress;

public class RegularImport extends ImportStrategy {
    private static final Logger LOG = LoggerFactory.getLogger(RegularImport.class);

    private static final int MAX_ATTEMPTS = 3;
    private static final String EXCEPTION_CLASS_NAME_PERMANENT_LOCKING_EXCEPTION = "PermanentLockingException";

    private final AtlasGraph           graph;
    private final AtlasEntityStore     entityStore;
    private final AtlasTypeRegistry    typeRegistry;
    private final EntityGraphRetriever entityGraphRetriever;
    private final boolean              directoryBasedImportConfigured;

    public RegularImport(AtlasGraph graph, AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
        this.graph                          = graph;
        this.entityStore                    = entityStore;
        this.typeRegistry                   = typeRegistry;
        this.entityGraphRetriever           = new EntityGraphRetriever(graph, typeRegistry);
        this.directoryBasedImportConfigured = StringUtils.isNotEmpty(AtlasConfiguration.IMPORT_TEMP_DIRECTORY.getString());
    }

    @VisibleForTesting
    static String getJsonArray(String json, String vertexGuid) {
        String quotedGuid = String.format("\"%s\"", vertexGuid);

        if (StringUtils.isEmpty(json)) {
            json = String.format("[%s]", quotedGuid);
        } else {
            json = json.replace("]", "").concat(",").concat(quotedGuid).concat("]");
        }

        return json;
    }

    @Override
    public EntityMutationResponse run(EntityImportStream entityStream, AtlasImportResult importResult) throws AtlasBaseException {
        LOG.debug("==> bulkImport()");

        if (entityStream == null || !entityStream.hasNext()) {
            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entities to create/update.");
        }

        Set<String>  processedGuids = new HashSet<>();
        float        currentPercent = 0f;
        List<String> residualList   = new ArrayList<>();

        EntityMutationResponse             ret                                = null;
        EntityImportStreamWithResidualList entityImportStreamWithResidualList = new EntityImportStreamWithResidualList(entityStream, residualList);

        while (entityImportStreamWithResidualList.hasNext()) {
            AtlasEntityWithExtInfo entityWithExtInfo = entityImportStreamWithResidualList.getNextEntityWithExtInfo();

            TypesUtil.Pair<EntityMutationResponse, Float> result = run(entityWithExtInfo, ret, importResult, processedGuids, entityStream.getPosition(), entityImportStreamWithResidualList.getStreamSize(), currentPercent, entityImportStreamWithResidualList.residualList);

            ret            = result.left;
            currentPercent = result.right;
        }

        importResult.getProcessedEntities().addAll(processedGuids);

        LOG.info("bulkImport(): done. Total number of entities (including referred entities) imported: {}", processedGuids.size());

        return ret;
    }

    @Override
    public TypesUtil.Pair<EntityMutationResponse, Float> run(AtlasEntityWithExtInfo entityWithExtInfo, EntityMutationResponse ret,
                                                             AtlasImportResult importResult, Set<String> processedGuids, int entityStreamPosition,
                                                             int streamSize, float currentPercent, List<String> residualList) throws AtlasBaseException {
        if (ret == null) {
            ret = new EntityMutationResponse();

            ret.setGuidAssignments(new HashMap<>());
        }

        AtlasEntity entity = entityWithExtInfo != null ? entityWithExtInfo.getEntity() : null;

        if (entity != null) {
            if (processedGuids == null) {
                processedGuids = new HashSet<>();
            }

            if (residualList == null) {
                residualList = new ArrayList<>();
            }

            boolean isVertexGuidUpdated = false;

            for (int attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
                try {
                    AtlasEntityStreamForImport oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);
                    EntityMutationResponse     resp            = entityStore.createOrUpdateForImport(oneEntityStream);

                    if (resp.getGuidAssignments() != null) {
                        ret.getGuidAssignments().putAll(resp.getGuidAssignments());
                    }

                    currentPercent = updateImportMetrics(entityWithExtInfo, resp, importResult, processedGuids, entityStreamPosition, streamSize, currentPercent);

                    break;
                } catch (AtlasBaseException e) {
                    if (!updateResidualList(e, residualList, entityWithExtInfo.getEntity().getGuid())) {
                        throw e;
                    }

                    break;
                } catch (AtlasSchemaViolationException e) {
                    LOG.debug("Entity: {}", entity.getGuid(), e);

                    if (!isVertexGuidUpdated) {
                        updateVertexGuid(entityWithExtInfo);

                        isVertexGuidUpdated = true;
                    } else {
                        LOG.error("Guid update failed: {}", entityWithExtInfo.getEntity().getGuid());

                        throw e;
                    }
                } catch (Throwable e) {
                    List<Throwable> throwableList = ExceptionUtils.getThrowableList(e);

                    if (!throwableList.isEmpty() && containsException(throwableList, EXCEPTION_CLASS_NAME_PERMANENT_LOCKING_EXCEPTION)) {
                        if (attempt < MAX_ATTEMPTS - 1) {
                            LOG.error("Caught {} , Retrying the transaction, attempt count is:{}", EXCEPTION_CLASS_NAME_PERMANENT_LOCKING_EXCEPTION, attempt);
                            continue;
                        } else {
                            throw e;
                        }
                    }

                    AtlasBaseException abe = new AtlasBaseException(e);

                    if (!updateResidualList(abe, residualList, entityWithExtInfo.getEntity().getGuid())) {
                        throw abe;
                    }

                    LOG.warn("Exception: {}", entity.getGuid(), e);

                    break;
                } finally {
                    RequestContext.get().clearCache();
                }
            }
        }

        return TypesUtil.Pair.of(ret, currentPercent);
    }

    private boolean containsException(final List<Throwable> exceptions, final String exceptionName) {
        return exceptions.stream().anyMatch(o -> o.getClass().getSimpleName().equals(exceptionName));
    }

    @GraphTransaction
    public void updateVertexGuid(AtlasEntityWithExtInfo entityWithExtInfo) {
        updateVertexGuid(entityWithExtInfo.getEntity());

        if (MapUtils.isNotEmpty(entityWithExtInfo.getReferredEntities())) {
            for (AtlasEntity entity : entityWithExtInfo.getReferredEntities().values()) {
                updateVertexGuid(entity);
            }
        }
    }

    public void updateVertexGuid(AtlasEntity entity) {
        String          entityGuid = entity.getGuid();
        AtlasObjectId   objectId   = entityGraphRetriever.toAtlasObjectIdWithoutGuid(entity);
        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
        String          vertexGuid;

        try {
            vertexGuid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(this.graph, entityType, objectId.getUniqueAttributes());
        } catch (AtlasBaseException e) {
            LOG.warn("Entity: {}: Does not exist!", objectId);

            return;
        }

        if (StringUtils.isEmpty(vertexGuid) || vertexGuid.equals(entityGuid)) {
            return;
        }

        AtlasVertex v = AtlasGraphUtilsV2.findByGuid(this.graph, vertexGuid);

        if (v == null) {
            return;
        }

        addHistoricalGuid(v, vertexGuid);

        AtlasGraphUtilsV2.setProperty(v, Constants.GUID_PROPERTY_KEY, entityGuid);

        LOG.warn("GUID Updated: Entity: {}: from: {}: to: {}", objectId, vertexGuid, entity.getGuid());
    }

    private void addHistoricalGuid(AtlasVertex v, String vertexGuid) {
        String existingJson = AtlasGraphUtilsV2.getProperty(v, HISTORICAL_GUID_PROPERTY_KEY, String.class);

        AtlasGraphUtilsV2.setProperty(v, HISTORICAL_GUID_PROPERTY_KEY, getJsonArray(existingJson, vertexGuid));
    }

    private boolean updateResidualList(AtlasBaseException e, List<String> lineageList, String guid) {
        if (!e.getAtlasErrorCode().getErrorCode().equals(AtlasErrorCode.INVALID_OBJECT_ID.getErrorCode())) {
            return false;
        }

        lineageList.add(guid);

        return true;
    }

    private float updateImportMetrics(AtlasEntity.AtlasEntityWithExtInfo currentEntity, EntityMutationResponse resp, AtlasImportResult importResult, Set<String> processedGuids, int currentIndex, int streamSize, float currentPercent) {
        if (!directoryBasedImportConfigured) {
            BulkImporterImpl.updateImportMetrics("entity:%s:created", resp.getCreatedEntities(), processedGuids, importResult);
            BulkImporterImpl.updateImportMetrics("entity:%s:updated", resp.getUpdatedEntities(), processedGuids, importResult);
            BulkImporterImpl.updateImportMetrics("entity:%s:deleted", resp.getDeletedEntities(), processedGuids, importResult);
        }

        String lastEntityImported = String.format("entity:last-imported:%s:[%s]:(%s)", currentEntity.getEntity().getTypeName(), currentIndex, currentEntity.getEntity().getGuid());

        return updateImportProgress(LOG, currentIndex, streamSize, currentPercent, lastEntityImported);
    }

    private static class EntityImportStreamWithResidualList {
        private final EntityImportStream stream;
        private final List<String>       residualList;
        private       boolean            navigateResidualList;
        private       int                currentResidualListIndex;

        public EntityImportStreamWithResidualList(EntityImportStream stream, List<String> residualList) {
            this.stream                   = stream;
            this.residualList             = residualList;
            this.navigateResidualList     = false;
            this.currentResidualListIndex = 0;
        }

        public AtlasEntity.AtlasEntityWithExtInfo getNextEntityWithExtInfo() {
            if (!navigateResidualList) {
                return stream.getNextEntityWithExtInfo();
            } else {
                stream.setPositionUsingEntityGuid(residualList.get(currentResidualListIndex++));
                return stream.getNextEntityWithExtInfo();
            }
        }

        public boolean hasNext() {
            if (!navigateResidualList) {
                boolean streamHasNext = stream.hasNext();

                navigateResidualList = (!streamHasNext);

                return streamHasNext ? streamHasNext : (currentResidualListIndex < residualList.size());
            } else {
                return (currentResidualListIndex < residualList.size());
            }
        }

        public int getStreamSize() {
            return stream.size() + residualList.size();
        }
    }
}
