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

import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasConstants;
import org.apache.atlas.kafka.NotificationProvider;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.atlas.model.notification.HookNotification;
import org.apache.atlas.model.notification.MessageSource;
import org.apache.atlas.notification.NotificationException;
import org.apache.atlas.notification.NotificationInterface;
import org.apache.atlas.utils.AtlasConfigurationUtil;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.util.ShutdownHookManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.regex.Pattern;

/**
 * A base class for atlas hooks.
 */
public abstract class AtlasHook {
    private static final Logger LOG = LoggerFactory.getLogger(AtlasHook.class);

    public static final String ATLAS_NOTIFICATION_ASYNCHRONOUS                    = "atlas.notification.hook.asynchronous";
    public static final String ATLAS_NOTIFICATION_ASYNCHRONOUS_MIN_THREADS        = "atlas.notification.hook.asynchronous.minThreads";
    public static final String ATLAS_NOTIFICATION_ASYNCHRONOUS_MAX_THREADS        = "atlas.notification.hook.asynchronous.maxThreads";
    public static final String ATLAS_NOTIFICATION_ASYNCHRONOUS_KEEP_ALIVE_TIME_MS = "atlas.notification.hook.asynchronous.keepAliveTimeMs";
    public static final String ATLAS_NOTIFICATION_ASYNCHRONOUS_QUEUE_SIZE         = "atlas.notification.hook.asynchronous.queueSize";
    public static final String ATLAS_NOTIFICATION_MAX_RETRIES                     = "atlas.notification.hook.retry.maxRetries";
    public static final String ATLAS_NOTIFICATION_RETRY_INTERVAL                  = "atlas.notification.hook.retry.interval";
    public static final String ATLAS_NOTIFICATION_FAILED_MESSAGES_FILENAME_KEY    = "atlas.notification.failed.messages.filename";
    public static final String ATLAS_NOTIFICATION_LOG_FAILED_MESSAGES_ENABLED_KEY = "atlas.notification.log.failed.messages";
    public static final String ATLAS_HOOK_FAILED_MESSAGES_LOG_DEFAULT_NAME        = "atlas_hook_failed_messages.log";
    public static final String CONF_METADATA_NAMESPACE                            = "atlas.metadata.namespace";
    public static final String CLUSTER_NAME_KEY                                   = "atlas.cluster.name";
    public static final String DEFAULT_CLUSTER_NAME                               = "primary";
    public static final String CONF_ATLAS_HOOK_MESSAGES_SORT_ENABLED              = "atlas.hook.messages.sort.enabled";
    public static final String ATLAS_HOOK_ENTITY_IGNORE_PATTERN                   = "atlas.hook.entity.ignore.pattern";
    public static final String ATTRIBUTE_QUALIFIED_NAME                           = "qualifiedName";

    public static final String ATTRIBUTE_INPUTS                                   = "inputs";
    public static final String ATTRIBUTE_OUTPUTS                                  = "outputs";

    public static final boolean isRESTNotificationEnabled;
    public static final boolean isHookMsgsSortEnabled;

    protected static final Configuration         atlasProperties;
    protected static final NotificationInterface notificationInterface;

    private static final String               metadataNamespace;
    private static final int                  SHUTDOWN_HOOK_WAIT_TIME_MS = 3000;
    private static final boolean              logFailedMessages;
    private static final FailedMessagesLogger failedMessagesLogger;
    private static final int                  notificationMaxRetries;
    private static final int                  notificationRetryInterval;
    private static final List<Pattern>        entitiesToIgnore = new ArrayList<>();
    private static final boolean              shouldPreprocess;
    private static final ExecutorService      executor;

    protected MessageSource source;

    public AtlasHook() {
        source = new MessageSource(getMessageSource());

        notificationInterface.init(this.getClass().getSimpleName(), failedMessagesLogger);
    }

    public AtlasHook(String name) {
        source = new MessageSource(getMessageSource());

        LOG.info("AtlasHook: Spool name: Passed from caller.: {}", name);

        notificationInterface.init(name, failedMessagesLogger);
    }

    /**
     * Notify atlas of the entity through message. The entity can be a
     * complex entity with reference to other entities.
     * De-duping of entities is done on server side depending on the
     * unique attribute on the entities.
     *
     * @param messages   hook notification messages
     * @param maxRetries maximum number of retries while sending message to messaging system
     */
    public static void notifyEntities(List<HookNotification> messages, UserGroupInformation ugi, int maxRetries, MessageSource source) {
        if (executor == null) { // send synchronously
            notifyEntitiesPostPreprocess(messages, ugi, maxRetries, source);
        } else {
            executor.submit(() -> notifyEntitiesPostPreprocess(messages, ugi, maxRetries, source));
        }
    }

