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

import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.security.SecurityUtil;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.CreateTopicsResult;
import org.apache.kafka.clients.admin.DescribeTopicsResult;
import org.apache.kafka.clients.admin.ListTopicsOptions;
import org.apache.kafka.clients.admin.ListTopicsResult;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.admin.TopicDescription;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicPartitionInfo;
import org.apache.kafka.common.errors.TopicExistsException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

import static org.apache.atlas.security.SecurityProperties.HADOOP_SECURITY_CREDENTIAL_PROVIDER_PATH;

public class KafkaUtils implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(KafkaUtils.class);

    public static final String ATLAS_KAFKA_PROPERTY_PREFIX     = "atlas.kafka";
    public static final String KAFKA_SASL_JAAS_CONFIG_PROPERTY = "sasl.jaas.config";
    public static final String JAAS_PASSWORD_SUFFIX            = "password";

    private static final String JAAS_CONFIG_PREFIX_PARAM                     = "atlas.jaas";
    private static final String JAAS_CONFIG_LOGIN_MODULE_NAME_PARAM          = "loginModuleName";
    private static final String JAAS_CONFIG_LOGIN_MODULE_CONTROL_FLAG_PARAM  = "loginModuleControlFlag";
    private static final String JAAS_DEFAULT_LOGIN_MODULE_CONTROL_FLAG       = "required";
    private static final String JAAS_VALID_LOGIN_MODULE_CONTROL_FLAG_OPTIONS = "optional|requisite|sufficient|required";
    private static final String JAAS_CONFIG_LOGIN_OPTIONS_PREFIX             = "option";
    private static final String JAAS_PRINCIPAL_PROP                          = "principal";
    private static final String JAAS_DEFAULT_CLIENT_NAME                     = "KafkaClient";
    private static final String JAAS_TICKET_BASED_CLIENT_NAME                = "ticketBased-KafkaClient";
    private static final String IMPORT_INTERNAL_TOPICS                       = "atlas.kafka.bridge.enable.internal.topics.import";
    private static final String JAAS_MASK_PASSWORD                           = "********";

    protected final Properties  kafkaConfiguration;
    protected final AdminClient adminClient;
    protected final boolean     importInternalTopics;

    public KafkaUtils(Configuration atlasConfiguration) {
        LOG.debug("==> KafkaUtils() ");

        this.kafkaConfiguration = ApplicationProperties.getSubsetAsProperties(atlasConfiguration, ATLAS_KAFKA_PROPERTY_PREFIX);

        setKafkaJAASProperties(atlasConfiguration, kafkaConfiguration);

        adminClient          = AdminClient.create(this.kafkaConfiguration);
        importInternalTopics = atlasConfiguration.getBoolean(IMPORT_INTERNAL_TOPICS, false);

        LOG.debug("<== KafkaUtils() ");
    }

    public static void setKafkaJAASProperties(Configuration configuration, Properties kafkaProperties) {
        LOG.debug("==> KafkaUtils.setKafkaJAASProperties()");

        if (kafkaProperties.containsKey(KAFKA_SASL_JAAS_CONFIG_PROPERTY)) {
            LOG.debug("JAAS config is already set, returning");

            return;
        }

        Properties jaasConfig = ApplicationProperties.getSubsetAsProperties(configuration, JAAS_CONFIG_PREFIX_PARAM);
        // JAAS Configuration is present then update set those properties in sasl.jaas.config

        if (jaasConfig != null && !jaasConfig.isEmpty()) {
            String jaasClientName = JAAS_DEFAULT_CLIENT_NAME;

            // Required for backward compatability for Hive CLI
            if (!isLoginKeytabBased() && isLoginTicketBased()) {
                LOG.debug("Checking if ticketBased-KafkaClient is set");

                // if ticketBased-KafkaClient property is not specified then use the default client name
                String        ticketBasedConfigPrefix = JAAS_CONFIG_PREFIX_PARAM + "." + JAAS_TICKET_BASED_CLIENT_NAME;
                Configuration ticketBasedConfig       = configuration.subset(ticketBasedConfigPrefix);

                if (ticketBasedConfig != null && !ticketBasedConfig.isEmpty()) {
                    LOG.debug("ticketBased-KafkaClient JAAS configuration is set, using it");

                    jaasClientName = JAAS_TICKET_BASED_CLIENT_NAME;
                } else {
                    LOG.info("UserGroupInformation.isLoginTicketBased is true, but no JAAS configuration found for client {}. Will use JAAS configuration of client {}", JAAS_TICKET_BASED_CLIENT_NAME, jaasClientName);
                }
            }

            String keyPrefix       = jaasClientName + ".";
            String keyParam        = keyPrefix + JAAS_CONFIG_LOGIN_MODULE_NAME_PARAM;
            String loginModuleName = jaasConfig.getProperty(keyParam);

            if (loginModuleName == null) {
                LOG.error("Unable to add JAAS configuration for client [{}] as it is missing param [{}]. Skipping JAAS config for [{}]", jaasClientName, keyParam, jaasClientName);

                return;
            }

            keyParam = keyPrefix + JAAS_CONFIG_LOGIN_MODULE_CONTROL_FLAG_PARAM;

            String controlFlag = jaasConfig.getProperty(keyParam);

            if (StringUtils.isEmpty(controlFlag)) {
                String validValues = JAAS_VALID_LOGIN_MODULE_CONTROL_FLAG_OPTIONS;

                controlFlag = JAAS_DEFAULT_LOGIN_MODULE_CONTROL_FLAG;

                LOG.warn("Unknown JAAS configuration value for ({}) = [{}], valid value are [{}] using the default value, REQUIRED", keyParam, controlFlag, validValues);
            }

            String        optionPrefix       = keyPrefix + JAAS_CONFIG_LOGIN_OPTIONS_PREFIX + ".";
            String        principalOptionKey = optionPrefix + JAAS_PRINCIPAL_PROP;
            String        passwordOptionKey  = optionPrefix + JAAS_PASSWORD_SUFFIX;
            int           optionPrefixLen    = optionPrefix.length();
            StringBuilder optionStringBuffer = new StringBuilder();

            for (String key : jaasConfig.stringPropertyNames()) {
                if (key.startsWith(optionPrefix)) {
                    String optionVal = jaasConfig.getProperty(key);

                    if (optionVal != null) {
                        optionVal = optionVal.trim();

                        try {
                            if (key.equalsIgnoreCase(principalOptionKey)) {
                                optionVal = org.apache.hadoop.security.SecurityUtil.getServerPrincipal(optionVal, (String) null);
                            }
                        } catch (IOException e) {
                            LOG.warn("Failed to build serverPrincipal. Using provided value:[{}]", optionVal);
                        }

                        if (key.equalsIgnoreCase(passwordOptionKey)) {
                            String jaasKafkaClientConfigurationProperty = "atlas.jaas.KafkaClient.option.password";

                            if (JAAS_MASK_PASSWORD.equals(configuration.getString(jaasKafkaClientConfigurationProperty))) {
                                try {
                                    optionVal = SecurityUtil.getPassword(configuration, jaasKafkaClientConfigurationProperty, HADOOP_SECURITY_CREDENTIAL_PROVIDER_PATH);
                                } catch (Exception e) {
                                    LOG.error("Error in getting secure password ", e);
                                }
                            }
                        }

                        optionVal = surroundWithQuotes(optionVal);

                        optionStringBuffer.append(String.format(" %s=%s", key.substring(optionPrefixLen), optionVal));
                    }
                }
            }

            String newJaasProperty = String.format("%s %s %s ;", loginModuleName.trim(), controlFlag, optionStringBuffer);

            kafkaProperties.put(KAFKA_SASL_JAAS_CONFIG_PROPERTY, newJaasProperty);
        }

        LOG.debug("<== KafkaUtils.setKafkaJAASProperties()");
    }

    public static boolean isLoginTicketBased() {
        boolean ret = false;

        try {
            ret = UserGroupInformation.isLoginTicketBased();
        } catch (Exception excp) {
            LOG.warn("Error in determining ticket-cache for KafkaClient-JAAS config", excp);
        }

        return ret;
    }

    static boolean isLoginKeytabBased() {
        boolean ret = false;

        try {
            ret = UserGroupInformation.isLoginKeytabBased();
        } catch (Exception excp) {
            LOG.warn("Error in determining keytab for KafkaClient-JAAS config", excp);
        }

        return ret;
    }

    static String surroundWithQuotes(String optionVal) {
        if (StringUtils.isEmpty(optionVal)) {
            return optionVal;
        }

        // Enclose property values in double quotes, so that Kafka can parse it.
        // Escape all double quotes that may occur in the property value.
        String doubleQuoteEscaped = optionVal.replace("\"", "\\\"");
        return String.format("\"%s\"", doubleQuoteEscaped);
    }

    public void createTopics(List<String> topicNames, int numPartitions, int replicationFactor) throws TopicExistsException, ExecutionException, InterruptedException {
        LOG.debug("==> createTopics() ");

        List<NewTopic> newTopicList = topicNames.stream()
                .map(topicName -> new NewTopic(topicName, numPartitions, (short) replicationFactor))
                .collect(Collectors.toList());

        CreateTopicsResult             createTopicsResult = adminClient.createTopics(newTopicList);
        Map<String, KafkaFuture<Void>> futureMap          = createTopicsResult.values();

        for (Map.Entry<String, KafkaFuture<Void>> futureEntry : futureMap.entrySet()) {
            KafkaFuture<Void> future = futureEntry.getValue();

            future.get();
        }

        LOG.debug("<== createTopics() ");
    }

    public List<String> listAllTopics() throws ExecutionException, InterruptedException {
        LOG.debug("==> KafkaUtils.listAllTopics() ");

        ListTopicsResult listTopicsResult = adminClient.listTopics((new ListTopicsOptions()).listInternal(importInternalTopics));
        List<String>     ret              = new ArrayList<>(listTopicsResult.names().get());

        LOG.debug("<== KafkaUtils.listAllTopics() ");

        return ret;
    }

    public Integer getPartitionCount(String topicName) throws ExecutionException, InterruptedException {
        LOG.debug("==> KafkaUtils.getPartitionCount({})", topicName);

        Integer                  ret           = null;
        List<TopicPartitionInfo> partitionList = getPartitionList(topicName);

        if (partitionList != null) {
            ret = partitionList.size();
        }

        LOG.debug("<== KafkaUtils.getPartitionCount returning for topic {} with count {}", topicName, ret);

        return ret;
    }

    public Integer getReplicationFactor(String topicName) throws ExecutionException, InterruptedException {
        LOG.debug("==> KafkaUtils.getReplicationFactor({})", topicName);

        Integer                  ret           = null;
        List<TopicPartitionInfo> partitionList = getPartitionList(topicName);

        if (partitionList != null) {
            ret = partitionList.stream().mapToInt(x -> x.replicas().size()).max().orElse(1);
        }

        LOG.debug("<== KafkaUtils.getReplicationFactor returning for topic {} with replicationFactor {}", topicName, ret);

        return ret;
    }

    public void close() {
        LOG.debug("==> KafkaUtils.close()");

        if (adminClient != null) {
            adminClient.close();
        }

        LOG.debug("<== KafkaUtils.close()");
    }

    private List<TopicPartitionInfo> getPartitionList(String topicName) throws ExecutionException, InterruptedException {
        List<TopicPartitionInfo> ret                  = null;
        DescribeTopicsResult     describeTopicsResult = adminClient.describeTopics(Collections.singleton(topicName));

        if (describeTopicsResult != null) {
            Map<String, KafkaFuture<TopicDescription>> futureMap = describeTopicsResult.values();

            for (Map.Entry<String, KafkaFuture<TopicDescription>> futureEntry : futureMap.entrySet()) {
                KafkaFuture<TopicDescription> topicDescriptionFuture = futureEntry.getValue();
                TopicDescription              topicDescription       = topicDescriptionFuture.get();

                ret = topicDescription.partitions();
            }
        }

        return ret;
    }
}
