/*
 * 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
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.kafka.streams.processor.internals;

import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.streams.errors.MissingSourceTopicException;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.errors.TaskAssignmentException;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder.TopicsInfo;
import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.streams.processor.internals.assignment.CopartitionedTopicsEnforcer;

import org.slf4j.Logger;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;
import java.util.Set;
import java.util.stream.Collectors;

import static org.apache.kafka.streams.processor.internals.TopologyMetadata.getTopologyNameOrElseUnnamed;

public class RepartitionTopics {

    private final InternalTopicManager internalTopicManager;
    private final TopologyMetadata topologyMetadata;
    private final Cluster clusterMetadata;
    private final CopartitionedTopicsEnforcer copartitionedTopicsEnforcer;
    private final Logger log;
    private final Map<TopicPartition, PartitionInfo> topicPartitionInfos = new HashMap<>();
    private final Map<Subtopology, Set<String>> missingInputTopicsBySubtopology = new HashMap<>();

    public RepartitionTopics(final TopologyMetadata topologyMetadata,
                             final InternalTopicManager internalTopicManager,
                             final CopartitionedTopicsEnforcer copartitionedTopicsEnforcer,
                             final Cluster clusterMetadata,
                             final String logPrefix) {
        this.topologyMetadata = topologyMetadata;
        this.internalTopicManager = internalTopicManager;
        this.clusterMetadata = clusterMetadata;
        this.copartitionedTopicsEnforcer = copartitionedTopicsEnforcer;
        final LogContext logContext = new LogContext(logPrefix);
        log = logContext.logger(getClass());
    }

    public void setup() {
        final Map<String, InternalTopicConfig> repartitionTopicMetadata = computeRepartitionTopicConfig(clusterMetadata);

        if (repartitionTopicMetadata.isEmpty()) {
            if (missingInputTopicsBySubtopology.isEmpty()) {
                log.info("Skipping the repartition topic validation since there are no repartition topics.");
            } else {
                log.info("Skipping the repartition topic validation since all topologies containing repartition"
                             + "topics are missing external user source topics and cannot be processed.");
            }
        } else {
            // ensure the co-partitioning topics within the group have the same number of partitions,
            // and enforce the number of partitions for those repartition topics to be the same if they
            // are co-partitioned as well.
            ensureCopartitioning(topologyMetadata.copartitionGroups(), repartitionTopicMetadata, clusterMetadata);

            // make sure the repartition source topics exist with the right number of partitions,
            // create these topics if necessary
            internalTopicManager.makeReady(repartitionTopicMetadata);

            // augment the metadata with the newly computed number of partitions for all the
            // repartition source topics
            for (final Map.Entry<String, InternalTopicConfig> entry : repartitionTopicMetadata.entrySet()) {
                final String topic = entry.getKey();
                final int numPartitions = entry.getValue().numberOfPartitions().orElse(-1);

                for (int partition = 0; partition < numPartitions; partition++) {
                    topicPartitionInfos.put(
                        new TopicPartition(topic, partition),
                        new PartitionInfo(topic, partition, null, new Node[0], new Node[0])
                    );
                }
            }
        }
    }

    public Set<String> topologiesWithMissingInputTopics() {
        return missingInputTopicsBySubtopology.keySet()
            .stream()
            .map(s -> getTopologyNameOrElseUnnamed(s.namedTopology))
            .collect(Collectors.toSet());
    }

    public Set<String> missingSourceTopics() {
        return missingInputTopicsBySubtopology.entrySet().stream()
                .map(entry -> entry.getValue())
                .flatMap(missingTopicSet -> missingTopicSet.stream())
                .collect(Collectors.toSet());
    }

    public Queue<StreamsException> missingSourceTopicExceptions() {
        return missingInputTopicsBySubtopology.entrySet().stream().map(entry -> {
            final Set<String> missingSourceTopics = entry.getValue();
            final int subtopologyId = entry.getKey().nodeGroupId;
            final String topologyName = entry.getKey().namedTopology;

            return new StreamsException(
                new MissingSourceTopicException(String.format(
                    "Missing source topics %s for subtopology %d of topology %s",
                    missingSourceTopics, subtopologyId, topologyName)),
                new TaskId(subtopologyId, 0, topologyName));
        }).collect(Collectors.toCollection(LinkedList::new));
    }

    public Map<TopicPartition, PartitionInfo> topicPartitionsInfo() {
        return Collections.unmodifiableMap(topicPartitionInfos);
    }

    /**
     * @param clusterMetadata  cluster metadata, eg which topics exist on the brokers
     */
    private Map<String, InternalTopicConfig> computeRepartitionTopicConfig(final Cluster clusterMetadata) {
        final Set<TopicsInfo> allTopicsInfo = new HashSet<>();
        final Map<String, InternalTopicConfig> allRepartitionTopicConfigs = new HashMap<>();
        for (final Map.Entry<String, Map<Subtopology, TopicsInfo>> topologyEntry : topologyMetadata.topologyToSubtopologyTopicsInfoMap().entrySet()) {
            final String topologyName = topologyMetadata.hasNamedTopologies() ? topologyEntry.getKey() : null;

            final Set<TopicsInfo> topicsInfoForTopology = new HashSet<>();
            final Set<String> missingSourceTopicsForTopology = new HashSet<>();
            final Map<String, InternalTopicConfig> repartitionTopicConfigsForTopology = new HashMap<>();

            for (final Map.Entry<Subtopology, TopicsInfo> subtopologyEntry : topologyEntry.getValue().entrySet()) {
                final TopicsInfo topicsInfo = subtopologyEntry.getValue();

                topicsInfoForTopology.add(topicsInfo);
                repartitionTopicConfigsForTopology.putAll(
                    topicsInfo.repartitionSourceTopics
                        .values()
                        .stream()
                        .collect(Collectors.toMap(InternalTopicConfig::name, topicConfig -> topicConfig)));

                final Set<String> missingSourceTopicsForSubtopology = computeMissingExternalSourceTopics(topicsInfo, clusterMetadata);
                missingSourceTopicsForTopology.addAll(missingSourceTopicsForSubtopology);
                if (!missingSourceTopicsForSubtopology.isEmpty()) {
                    final Subtopology subtopology = subtopologyEntry.getKey();
                    missingInputTopicsBySubtopology.put(subtopology, missingSourceTopicsForSubtopology);
                    log.error("Subtopology {} has missing source topics {} and will be excluded from the current assignment, "
                        + "this can be due to the consumer client's metadata being stale or because they have "
                        + "not been created yet. Please verify that you have created all input topics; if they "
                        + "do exist, you just need to wait for the metadata to be updated, at which time a new "
                        + "rebalance will be kicked off automatically and the topology will be retried at that time.",
                        subtopology.nodeGroupId, missingSourceTopicsForSubtopology);
                }
            }

            if (missingSourceTopicsForTopology.isEmpty()) {
                allRepartitionTopicConfigs.putAll(repartitionTopicConfigsForTopology);
                allTopicsInfo.addAll(topicsInfoForTopology);
            } else {
                log.debug("Skipping repartition topic validation for entire topology {} due to missing source topics {}",
                    topologyName, missingSourceTopicsForTopology);
            }
        }
        setRepartitionSourceTopicPartitionCount(allRepartitionTopicConfigs, allTopicsInfo, clusterMetadata);

        return allRepartitionTopicConfigs;
    }

    private void ensureCopartitioning(final Collection<Set<String>> copartitionGroups,
                                      final Map<String, InternalTopicConfig> repartitionTopicMetadata,
                                      final Cluster clusterMetadata) {
        for (final Set<String> copartitionGroup : copartitionGroups) {
            copartitionedTopicsEnforcer.enforce(copartitionGroup, repartitionTopicMetadata, clusterMetadata);
        }
    }

    private Set<String> computeMissingExternalSourceTopics(final TopicsInfo topicsInfo,
                                                           final Cluster clusterMetadata) {
        final Set<String> missingExternalSourceTopics = new HashSet<>(topicsInfo.sourceTopics);
        missingExternalSourceTopics.removeAll(topicsInfo.repartitionSourceTopics.keySet());
        missingExternalSourceTopics.removeAll(clusterMetadata.topics());
        return missingExternalSourceTopics;
    }

    /**
     * Computes the number of partitions and sets it for each repartition topic in repartitionTopicMetadata
     */
    private void setRepartitionSourceTopicPartitionCount(final Map<String, InternalTopicConfig> repartitionTopicMetadata,
                                                         final Collection<TopicsInfo> topicGroups,
                                                         final Cluster clusterMetadata) {
        boolean partitionCountNeeded;
        do {
            partitionCountNeeded = false;
            boolean progressMadeThisIteration = false;  // avoid infinitely looping without making any progress on unknown repartitions

            for (final TopicsInfo topicsInfo : topicGroups) {
                for (final String repartitionSourceTopic : topicsInfo.repartitionSourceTopics.keySet()) {
                    final Optional<Integer> repartitionSourceTopicPartitionCount =
                        repartitionTopicMetadata.get(repartitionSourceTopic).numberOfPartitions();

                    if (!repartitionSourceTopicPartitionCount.isPresent()) {
                        final Integer numPartitions = computePartitionCount(
                            repartitionTopicMetadata,
                            topicGroups,
                            clusterMetadata,
                            repartitionSourceTopic
                        );

                        if (numPartitions == null) {
                            partitionCountNeeded = true;
                            log.trace("Unable to determine number of partitions for {}, another iteration is needed",
                                repartitionSourceTopic);
                        } else {
                            log.trace("Determined number of partitions for {} to be {}", repartitionSourceTopic, numPartitions);
                            repartitionTopicMetadata.get(repartitionSourceTopic).setNumberOfPartitions(numPartitions);
                            progressMadeThisIteration = true;
                        }
                    }
                }
            }
            if (!progressMadeThisIteration && partitionCountNeeded) {
                log.error("Unable to determine the number of partitions of all repartition topics, most likely a source topic is missing or pattern doesn't match any topics\n" +
                    "topic groups: {}\n" +
                    "cluster topics: {}.", topicGroups, clusterMetadata.topics());
                throw new TaskAssignmentException("Failed to compute number of partitions for all repartition topics, " +
                    "make sure all user input topics are created and all Pattern subscriptions match at least one topic in the cluster");
            }
        } while (partitionCountNeeded);
    }

    private Integer computePartitionCount(final Map<String, InternalTopicConfig> repartitionTopicMetadata,
                                          final Collection<TopicsInfo> topicGroups,
                                          final Cluster clusterMetadata,
                                          final String repartitionSourceTopic) {
        Integer partitionCount = null;
        // try set the number of partitions for this repartition topic if it is not set yet
        for (final TopicsInfo topicsInfo : topicGroups) {
            final Set<String> sinkTopics = topicsInfo.sinkTopics;

            if (sinkTopics.contains(repartitionSourceTopic)) {
                // if this topic is one of the sink topics of this topology,
                // use the maximum of all its source topic partitions as the number of partitions
                for (final String upstreamSourceTopic : topicsInfo.sourceTopics) {
                    Integer numPartitionsCandidate = null;
                    // It is possible the sourceTopic is another internal topic, i.e,
                    // map().join().join(map())
                    if (repartitionTopicMetadata.containsKey(upstreamSourceTopic)) {
                        if (repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().isPresent()) {
                            numPartitionsCandidate =
                                repartitionTopicMetadata.get(upstreamSourceTopic).numberOfPartitions().get();
                        }
                    } else {
                        final Integer count = clusterMetadata.partitionCountForTopic(upstreamSourceTopic);
                        if (count == null) {
                            throw new TaskAssignmentException(
                                "No partition count found for source topic "
                                    + upstreamSourceTopic
                                    + ", but it should have been."
                            );
                        }
                        numPartitionsCandidate = count;
                    }

                    if (numPartitionsCandidate != null) {
                        if (partitionCount == null || numPartitionsCandidate > partitionCount) {
                            partitionCount = numPartitionsCandidate;
                        }
                    }
                }
            }
        }
        return partitionCount;
    }
}