    /**
     * Returns the logged in user.
     *
     * @return
     */
    public static String getUser() {
        return getUser(null, null);
    }

    public static String getUser(String userName) {
        return getUser(userName, null);
    }

    /**
     * Returns the user. Order of preference:
     * 1. Given userName
     * 2. ugi.getShortUserName()
     * 3. UserGroupInformation.getCurrentUser().getShortUserName()
     * 4. System.getProperty("user.name")
     */

    public static String getUser(String userName, UserGroupInformation ugi) {
        if (StringUtils.isNotEmpty(userName)) {
            LOG.debug("Returning userName {}", userName);

            return userName;
        }

        if (ugi != null && StringUtils.isNotEmpty(ugi.getShortUserName())) {
            LOG.debug("Returning ugi.getShortUserName {}", userName);

            return ugi.getShortUserName();
        }

        try {
            return UserGroupInformation.getCurrentUser().getShortUserName();
        } catch (IOException e) {
            LOG.warn("Failed for UserGroupInformation.getCurrentUser() ", e);

            return System.getProperty("user.name");
        }
    }

    protected static boolean isMatch(String qualifiedName, List<Pattern> patterns) {
        return patterns.stream().anyMatch((Pattern pattern) -> pattern.matcher(qualifiedName).matches());
    }

    @VisibleForTesting
    static void notifyEntitiesInternal(List<HookNotification> messages, int maxRetries, UserGroupInformation ugi, NotificationInterface notificationInterface, boolean shouldLogFailedMessages, FailedMessagesLogger logger, MessageSource source) {
        if (messages == null || messages.isEmpty()) {
            return;
        }

        final int maxAttempts         = maxRetries < 1 ? 1 : maxRetries;
        Exception notificationFailure = null;

        for (int numAttempt = 1; numAttempt <= maxAttempts; numAttempt++) {
            if (numAttempt > 1) { // retry attempt
                try {
                    LOG.debug("Sleeping for {} ms before retry", notificationRetryInterval);

                    Thread.sleep(notificationRetryInterval);
                } catch (InterruptedException ie) {
                    LOG.error("Notification hook thread sleep interrupted");

                    break;
                }
            }

            try {
                if (ugi == null) {
                    notificationInterface.send(NotificationInterface.NotificationType.HOOK, messages, source);
                } else {
                    PrivilegedExceptionAction<Object> privilegedNotify = () -> {
                        notificationInterface.send(NotificationInterface.NotificationType.HOOK, messages, source);

                        return messages;
                    };

                    ugi.doAs(privilegedNotify);
                }

                notificationFailure = null; // notification sent successfully, reset error

                break;
            } catch (Exception e) {
                notificationFailure = e;

                LOG.error("Failed to send notification - attempt #{}; error={}", numAttempt, e.getMessage());
            }
        }

        if (notificationFailure != null) {
            if (shouldLogFailedMessages && notificationFailure instanceof NotificationException) {
                final List<String> failedMessages = ((NotificationException) notificationFailure).getFailedMessages();

                for (String msg : failedMessages) {
                    logger.log(msg);
                }
            }

            LOG.error("Giving up after {} failed attempts to send notification to Atlas: {}", maxAttempts, messages, notificationFailure);
        }
    }

    public abstract String getMessageSource();

    public String getMetadataNamespace() {
        return metadataNamespace;
    }

    /**
     * Notify atlas of the entity through message. The entity can be a
     * complex entity with reference to other entities.
     * De-duping of entities is done on server side depending on the
     * unique attribute on the entities.
     *
     * @param messages hook notification messages
     */
    protected void notifyEntities(List<HookNotification> messages, UserGroupInformation ugi) {
        notifyEntities(messages, ugi, notificationMaxRetries, source);
    }

    private static AtlasEntity.AtlasEntitiesWithExtInfo getAtlasEntitiesWithExtInfo(HookNotification hookNotification) {
        final AtlasEntity.AtlasEntitiesWithExtInfo entitiesWithExtInfo;

        switch (hookNotification.getType()) {
            case ENTITY_CREATE_V2:
                entitiesWithExtInfo = ((HookNotification.EntityCreateRequestV2) hookNotification).getEntities();
                break;
            case ENTITY_FULL_UPDATE_V2:
                entitiesWithExtInfo = ((HookNotification.EntityUpdateRequestV2) hookNotification).getEntities();
                break;
            default:
                entitiesWithExtInfo = null;
        }

        return entitiesWithExtInfo;
    }

