/**
 * 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 org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.impexp.AtlasImportRequest;
import org.apache.atlas.model.impexp.AtlasImportResult;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.repository.converters.AtlasFormatConverters;
import org.apache.atlas.repository.converters.AtlasInstanceConverter;
import org.apache.atlas.repository.graph.AtlasGraphProvider;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.migration.DataMigrationStatusService;
import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
import org.apache.atlas.repository.store.graph.v1.DeleteHandlerDelegate;
import org.apache.atlas.repository.store.graph.v2.AtlasEntityStoreV2;
import org.apache.atlas.repository.store.graph.v2.AtlasRelationshipStoreV2;
import org.apache.atlas.repository.store.graph.v2.EntityGraphMapper;
import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
import org.apache.atlas.repository.store.graph.v2.EntityImportStream;
import org.apache.atlas.repository.store.graph.v2.IAtlasEntityChangeNotifier;
import org.apache.atlas.repository.store.graph.v2.bulkimport.pc.EntityConsumerBuilder;
import org.apache.atlas.repository.store.graph.v2.bulkimport.pc.EntityCreationManager;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.utils.AtlasStringUtil;
import org.apache.atlas.v1.typesystem.types.utils.TypesUtil;
import org.apache.commons.lang3.NotImplementedException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Set;

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

    private final AtlasGraph         graph;
    private final AtlasGraphProvider graphProvider;
    private final AtlasTypeRegistry  typeRegistry;

    public MigrationImport(AtlasGraph graph, AtlasGraphProvider graphProvider, AtlasTypeRegistry typeRegistry) {
        this.graph         = graph;
        this.graphProvider = graphProvider;
        this.typeRegistry  = typeRegistry;

        LOG.info("MigrationImport: Using bulkLoading...");
    }

    public EntityMutationResponse run(EntityImportStream entityStream, AtlasImportResult importResult) throws AtlasBaseException {
        if (entityStream == null || !entityStream.hasNext()) {
            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entities to create/update.");
        }

        if (importResult.getRequest() == null) {
            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "importResult should contain request");
        }

        DataMigrationStatusService dataMigrationStatusService = createMigrationStatusService(importResult);

        long                   index           = 0;
        int                    streamSize      = entityStream.size();
        EntityMutationResponse ret             = new EntityMutationResponse();
        EntityCreationManager  creationManager = createEntityCreationManager(importResult, dataMigrationStatusService);

        try {
            LOG.info("Migration Import: Size: {}: Starting...", streamSize);

            index = creationManager.read(entityStream);

            creationManager.drain();
            creationManager.extractResults();
        } catch (Exception ex) {
            LOG.error("Migration Import: Error: Current position: {}", index, ex);
        } finally {
            shutdownEntityCreationManager(creationManager);
        }

        LOG.info("Migration Import: Size: {}: Done!", streamSize);

        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 {
        throw new NotImplementedException("MigrationImport: entity import not implemented");
    }

    private DataMigrationStatusService createMigrationStatusService(AtlasImportResult importResult) {
        DataMigrationStatusService dataMigrationStatusService = new DataMigrationStatusService();

        dataMigrationStatusService.init(AtlasStringUtil.getOption(importResult.getRequest().getOptions(), AtlasImportRequest.OPTION_KEY_MIGRATION_FILE_NAME));

        return dataMigrationStatusService;
    }

    private EntityCreationManager createEntityCreationManager(AtlasImportResult importResult, DataMigrationStatusService dataMigrationStatusService) {
        AtlasGraph           graphBulk                = graphProvider.getBulkLoading();
        EntityGraphRetriever entityGraphRetriever     = new EntityGraphRetriever(this.graph, typeRegistry);
        EntityGraphRetriever entityGraphRetrieverBulk = new EntityGraphRetriever(graphBulk, typeRegistry);
        AtlasEntityStoreV2   entityStore              = createEntityStore(this.graph, typeRegistry);
        AtlasEntityStoreV2   entityStoreBulk          = createEntityStore(graphBulk, typeRegistry);

        int     batchSize         = importResult.getRequest().getOptionKeyBatchSize();
        int     numWorkers        = getNumWorkers(importResult.getRequest().getOptionKeyNumWorkers());
        boolean isMigrationImport = false;

        if (importResult.getRequest().getOptions().get("migration") != null) {
            isMigrationImport = Boolean.parseBoolean(importResult.getRequest().getOptions().get("migration"));
        }

        EntityConsumerBuilder consumerBuilder = new EntityConsumerBuilder(typeRegistry, this.graph, entityStore, entityGraphRetriever, graphBulk, entityStoreBulk, entityGraphRetrieverBulk, batchSize, isMigrationImport);

        LOG.info("MigrationImport: EntityCreationManager: Created!");

        return new EntityCreationManager(consumerBuilder, batchSize, numWorkers, importResult, dataMigrationStatusService);
    }

    private static int getNumWorkers(int numWorkersFromOptions) {
        int ret = (numWorkersFromOptions > 0) ? numWorkersFromOptions : 1;

        LOG.info("Migration Import: Setting numWorkers: {}", ret);

        return ret;
    }

    private AtlasEntityStoreV2 createEntityStore(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
        FullTextMapperV2Nop        fullTextMapperV2     = new FullTextMapperV2Nop();
        IAtlasEntityChangeNotifier entityChangeNotifier = new EntityChangeNotifierNop();
        DeleteHandlerDelegate      deleteDelegate       = new DeleteHandlerDelegate(graph, typeRegistry, null);
        AtlasFormatConverters      formatConverters     = new AtlasFormatConverters(typeRegistry);

        AtlasInstanceConverter instanceConverter = new AtlasInstanceConverter(graph, typeRegistry, formatConverters);
        AtlasRelationshipStore relationshipStore = new AtlasRelationshipStoreV2(graph, typeRegistry, deleteDelegate, entityChangeNotifier);
        EntityGraphMapper      entityGraphMapper = new EntityGraphMapper(deleteDelegate, typeRegistry, graph, relationshipStore, entityChangeNotifier, instanceConverter, fullTextMapperV2, null);

        return new AtlasEntityStoreV2(graph, deleteDelegate, typeRegistry, entityChangeNotifier, entityGraphMapper);
    }

    private void shutdownEntityCreationManager(EntityCreationManager creationManager) {
        try {
            creationManager.shutdown();
        } catch (InterruptedException e) {
            LOG.error("Migration Import: Shutdown: Interrupted!", e);
        }
    }
}
