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

import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasException;
import org.apache.atlas.RequestContext;
import org.apache.atlas.annotation.GraphTransaction;
import org.apache.atlas.discovery.AtlasDiscoveryService;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.discovery.SearchParameters;
import org.apache.atlas.model.tasks.AtlasTask;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.tasks.AbstractTask;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static org.apache.atlas.model.tasks.AtlasTask.Status.COMPLETE;
import static org.apache.atlas.model.tasks.AtlasTask.Status.FAILED;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_ENTITY_TYPES_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_EXCLUDE_ENTITY_TYPES_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_SUBTYPES_INCLUDED_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_TYPE_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_REDUCTION_TYPE_NAME;
import static org.apache.atlas.repository.Constants.AtlasAuditAgingType;
import static org.apache.atlas.repository.Constants.PROPERTY_KEY_AUDIT_REDUCTION_NAME;
import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.setEncodedProperty;
import static org.apache.atlas.repository.store.graph.v2.tasks.AuditReductionTaskFactory.AGING_TYPE_PROPERTY_KEY_MAP;
import static org.apache.atlas.repository.store.graph.v2.tasks.AuditReductionTaskFactory.ATLAS_AUDIT_REDUCTION;

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

    private static final String VALUE_DELIMITER  = ",";
    private static final String ALL_ENTITY_TYPES = "_ALL_ENTITY_TYPES";
    private static final int    SEARCH_OFFSET    = 0;
    private static final int    SEARCH_LIMIT     = AtlasConfiguration.ATLAS_AUDIT_AGING_SEARCH_MAX_LIMIT.getInt();

    private final AtlasDiscoveryService discoveryService;
    private final AtlasTypeRegistry     typeRegistry;
    private final AtlasGraph            graph;

    public AuditReductionEntityRetrievalTask(AtlasTask task, AtlasGraph graph, AtlasDiscoveryService discoveryService, AtlasTypeRegistry typeRegistry) {
        super(task);
        this.graph            = graph;
        this.discoveryService = discoveryService;
        this.typeRegistry     = typeRegistry;
    }

    @Override
    public AtlasTask.Status perform() throws Exception {
        RequestContext.clear();

        Map<String, Object> params = getTaskDef().getParameters();

        if (MapUtils.isEmpty(params)) {
            LOG.warn("Task: {}: Unable to process task: Parameters is not readable!", getTaskGuid());
            return FAILED;
        }

        String userName = getTaskDef().getCreatedBy();

        if (StringUtils.isEmpty(userName)) {
            LOG.warn("Task: {}: Unable to process task as user name is empty!", getTaskGuid());

            return FAILED;
        }

        RequestContext.get().setUser(userName, null);

        try {
            run(params);

            setStatus(COMPLETE);
        } catch (Exception e) {
            LOG.error("Task: {}: Error performing task!", getTaskGuid(), e);

            setStatus(FAILED);

            throw e;
        } finally {
            RequestContext.clear();
        }

        return getStatus();
    }

    protected void run(Map<String, Object> parameters) throws AtlasBaseException, IOException, AtlasException {
        try {
            AtlasTask auditAgingTask = createAgingTaskWithEligibleGUIDs(parameters);

            if (auditAgingTask != null) {
                LOG.info("{} task created for audit aging type-{}", ATLAS_AUDIT_REDUCTION, parameters.get(AUDIT_AGING_TYPE_KEY));
            }
        } catch (Exception e) {
            LOG.error("Error while retrieving entities eligible for audit aging and creating audit aging tasks", e);
        }
    }

    protected AtlasTask createAgingTaskWithEligibleGUIDs(Map<String, Object> parameters) throws AtlasBaseException {
        Set<String>         entityTypes      = new HashSet<>(((Collection<String>) parameters.get(AUDIT_AGING_ENTITY_TYPES_KEY)));
        AtlasAuditAgingType auditAgingType   = (AtlasAuditAgingType) parameters.get(AUDIT_AGING_TYPE_KEY);
        boolean             subTypesIncluded = (boolean) parameters.get(AUDIT_AGING_SUBTYPES_INCLUDED_KEY);

        SearchParameters searchEntitiesToReduceAudit = new SearchParameters();

        searchEntitiesToReduceAudit.setTypeName(ALL_ENTITY_TYPES);
        searchEntitiesToReduceAudit.setOffset(SEARCH_OFFSET);
        searchEntitiesToReduceAudit.setLimit(SEARCH_LIMIT);
        searchEntitiesToReduceAudit.setIncludeSubTypes(subTypesIncluded);

        if (CollectionUtils.isNotEmpty(entityTypes)) {
            if (!validateTypesAndIncludeSubTypes(entityTypes, auditAgingType, subTypesIncluded)) {
                LOG.error("All entity type names provided for audit aging type-{} are invalid", auditAgingType);

                return null;
            }

            String queryString = String.join(VALUE_DELIMITER, entityTypes);

            if (auditAgingType == AtlasAuditAgingType.DEFAULT && StringUtils.isNotEmpty(queryString)) {
                queryString = "!" + queryString;
            }

            searchEntitiesToReduceAudit.setQuery(queryString);
        }

        LOG.info("Getting GUIDs eligible for Audit aging type-{} with SearchParameters: {}", auditAgingType, searchEntitiesToReduceAudit);

        Set<String> guids = discoveryService.searchGUIDsWithParameters(auditAgingType, entityTypes, searchEntitiesToReduceAudit);

        AtlasVertex auditReductionVertex = getOrCreateVertex();

        AtlasTask ageoutTask = updateVertexWithGuidsAndCreateAgingTask(auditReductionVertex, AGING_TYPE_PROPERTY_KEY_MAP.get(auditAgingType), guids, parameters);

        /** For DEFAULT audit aging, "entityTypes" should be excluded from _ALL_ENTITY_TYPES i.e., negating the queryString
         *  Including AUDIT_AGING_EXCLUDE_ENTITY_TYPES_KEY to indicate the same in AtlasTask response to user
         */
        if (ageoutTask != null) {
            if (auditAgingType == AtlasAuditAgingType.DEFAULT && CollectionUtils.isNotEmpty(entityTypes)) {
                ageoutTask.getParameters().put(AUDIT_AGING_EXCLUDE_ENTITY_TYPES_KEY, true);
            } else {
                ageoutTask.getParameters().put(AUDIT_AGING_EXCLUDE_ENTITY_TYPES_KEY, false);
            }
        }

        return ageoutTask;
    }

    private boolean validateTypesAndIncludeSubTypes(Set<String> entityTypes, AtlasAuditAgingType auditAgingType, boolean subTypesIncluded) throws AtlasBaseException {
        Collection<String> allEntityTypeNames     = typeRegistry.getAllEntityDefNames();
        Set<String>        entityTypesToSearch    = new HashSet<>();
        Set<String>        invalidEntityTypeNames = new HashSet<>();

        entityTypes.forEach(entityType -> {
            if (entityType.endsWith("*")) {
                String suffix = entityType.replace("*", "");

                entityTypesToSearch.addAll(allEntityTypeNames.stream().filter(e -> e.startsWith(suffix)).collect(Collectors.toSet()));
            } else if (allEntityTypeNames.contains(entityType)) {
                entityTypesToSearch.add(entityType);
            } else {
                invalidEntityTypeNames.add(entityType);
            }
        });

        if (auditAgingType != AtlasAuditAgingType.DEFAULT) {
            if (CollectionUtils.isNotEmpty(invalidEntityTypeNames)) {
                LOG.warn("Invalid entity type name(s) {} provided for aging type-{}", String.join(VALUE_DELIMITER, invalidEntityTypeNames), auditAgingType);
            }

            if (CollectionUtils.isEmpty(entityTypesToSearch)) {
                return false;
            }
        }

        entityTypes.clear();
        entityTypes.addAll(subTypesIncluded ? AtlasEntityType.getEntityTypesAndAllSubTypes(entityTypesToSearch, typeRegistry) : entityTypesToSearch);

        return true;
    }

    @GraphTransaction
    private AtlasTask updateVertexWithGuidsAndCreateAgingTask(AtlasVertex vertex, String vertexProperty, Set<String> guids, Map<String, Object> params) throws AtlasBaseException {
        List<String> guidsEligibleForAuditReduction = vertex.getProperty(vertexProperty, List.class);

        if (CollectionUtils.isEmpty(guidsEligibleForAuditReduction) && CollectionUtils.isEmpty(guids)) {
            return null;
        }

        if (CollectionUtils.isEmpty(guidsEligibleForAuditReduction)) {
            guidsEligibleForAuditReduction = new ArrayList<>();
        }

        if (CollectionUtils.isNotEmpty(guids)) {
            guidsEligibleForAuditReduction.addAll(guids);
            setEncodedProperty(vertex, vertexProperty, guidsEligibleForAuditReduction);
        }

        return discoveryService.createAndQueueAuditReductionTask(params, ATLAS_AUDIT_REDUCTION);
    }

    private AtlasVertex getOrCreateVertex() {
        AtlasGraphQuery       query                = graph.query().has(PROPERTY_KEY_AUDIT_REDUCTION_NAME, AUDIT_REDUCTION_TYPE_NAME);
        Iterator<AtlasVertex> results              = query.vertices().iterator();
        AtlasVertex           auditReductionVertex = results.hasNext() ? results.next() : null;

        if (auditReductionVertex == null) {
            auditReductionVertex = graph.addVertex();

            setEncodedProperty(auditReductionVertex, PROPERTY_KEY_AUDIT_REDUCTION_NAME, AUDIT_REDUCTION_TYPE_NAME);
        }

        return auditReductionVertex;
    }
}
