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

import org.apache.atlas.GraphTransactionInterceptor;
import org.apache.atlas.RequestContext;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.pc.WorkItemConsumer;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
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.BulkImporterImpl;
import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
import org.apache.atlas.repository.store.graph.v2.EntityStream;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.utils.AtlasPerfMetrics;
import org.apache.commons.collections.MapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.atomic.AtomicLong;

public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWithExtInfo> {
    private static final Logger LOG                    = LoggerFactory.getLogger(EntityConsumer.class);
    private static final int    MAX_COMMIT_RETRY_COUNT = 3;

    private final int                                      batchSize;
    private final AtlasEntityStore                         entityStore;
    private final AtlasGraph                               atlasGraphBulk;
    private final AtlasEntityStore                         entityStoreBulk;
    private final AtlasTypeRegistry                        typeRegistry;
    private final EntityGraphRetriever                     entityRetrieverBulk;
    private final boolean                                  isMigrationImport;
    private final AtomicLong                               counter      = new AtomicLong(1);
    private final AtomicLong                               currentBatch = new AtomicLong(1);
    private final AtlasGraph                               atlasGraph;
    private final List<AtlasEntity.AtlasEntityWithExtInfo> entityBuffer = new ArrayList<>();
    private final List<String>                             localResults = new ArrayList<>();

    public EntityConsumer(AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph, AtlasEntityStore entityStore,
            AtlasGraph atlasGraphBulk, AtlasEntityStore entityStoreBulk, EntityGraphRetriever entityRetrieverBulk,
            BlockingQueue<AtlasEntity.AtlasEntityWithExtInfo> queue, int batchSize, boolean isMigrationImport) {
        super(queue);

        this.typeRegistry        = typeRegistry;
        this.atlasGraph          = atlasGraph;
        this.entityStore         = entityStore;
        this.atlasGraphBulk      = atlasGraphBulk;
        this.entityStoreBulk     = entityStoreBulk;
        this.entityRetrieverBulk = entityRetrieverBulk;
        this.batchSize           = batchSize;
        this.isMigrationImport   = isMigrationImport;
    }

    @Override
    protected void commitDirty() {
        super.commitDirty();

        LOG.info("Total: Commit: {}", counter.get());

        counter.set(0);
    }

    @Override
    protected void doCommit() {
        for (int retryCount = 1; retryCount <= MAX_COMMIT_RETRY_COUNT; retryCount++) {
            if (commitWithRetry(retryCount)) {
                return;
            }
        }

        LOG.error("Retries exceeded! Potential data loss! Please correct data and re-attempt. Buffer: {}: Counter: {}", entityBuffer.size(), counter.get());

        clear();
    }

    @Override
    protected void processItem(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
        int delta = MapUtils.isEmpty(entityWithExtInfo.getReferredEntities()) ? 1 : entityWithExtInfo.getReferredEntities().size() + 1;

        long currentCount = counter.addAndGet(delta);

        currentBatch.addAndGet(delta);

        try {
            processEntity(entityWithExtInfo, currentCount);

            attemptCommit();
        } catch (Exception e) {
            LOG.info("Invalid entities. Possible data loss: Please correct and re-submit!", e);
        }
    }

    private void processEntity(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo, long currentCount) {
        RequestContext.get().setImportInProgress(true);
        RequestContext.get().setCreateShellEntityForNonExistingReference(true);
        RequestContext.get().setMigrationInProgress(this.isMigrationImport);

        try {
            LOG.debug("Processing: {}", currentCount);

            importUsingBulkEntityStore(entityWithExtInfo);
        } catch (IllegalStateException | IllegalArgumentException e) {
            LOG.warn("{}: {} - {}", e.getClass().getSimpleName(), entityWithExtInfo.getEntity().getTypeName(), entityWithExtInfo.getEntity().getGuid(), e);

            importUsingRegularEntityStore(entityWithExtInfo, e);
        } catch (AtlasBaseException e) {
            LOG.warn("AtlasBaseException: {} - {}", entityWithExtInfo.getEntity().getTypeName(), entityWithExtInfo.getEntity().getGuid(), e);
        } catch (AtlasSchemaViolationException e) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Entity: {}", entityWithExtInfo.getEntity().getGuid(), e);
            }

            BulkImporterImpl.updateVertexGuid(this.atlasGraphBulk, typeRegistry, entityRetrieverBulk, entityWithExtInfo.getEntity());
        }
    }

    private void importUsingBulkEntityStore(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
        EntityStream           oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);
        EntityMutationResponse result          = entityStoreBulk.createOrUpdateForImportNoCommit(oneEntityStream);

        localResults.add(entityWithExtInfo.getEntity().getGuid());
        entityBuffer.add(entityWithExtInfo);
    }

    private void importUsingRegularEntityStore(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo, Exception ex) {
        commitValidatedEntities(ex);
        performRegularImport(entityWithExtInfo);
    }

    private void performRegularImport(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
        synchronized (atlasGraph) {
            try {
                LOG.info("Regular: EntityStore: {}: Starting...", this.counter.get());

                AtlasEntityStreamForImport oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);

                this.entityStore.createOrUpdateForImportNoCommit(oneEntityStream);

                atlasGraph.commit();

                localResults.add(entityWithExtInfo.getEntity().getGuid());

                dispatchResults();
            } catch (Exception e) {
                atlasGraph.rollback();

                LOG.error("Regular: EntityStore: Rollback!: Entity creation using regular (non-bulk) failed! Please correct entity and re-submit!", e);
            } finally {
                LOG.info("Regular: EntityStore: {}: Commit: Done!", this.counter.get());

                atlasGraph.commit();
                addResult(entityWithExtInfo.getEntity().getGuid());
                clear();

                LOG.info("Regular: EntityStore: {}: Done!", this.counter.get());
            }
        }
    }

    private void commitValidatedEntities(Exception ex) {
        try {
            LOG.info("Validated Entities: Commit: Starting...");

            rollbackPauseRetry(1, ex);

            doCommit();
        } finally {
            LOG.info("Validated Entities: Commit: Done!");
        }
    }

    private void attemptCommit() {
        if (currentBatch.get() < batchSize) {
            return;
        }

        doCommit();
    }

    private boolean commitWithRetry(int retryCount) {
        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("commitWithRetry");

        try {
            atlasGraphBulk.commit();

            if (LOG.isDebugEnabled()) {
                LOG.debug("Commit: Done!: Buffer: {}: Batch: {}: Counter: {}", entityBuffer.size(), currentBatch.get(), counter.get());
            }

            dispatchResults();
            return true;
        } catch (Exception ex) {
            rollbackPauseRetry(retryCount, ex);

            return false;
        } finally {
            RequestContext.get().endMetricRecord(metric);
        }
    }

    private void rollbackPauseRetry(int retryCount, Exception ex) {
        bulkGraphRollback(retryCount);

        LOG.warn("Rollback: Done! Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount);

        pause(retryCount);

        String exceptionClass = ex.getClass().getSimpleName();

        if (!exceptionClass.equals("JanusGraphException") && !exceptionClass.equals("PermanentLockingException")) {
            LOG.warn("Commit error! Will pause and retry: Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount, ex);
        }

        retryProcessEntity(retryCount);
    }

    private void bulkGraphRollback(int retryCount) {
        try {
            atlasGraphBulk.rollback();
            clearCache();
        } catch (Exception e) {
            LOG.error("Rollback: Exception! Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount);
        }
    }

    private void retryProcessEntity(int retryCount) {
        if (LOG.isDebugEnabled() || retryCount > 1) {
            LOG.info("Replaying: Starting!: Buffer: {}: Retry count: {}", entityBuffer.size(), retryCount);
        }

        List<AtlasEntity.AtlasEntityWithExtInfo> localBuffer = new ArrayList<>(entityBuffer);

        entityBuffer.clear();

        for (AtlasEntity.AtlasEntityWithExtInfo e : localBuffer) {
            processEntity(e, counter.get());
        }

        LOG.info("Replaying: Done!: Buffer: {}: Retry count: {}", entityBuffer.size(), retryCount);
    }

    private void dispatchResults() {
        localResults.forEach(this::addResult);

        clear();
    }

    private void pause(int retryCount) {
        try {
            Thread.sleep(1000L * retryCount);
        } catch (InterruptedException e) {
            LOG.error("pause: Interrupted!", e);
        }
    }

    private void clear() {
        localResults.clear();
        entityBuffer.clear();
        clearCache();
        currentBatch.set(0);
    }

    private void clearCache() {
        GraphTransactionInterceptor.clearCache();
        RequestContext.get().clearCache();
    }
}
