/**
 * 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.AtlasException;
import org.apache.atlas.RequestContext;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.audit.EntityAuditEventV2;
import org.apache.atlas.model.tasks.AtlasTask;
import org.apache.atlas.repository.Constants.AtlasAuditAgingType;
import org.apache.atlas.repository.audit.EntityAuditRepository;
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.commons.collections4.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.Arrays;
import java.util.Collection;
import java.util.HashMap;
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_ACTION_TYPES_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_COUNT_KEY;
import static org.apache.atlas.repository.Constants.AUDIT_AGING_TTL_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.CREATE_EVENTS_AGEOUT_ALLOWED_KEY;
import static org.apache.atlas.repository.Constants.PROPERTY_KEY_AUDIT_REDUCTION_NAME;
import static org.apache.atlas.repository.store.graph.v2.tasks.AuditReductionTaskFactory.AGING_TYPE_PROPERTY_KEY_MAP;

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

    private static final int GUID_BATCH_SIZE_PER_AGE_OUT_TASK = 100;

    private final EntityAuditRepository auditRepository;
    private final AtlasGraph            graph;

    public AuditReductionTask(AtlasTask task, EntityAuditRepository auditRepository, AtlasGraph graph) {
        super(task);

        this.auditRepository = auditRepository;
        this.graph           = graph;
    }

    @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();
    }

    public AtlasVertex findVertex() {
        AtlasGraphQuery       query   = graph.query().has(PROPERTY_KEY_AUDIT_REDUCTION_NAME, AUDIT_REDUCTION_TYPE_NAME);
        Iterator<AtlasVertex> results = query.vertices().iterator();

        return results.hasNext() ? results.next() : null;
    }

    protected void run(Map<String, Object> parameters) throws AtlasBaseException, IOException, AtlasException {
        AtlasVertex vertex = findVertex();

        if (vertex == null) {
            return;
        }

        Map<String, List<EntityAuditEventV2>> entitiesWithSucceededAgeout = new HashMap<>();

        AtlasAuditAgingType auditAgingType            = AtlasAuditAgingType.valueOf(String.valueOf(parameters.get(AUDIT_AGING_TYPE_KEY)));
        Set<String>         actionTypes               = new HashSet<>(((Collection<String>) parameters.get(AUDIT_AGING_ACTION_TYPES_KEY)));
        int                 auditCountInput           = (int) parameters.get(AUDIT_AGING_COUNT_KEY);
        short               auditCount                = auditCountInput > Short.MAX_VALUE ? Short.MAX_VALUE : auditCountInput < Short.MIN_VALUE ? Short.MIN_VALUE : (short) auditCountInput;
        int                 ttl                       = (int) parameters.get(AUDIT_AGING_TTL_KEY);
        boolean             createEventsAgeoutAllowed = (boolean) parameters.get(CREATE_EVENTS_AGEOUT_ALLOWED_KEY);
        String              vertexPropertyKeyForGuids = AGING_TYPE_PROPERTY_KEY_MAP.get(auditAgingType);

        List<String> entityGuidsEligibleForAuditAgeout = vertex.getProperty(vertexPropertyKeyForGuids, List.class);
        int          guidsCount                        = CollectionUtils.isNotEmpty(entityGuidsEligibleForAuditAgeout) ? entityGuidsEligibleForAuditAgeout.size() : 0;
        int          batchIndex                        = 1;

        Set<EntityAuditEventV2.EntityAuditActionV2> entityAuditActions = actionTypes.stream().map(EntityAuditEventV2.EntityAuditActionV2::fromString).collect(Collectors.toSet());

        for (int startIndex = 0; startIndex < guidsCount; ) {
            int          endIndex   = startIndex + GUID_BATCH_SIZE_PER_AGE_OUT_TASK < guidsCount ? startIndex + GUID_BATCH_SIZE_PER_AGE_OUT_TASK : guidsCount;
            List<String> guidsBatch = entityGuidsEligibleForAuditAgeout.subList(startIndex, endIndex);

            for (String guid : guidsBatch) {
                List<EntityAuditEventV2> deletedAuditEvents = auditRepository.deleteEventsV2(guid, entityAuditActions, auditCount, ttl, createEventsAgeoutAllowed, auditAgingType);

                entitiesWithSucceededAgeout.put(guid, deletedAuditEvents);
            }

            if (LOG.isDebugEnabled()) {
                LOG.debug("{} Audit aging completed for batch-{} with guids: {}", auditAgingType, batchIndex, Arrays.toString(entitiesWithSucceededAgeout.keySet().toArray()));
            }

            entitiesWithSucceededAgeout.clear();

            startIndex = endIndex;

            batchIndex++;

            List<String> remainingGuids = startIndex < guidsCount ? new ArrayList<>(entityGuidsEligibleForAuditAgeout.subList(startIndex, guidsCount)) : null;

            vertex.setProperty(vertexPropertyKeyForGuids, remainingGuids);
        }
    }
}