    public static String getQualifiedName(Object obj) {
        Map<String, Object> attributes = null;

        if (obj instanceof AtlasObjectId) {
            attributes = ((AtlasObjectId) obj).getUniqueAttributes();
        } else if (obj instanceof Map) {
            attributes = (Map) ((Map) obj).get(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES);
        } else if (obj instanceof AtlasEntity) {
            attributes = ((AtlasEntity) obj).getAttributes();
        } else if (obj instanceof AtlasEntity.AtlasEntityWithExtInfo) {
            attributes = ((AtlasEntity.AtlasEntityWithExtInfo) obj).getEntity().getAttributes();
        }

        Object ret = attributes != null ? attributes.get(ATTRIBUTE_QUALIFIED_NAME) : null;

        return ret != null ? ret.toString() : null;
    }

    private static void filterProcessRelatedEntities(Object obj) {
        if (obj == null || !(obj instanceof Collection)) {
            return;
        }

        Collection objList = (Collection) obj;
        List toRemove = new ArrayList();

        for (Object entity : objList) {
            String qualifiedName = getQualifiedName(entity);

            if (isMatch(qualifiedName, entitiesToIgnore)) {
                toRemove.add(entity);

                LOG.info("Ignored entity {}", qualifiedName);
            }
        }

        objList.removeAll(toRemove);
    }

    private static void filterRelationshipAttributes(Map<String, Object> relationshipAttributes) {
        if (relationshipAttributes == null) {
            return;
        }

        List<String> keysToRemove = new ArrayList<>();
        for (Map.Entry<String, Object> entry : relationshipAttributes.entrySet()) {
            Object obj = entry.getValue();

            if (obj instanceof Collection) {
                Collection entities = (Collection) obj;

                entities.removeIf((Object entity) -> {
                    String qualifiedName = getQualifiedName(entity);

                    return qualifiedName != null && isMatch(qualifiedName, entitiesToIgnore);
                });
            } else {
                String qualifiedName = getQualifiedName(obj);

                if (qualifiedName != null && isMatch(qualifiedName, entitiesToIgnore)) {
                    keysToRemove.add(entry.getKey());
                }
            }
        }

        for (String key : keysToRemove) {
            relationshipAttributes.remove(key);

            LOG.info("Ignored entity {}", key);
        }
    }

    private static void filterEntityAttributes(AtlasEntity entity) {
        Object inputs = entity.getAttribute(ATTRIBUTE_INPUTS);
        Object outputs = entity.getAttribute(ATTRIBUTE_OUTPUTS);

        filterProcessRelatedEntities(inputs);
        filterProcessRelatedEntities(outputs);

        filterRelationshipAttributes(entity.getRelationshipAttributes());
    }

    private static void preprocessEntities(List<HookNotification> hookNotifications) {
        for (int i = 0; i < hookNotifications.size(); i++) {
            HookNotification hookNotification = hookNotifications.get(i);

            AtlasEntity.AtlasEntitiesWithExtInfo entitiesWithExtInfo = getAtlasEntitiesWithExtInfo(hookNotification);

            if (entitiesWithExtInfo == null) {
                continue;
            }

            List<AtlasEntity> entities = entitiesWithExtInfo.getEntities();

            entities = ((entities != null) ? entities : Collections.emptyList());

            entities.removeIf((AtlasEntity entity) -> isMatch(entity.getAttribute(ATTRIBUTE_QUALIFIED_NAME).toString(), entitiesToIgnore));

            for (AtlasEntity entity : entities) {
                filterEntityAttributes(entity);
            }

            Map<String, AtlasEntity> referredEntitiesMap = entitiesWithExtInfo.getReferredEntities();

            referredEntitiesMap = ((referredEntitiesMap != null) ? referredEntitiesMap : Collections.emptyMap());

            referredEntitiesMap.entrySet().removeIf((Map.Entry<String, AtlasEntity> entry) -> isMatch(entry.getValue().getAttribute(ATTRIBUTE_QUALIFIED_NAME).toString(), entitiesToIgnore));

            for (Map.Entry<String, AtlasEntity> entry : referredEntitiesMap.entrySet()) {
                filterEntityAttributes(entry.getValue());
            }

            if (CollectionUtils.isEmpty(entities) && CollectionUtils.isEmpty(referredEntitiesMap.values())) {
                hookNotifications.remove(i--);

                LOG.info("ignored message: {}", hookNotification);
            }
        }
    }

    private static void notifyEntitiesPostPreprocess(List<HookNotification> messages, UserGroupInformation ugi, int maxRetries, MessageSource source) {
        if (shouldPreprocess) {
            preprocessEntities(messages);
        }

        if (CollectionUtils.isNotEmpty(messages)) {
            notifyEntitiesInternal(messages, maxRetries, ugi, notificationInterface, logFailedMessages, failedMessagesLogger, source);
        }
    }

