/**
 * 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.RequestContext;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.glossary.GlossaryService;
import org.apache.atlas.model.impexp.AtlasExportRequest;
import org.apache.atlas.model.impexp.AtlasExportResult;
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.typedef.AtlasBusinessMetadataDef;
import org.apache.atlas.model.typedef.AtlasClassificationDef;
import org.apache.atlas.model.typedef.AtlasEntityDef;
import org.apache.atlas.model.typedef.AtlasEnumDef;
import org.apache.atlas.model.typedef.AtlasRelationshipDef;
import org.apache.atlas.model.typedef.AtlasStructDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.graph.GraphHelper;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
import org.apache.atlas.repository.util.UniqueList;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.inject.Inject;

import java.util.ArrayList;
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.FETCH_TYPE_CONNECTED;
import static org.apache.atlas.model.impexp.AtlasExportRequest.FETCH_TYPE_FULL;
import static org.apache.atlas.model.impexp.AtlasExportRequest.FETCH_TYPE_INCREMENTAL;
import static org.apache.atlas.repository.Constants.GUID_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY;

@Component
public class ExportService {
    private static final Logger LOG = LoggerFactory.getLogger(ExportService.class);

    private final AtlasTypeRegistry               typeRegistry;
    private final StartEntityFetchByExportRequest startEntityFetchByExportRequest;
    private final EntitiesExtractor               entitiesExtractor;
    private final EntityGraphRetriever            entityGraphRetriever;
    private final HdfsPathEntityCreator           hdfsPathEntityCreator;
    private final GlossaryService                 glossaryService;
    private final AuditsWriter                    auditsWriter;
    private       ExportTypeProcessor             exportTypeProcessor;
    private static final String                   ATLAS_TYPE_HIVE_DB = "hive_db";

    @Inject
    public ExportService(final AtlasTypeRegistry typeRegistry, AtlasGraph graph, AuditsWriter auditsWriter, HdfsPathEntityCreator hdfsPathEntityCreator, GlossaryService glossaryService) {
        this.typeRegistry                    = typeRegistry;
        this.entityGraphRetriever            = new EntityGraphRetriever(graph, this.typeRegistry);
        this.auditsWriter                    = auditsWriter;
        this.hdfsPathEntityCreator           = hdfsPathEntityCreator;
        this.glossaryService                 = glossaryService;
        this.startEntityFetchByExportRequest = new StartEntityFetchByExportRequest(graph, typeRegistry, AtlasGremlinQueryProvider.getInstance());
        this.entitiesExtractor               = new EntitiesExtractor(graph, typeRegistry);
    }

    public AtlasExportResult run(ZipSink exportSink, AtlasExportRequest request, String userName, String hostName, String requestingIP) throws AtlasBaseException {
        long              startTime = System.currentTimeMillis();
        AtlasExportResult result    = new AtlasExportResult(request, userName, requestingIP, hostName, startTime, getCurrentChangeMarker());
        ExportContext     context   = new ExportContext(result, exportSink);

        exportTypeProcessor = new ExportTypeProcessor(typeRegistry, glossaryService);

        try {
            LOG.info("==> export(user={}, from={})", userName, requestingIP);

            AtlasExportResult.OperationStatus[] statuses = processItems(request, context);

            processTypesDef(context);

            long endTime = System.currentTimeMillis();

            updateSinkWithOperationMetrics(userName, context, statuses, startTime, endTime);
        } catch (Exception ex) {
            LOG.error("Operation failed: ", ex);
        } finally {
            entitiesExtractor.close();

            LOG.info("<== export(user={}, from={}): status {}: changeMarker: {}", userName, requestingIP, context.result.getOperationStatus(), context.result.getChangeMarker());

            context.clear();
            result.clear();
        }

        return context.result;
    }

    public void processEntity(AtlasEntityWithExtInfo entityWithExtInfo, ExportContext context) throws AtlasBaseException {
        exportTypeProcessor.addTypes(entityWithExtInfo.getEntity(), context);

        if (MapUtils.isNotEmpty(context.termsGlossary)) {
            addGlossaryEntities(context);
        }

        addEntity(entityWithExtInfo, context);

        context.guidsProcessed.add(entityWithExtInfo.getEntity().getGuid());

        entitiesExtractor.get(entityWithExtInfo.getEntity(), context);

        if (entityWithExtInfo.getReferredEntities() != null) {
            for (AtlasEntity e : entityWithExtInfo.getReferredEntities().values()) {
                exportTypeProcessor.addTypes(e, context);
                entitiesExtractor.get(e, context);
            }

            context.guidsProcessed.addAll(entityWithExtInfo.getReferredEntities().keySet());
        }
    }

    @VisibleForTesting
    AtlasExportResult.OperationStatus getOverallOperationStatus(AtlasExportResult.OperationStatus... statuses) {
        AtlasExportResult.OperationStatus overall = (statuses.length == 0) ? AtlasExportResult.OperationStatus.FAIL : statuses[0];

        for (AtlasExportResult.OperationStatus s : statuses) {
            if (overall != s) {
                overall = AtlasExportResult.OperationStatus.PARTIAL_SUCCESS;
            }
        }

        return overall;
    }

    private long getCurrentChangeMarker() {
        return RequestContext.earliestActiveRequestTime();
    }

    private void updateSinkWithOperationMetrics(String userName, ExportContext context, AtlasExportResult.OperationStatus[] statuses, long startTime, long endTime) throws AtlasBaseException {
        int duration = getOperationDuration(startTime, endTime);

        context.result.setSourceClusterName(AuditsWriter.getCurrentClusterName());

        context.sink.setExportOrder(context.entityCreationOrder.getList());
        context.sink.setTypesDef(context.result.getData().getTypesDef());

        context.result.setOperationStatus(getOverallOperationStatus(statuses));
        context.result.incrementMeticsCounter("duration", duration);

        auditsWriter.write(userName, context.result, startTime, endTime, context.entityCreationOrder.getList());

        context.result.setData(null);
        context.sink.setResult(context.result);
    }

    private int getOperationDuration(long startTime, long endTime) {
        return (int) (endTime - startTime);
    }

    private void processTypesDef(ExportContext context) {
        AtlasTypesDef typesDef = context.result.getData().getTypesDef();

        for (String entityType : context.entityTypes) {
            AtlasEntityDef entityDef = typeRegistry.getEntityDefByName(entityType);

            typesDef.getEntityDefs().add(entityDef);
        }

        for (String classificationType : context.classificationTypes) {
            AtlasClassificationDef classificationDef = typeRegistry.getClassificationDefByName(classificationType);

            typesDef.getClassificationDefs().add(classificationDef);
        }

        for (String structType : context.structTypes) {
            AtlasStructDef structDef = typeRegistry.getStructDefByName(structType);

            typesDef.getStructDefs().add(structDef);
        }

        for (String enumType : context.enumTypes) {
            AtlasEnumDef enumDef = typeRegistry.getEnumDefByName(enumType);

            typesDef.getEnumDefs().add(enumDef);
        }

        for (String relationshipType : context.relationshipTypes) {
            AtlasRelationshipDef relationshipDef = typeRegistry.getRelationshipDefByName(relationshipType);

            typesDef.getRelationshipDefs().add(relationshipDef);
        }

        for (String bm : context.businessMetadataTypes) {
            AtlasBusinessMetadataDef bmDef = typeRegistry.getBusinessMetadataDefByName(bm);

            typesDef.getBusinessMetadataDefs().add(bmDef);
        }
    }

    private AtlasExportResult.OperationStatus[] processItems(AtlasExportRequest request, ExportContext context) {
        AtlasExportResult.OperationStatus[] statuses      = new AtlasExportResult.OperationStatus[request.getItemsToExport().size()];
        List<AtlasObjectId>                 itemsToExport = request.getItemsToExport();

        for (int i = 0; i < itemsToExport.size(); i++) {
            AtlasObjectId item = itemsToExport.get(i);

            statuses[i] = processObjectId(item, context);
        }

        return statuses;
    }

    private AtlasExportResult.OperationStatus processObjectId(AtlasObjectId item, ExportContext context) {
        LOG.debug("==> processObjectId({})", item);

        try {
            List<String> entityGuids = getStartingEntity(item, context);

            if (entityGuids.isEmpty()) {
                return AtlasExportResult.OperationStatus.FAIL;
            }

            entitiesExtractor.setExtractor(typeRegistry.getEntityDefByName(item.getTypeName()));

            for (String guid : entityGuids) {
                AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(guid);
                String typeName = GraphHelper.getTypeName(vertex);
                context.startingEntityType = typeName;
                processEntityGuid(guid, context);
            }

            while (!context.guidsToProcess.isEmpty()) {
                while (!context.guidsToProcess.isEmpty()) {
                    String guid = context.guidsToProcess.remove(0);

                    processEntityGuid(guid, context);
                }

                if (!context.lineageToProcess.isEmpty()) {
                    context.guidsToProcess.addAll(context.lineageToProcess);
                    context.lineageProcessed.addAll(context.lineageToProcess.getList());
                    context.lineageToProcess.clear();
                }

                context.isSkipConnectedFetch = false;
            }
        } catch (AtlasBaseException excp) {
            LOG.error("Fetching entity failed for: {}", item, excp);

            return AtlasExportResult.OperationStatus.FAIL;
        }

        LOG.debug("<== processObjectId({})", item);

        return AtlasExportResult.OperationStatus.SUCCESS;
    }

    private List<String> getStartingEntity(AtlasObjectId item, ExportContext context) throws AtlasBaseException {
        if (item.getTypeName().equalsIgnoreCase(HdfsPathEntityCreator.HDFS_PATH_TYPE)) {
            hdfsPathEntityCreator.getCreateEntity(item);
        }

        return startEntityFetchByExportRequest.get(context.result.getRequest(), item);
    }

    private void processEntityGuid(String guid, ExportContext context) throws AtlasBaseException {
        LOG.debug("==> processEntityGuid({})", guid);

        if (context.guidsProcessed.contains(guid)) {
            return;
        }

        if (context.fetchType == ExportFetchType.INCREMENTAL && context.startingEntityType.equals(ATLAS_TYPE_HIVE_DB) && !context.skipLineage) {
            AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(guid);

            processVertex(context, vertex, guid);
        } else {
            AtlasEntityWithExtInfo entityWithExtInfo = entityGraphRetriever.toAtlasEntityWithExtInfo(guid);

            processEntity(entityWithExtInfo, context);
        }

        LOG.debug("<== processEntityGuid({})", guid);
    }

    public void processVertex(ExportContext context, AtlasVertex vertex, String guid) throws AtlasBaseException {
        if (MapUtils.isNotEmpty(context.termsGlossary)) {
            addGlossaryEntities(context);
        }

        addVertex(vertex, guid, context);

        context.guidsProcessed.add(guid);

        extractConnectedVertices(vertex, context);
    }

    public void extractConnectedVertices(AtlasVertex vertex, ExportContext context) {
        List<AtlasVertex> connectedVertices = entityGraphRetriever.findAllConnectedVertices(vertex);

        if (CollectionUtils.isNotEmpty(connectedVertices)) {
            for (AtlasVertex e : connectedVertices) {
                String typeName = GraphHelper.getTypeName(e);

                if (typeRegistry.getEntityTypeByName(typeName) != null) {
                    String guid = AtlasGraphUtilsV2.getEncodedProperty(e, GUID_PROPERTY_KEY, String.class);

                    if (!context.guidsProcessed.contains(guid)) {
                        context.guidsToProcess.add(guid);
                    }
                }
            }
        }
    }

    private void addVertex(AtlasVertex vertex, String guid, ExportContext context) throws AtlasBaseException {
        if (context.sink.hasEntity(guid)) {
            return;
        }

        LOG.info("export: Guid in process: {}", guid);
        if (context.doesTimestampQualify(vertex)) {
            AtlasEntityWithExtInfo entityWithExtInfo = entityGraphRetriever.toAtlasEntityWithExtInfo(guid);
            exportTypeProcessor.addTypes(entityWithExtInfo.getEntity(), context);
            context.addToSink(entityWithExtInfo);

            context.result.incrementMeticsCounter(String.format("entity:%s", entityWithExtInfo.getEntity().getTypeName()));
            if (entityWithExtInfo.getReferredEntities() != null) {
                for (AtlasEntity e : entityWithExtInfo.getReferredEntities().values()) {
                    context.result.incrementMeticsCounter(String.format("entity:%s", e.getTypeName()));
                }
            }

            context.result.incrementMeticsCounter("entity:withExtInfo");
        }
        context.reportProgress();
    }

    private void addGlossaryEntities(ExportContext context) {
        try {
            for (String termGuid : context.termsGlossary.keySet()) {
                try {
                    String glossaryGuid = context.termsGlossary.get(termGuid);

                    if (!context.sink.hasEntity(glossaryGuid)) {
                        AtlasEntity glossary = entityGraphRetriever.toAtlasEntity(glossaryGuid);

                        addEntity(new AtlasEntityWithExtInfo(glossary), context);
                    }

                    if (!context.sink.hasEntity(termGuid)) {
                        AtlasEntity term = entityGraphRetriever.toAtlasEntity(termGuid);

                        addEntity(new AtlasEntityWithExtInfo(term), context);
                    }
                } catch (AtlasBaseException exception) {
                    LOG.error("Error fetching Glossary for term: {}", termGuid);
                }
            }
        } finally {
            context.clearTerms();
        }
    }

    private void addEntity(AtlasEntityWithExtInfo entityWithExtInfo, ExportContext context) throws AtlasBaseException {
        if (context.sink.hasEntity(entityWithExtInfo.getEntity().getGuid())) {
            return;
        }

        if (context.doesTimestampQualify(entityWithExtInfo.getEntity())) {
            context.addToSink(entityWithExtInfo);

            context.result.incrementMeticsCounter(String.format("entity:%s", entityWithExtInfo.getEntity().getTypeName()));

            if (entityWithExtInfo.getReferredEntities() != null) {
                for (AtlasEntity e : entityWithExtInfo.getReferredEntities().values()) {
                    context.result.incrementMeticsCounter(String.format("entity:%s", e.getTypeName()));
                }
            }

            context.result.incrementMeticsCounter("entity:withExtInfo");
        } else {
            List<AtlasEntity> entities = context.getEntitiesWithModifiedTimestamp(entityWithExtInfo);

            for (AtlasEntity e : entities) {
                context.addToSink(new AtlasEntityWithExtInfo(e));
                context.result.incrementMeticsCounter(String.format("entity:%s", e.getTypeName()));
            }
        }

        context.reportProgress();
    }

    public enum TraversalDirection {
        UNKNOWN,
        INWARD,
        OUTWARD,
        BOTH
    }

    public enum ExportFetchType {
        FULL(FETCH_TYPE_FULL),
        CONNECTED(FETCH_TYPE_CONNECTED),
        INCREMENTAL(FETCH_TYPE_INCREMENTAL);

        final String str;

        ExportFetchType(String s) {
            this.str = s;
        }

        public static ExportFetchType from(String s) {
            for (ExportFetchType b : ExportFetchType.values()) {
                if (b.str.equalsIgnoreCase(s)) {
                    return b;
                }
            }

            return FULL;
        }
    }

    static class ExportContext {
        private static final int    REPORTING_THREASHOLD  = 1000;
        private static final String ATLAS_TYPE_HIVE_DB    = "hive_db";
        private static final String ATLAS_TYPE_HIVE_TABLE = "hive_table";

        final UniqueList<String>              entityCreationOrder   = new UniqueList<>();
        final Set<String>                     guidsProcessed        = new HashSet<>();
        final UniqueList<String>              guidsToProcess        = new UniqueList<>();
        final UniqueList<String>              lineageToProcess      = new UniqueList<>();
        final Set<String>                     lineageProcessed      = new HashSet<>();
        final Map<String, TraversalDirection> guidDirection         = new HashMap<>();
        final Set<String>                     entityTypes           = new HashSet<>();
        final Set<String>                     classificationTypes   = new HashSet<>();
        final Set<String>                     structTypes           = new HashSet<>();
        final Set<String>                     enumTypes             = new HashSet<>();
        final Set<String>                     relationshipTypes     = new HashSet<>();
        final Set<String>                     businessMetadataTypes = new HashSet<>();
        final Map<String, String>             termsGlossary         = new HashMap<>();
        final AtlasExportResult               result;
        final ExportFetchType                 fetchType;
        final boolean                         skipLineage;
        final long                            changeMarker;

        private final ZipSink sink;
        private final boolean isHiveDBIncremental;
        private final boolean isHiveTableIncremental;

        boolean isSkipConnectedFetch;
        private int progressReportCount;
        public String startingEntityType;

        ExportContext(AtlasExportResult result, ZipSink sink) {
            this.result = result;
            this.sink   = sink;

            fetchType                   = ExportFetchType.from(result.getRequest().getFetchTypeOptionValue());
            skipLineage                 = result.getRequest().getSkipLineageOptionValue();
            this.changeMarker           = result.getRequest().getChangeTokenFromOptions();
            this.isHiveDBIncremental    = checkHiveDBIncrementalSkipLineage(result.getRequest());
            this.isHiveTableIncremental = checkHiveTableIncremental(result.getRequest());
            this.isSkipConnectedFetch   = false;
        }

        public List<AtlasEntity> getEntitiesWithModifiedTimestamp(AtlasEntityWithExtInfo entityWithExtInfo) {
            if (fetchType != ExportFetchType.INCREMENTAL) {
                return new ArrayList<>();
            }

            List<AtlasEntity> ret = new ArrayList<>();
            if (doesTimestampQualify(entityWithExtInfo.getEntity())) {
                ret.add(entityWithExtInfo.getEntity());
                return ret;
            }

            if (entityWithExtInfo.getReferredEntities() != null) {
                for (AtlasEntity entity : entityWithExtInfo.getReferredEntities().values()) {
                    if ((doesTimestampQualify(entity))) {
                        ret.add(entity);
                    }
                }
            }

            return ret;
        }

        public void clear() {
            guidsToProcess.clear();
            guidsProcessed.clear();
            guidDirection.clear();
            startingEntityType = null;
        }

        public void addToBeProcessed(boolean isSuperTypeProcess, String guid, TraversalDirection direction) {
            if (isSuperTypeProcess) {
                lineageToProcess.add(guid);
            } else {
                guidsToProcess.add(guid);
            }

            guidDirection.put(guid, direction);
        }

        public void reportProgress() {
            if ((guidsProcessed.size() - progressReportCount) > REPORTING_THREASHOLD) {
                progressReportCount = guidsProcessed.size();

                LOG.info("export(): in progress.. number of entities exported: {}", this.guidsProcessed.size());
            }
        }

        public boolean doesTimestampQualify(AtlasEntity entity) {
            if (fetchType != ExportFetchType.INCREMENTAL) {
                return true;
            }

            return changeMarker <= entity.getUpdateTime().getTime();
        }

        public boolean doesTimestampQualify(AtlasVertex vertex) {
            if (fetchType != ExportFetchType.INCREMENTAL) {
                return true;
            }

            Long updatedTime = AtlasGraphUtilsV2.getEncodedProperty(vertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class);
            return changeMarker <= updatedTime;
        }

        public boolean getSkipLineage() {
            return skipLineage;
        }

        public void addToSink(AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
            addToEntityCreationOrder(entityWithExtInfo.getEntity().getGuid());
            sink.add(entityWithExtInfo);
        }

        public boolean isHiveDBIncrementalSkipLineage() {
            return isHiveDBIncremental;
        }

        public boolean isHiveTableIncrementalSkipLineage() {
            return isHiveTableIncremental;
        }

        public boolean isHiveTableIncremental() {
            return isHiveTableIncremental;
        }

        public void addToEntityCreationOrder(String guid) {
            entityCreationOrder.add(guid);
        }

        public void clearTerms() {
            termsGlossary.clear();
        }

        private boolean checkHiveDBIncrementalSkipLineage(AtlasExportRequest request) {
            if (CollectionUtils.isEmpty(request.getItemsToExport())) {
                return false;
            }

            return request.getItemsToExport().get(0).getTypeName().equalsIgnoreCase(ATLAS_TYPE_HIVE_DB) &&
                    request.getFetchTypeOptionValue().equalsIgnoreCase(AtlasExportRequest.FETCH_TYPE_INCREMENTAL) &&
                    request.getSkipLineageOptionValue();
        }

        private boolean checkHiveTableIncremental(AtlasExportRequest request) {
            if (CollectionUtils.isEmpty(request.getItemsToExport())) {
                return false;
            }

            return request.getItemsToExport().get(0).getTypeName().equalsIgnoreCase(ATLAS_TYPE_HIVE_TABLE) &&
                    request.getFetchTypeOptionValue().equalsIgnoreCase(AtlasExportRequest.FETCH_TYPE_INCREMENTAL);
        }
    }
}