    private static String getMetadataNamespace(Configuration config) {
        return AtlasConfigurationUtil.getRecentString(config, CONF_METADATA_NAMESPACE, getClusterName(config));
    }

    private static String getClusterName(Configuration config) {
        return config.getString(CLUSTER_NAME_KEY, DEFAULT_CLUSTER_NAME);
    }

    static {
        Configuration conf = null;

        try {
            conf = ApplicationProperties.get();
        } catch (Exception e) {
            LOG.info("Failed to load application properties", e);
        }

        atlasProperties   = conf;
        logFailedMessages = atlasProperties.getBoolean(ATLAS_NOTIFICATION_LOG_FAILED_MESSAGES_ENABLED_KEY, true);

        if (logFailedMessages) {
            failedMessagesLogger = new FailedMessagesLogger();
        } else {
            failedMessagesLogger = null;
        }

        isRESTNotificationEnabled = AtlasConfiguration.NOTIFICATION_HOOK_REST_ENABLED.getBoolean();
        isHookMsgsSortEnabled     = atlasProperties.getBoolean(CONF_ATLAS_HOOK_MESSAGES_SORT_ENABLED, isRESTNotificationEnabled);
        metadataNamespace         = getMetadataNamespace(atlasProperties);
        notificationMaxRetries    = atlasProperties.getInt(ATLAS_NOTIFICATION_MAX_RETRIES, 3);
        notificationRetryInterval = atlasProperties.getInt(ATLAS_NOTIFICATION_RETRY_INTERVAL, 1000);
        notificationInterface     = NotificationProvider.get();

        String[] patternsToIgnoreEntities = atlasProperties.getStringArray(ATLAS_HOOK_ENTITY_IGNORE_PATTERN);

        if (patternsToIgnoreEntities != null) {
            for (String pattern : patternsToIgnoreEntities) {
                try {
                    entitiesToIgnore.add(Pattern.compile(pattern));
                } catch (Throwable t) {
                    LOG.warn("failed to compile pattern {}", pattern, t);
                    LOG.warn("Ignoring invalid pattern in configuration {}: {}", ATLAS_HOOK_ENTITY_IGNORE_PATTERN, pattern);
                }
            }

            LOG.info("{}={}", ATLAS_HOOK_ENTITY_IGNORE_PATTERN, entitiesToIgnore);
        }

        shouldPreprocess = CollectionUtils.isNotEmpty(entitiesToIgnore);

        String currentUser = "";

        try {
            currentUser = getUser();
        } catch (Exception excp) {
            LOG.warn("Error in determining current user", excp);
        }

        notificationInterface.setCurrentUser(currentUser);

        boolean isAsync = atlasProperties.getBoolean(ATLAS_NOTIFICATION_ASYNCHRONOUS, Boolean.TRUE);

        if (isAsync) {
            int  minThreads      = atlasProperties.getInt(ATLAS_NOTIFICATION_ASYNCHRONOUS_MIN_THREADS, 1);
            int  maxThreads      = atlasProperties.getInt(ATLAS_NOTIFICATION_ASYNCHRONOUS_MAX_THREADS, 1);
            long keepAliveTimeMs = atlasProperties.getLong(ATLAS_NOTIFICATION_ASYNCHRONOUS_KEEP_ALIVE_TIME_MS, 10000);
            int  queueSize       = atlasProperties.getInt(ATLAS_NOTIFICATION_ASYNCHRONOUS_QUEUE_SIZE, 10000);

            executor = new ThreadPoolExecutor(minThreads, maxThreads, keepAliveTimeMs, TimeUnit.MILLISECONDS,
                    new LinkedBlockingDeque<>(queueSize),
                    new ThreadFactoryBuilder().setNameFormat("Atlas Notifier %d").setDaemon(true).build());

            ShutdownHookManager.get().addShutdownHook(new Thread() {
                @Override
                public void run() {
                    try {
                        LOG.info("==> Shutdown of Atlas Hook");

                        notificationInterface.close();
                        executor.shutdown();

                        boolean ignored = executor.awaitTermination(SHUTDOWN_HOOK_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
                    } catch (InterruptedException excp) {
                        LOG.info("Interrupt received in shutdown.", excp);
                    } finally {
                        LOG.info("<== Shutdown of Atlas Hook");
                    }
                }
            }, AtlasConstants.ATLAS_SHUTDOWN_HOOK_PRIORITY);
        } else {
            executor = null;
        }

        LOG.info("Created Atlas Hook");
    }
}
