/*
 * 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.coordinator.group.streams.topics;

import org.apache.kafka.common.message.StreamsGroupHeartbeatResponseData;
import org.apache.kafka.coordinator.group.streams.StreamsGroup;
import org.apache.kafka.coordinator.group.streams.StreamsGroupMember;
import org.apache.kafka.coordinator.group.streams.TopicMetadata;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class EndpointToPartitionsManager {

    private EndpointToPartitionsManager() {
    }

    public static StreamsGroupHeartbeatResponseData.EndpointToPartitions endpointToPartitions(final StreamsGroupMember streamsGroupMember,
                                                                                              final StreamsGroupHeartbeatResponseData.Endpoint responseEndpoint,
                                                                                              final StreamsGroup streamsGroup) {
        StreamsGroupHeartbeatResponseData.EndpointToPartitions endpointToPartitions = new StreamsGroupHeartbeatResponseData.EndpointToPartitions();
        Map<String, Set<Integer>> activeTasks = streamsGroupMember.assignedTasks().activeTasks();
        Map<String, Set<Integer>> standbyTasks = streamsGroupMember.assignedTasks().standbyTasks();
        endpointToPartitions.setUserEndpoint(responseEndpoint);
        Map<String, ConfiguredSubtopology> configuredSubtopologies = streamsGroup.configuredTopology().flatMap(ConfiguredTopology::subtopologies).get();
        List<StreamsGroupHeartbeatResponseData.TopicPartition> activeTopicPartitions = topicPartitions(activeTasks, configuredSubtopologies, streamsGroup.partitionMetadata());
        List<StreamsGroupHeartbeatResponseData.TopicPartition> standbyTopicPartitions = topicPartitions(standbyTasks, configuredSubtopologies, streamsGroup.partitionMetadata());
        endpointToPartitions.setActivePartitions(activeTopicPartitions);
        endpointToPartitions.setStandbyPartitions(standbyTopicPartitions);
        return endpointToPartitions;
    }

    private static List<StreamsGroupHeartbeatResponseData.TopicPartition> topicPartitions(final Map<String, Set<Integer>> tasks,
                                                                                          final Map<String, ConfiguredSubtopology> configuredSubtopologies,
                                                                                          final Map<String, TopicMetadata> groupTopicMetadata) {
        List<StreamsGroupHeartbeatResponseData.TopicPartition> topicPartitionsForTasks = new ArrayList<>();
        for (Map.Entry<String, Set<Integer>> taskEntry : tasks.entrySet()) {
            String subtopologyId = taskEntry.getKey();
            ConfiguredSubtopology configuredSubtopology = configuredSubtopologies.get(subtopologyId);
            Set<String> sourceTopics = configuredSubtopology.sourceTopics();
            Set<String> repartitionSourceTopics = configuredSubtopology.repartitionSourceTopics().keySet();
            Set<String> allSourceTopic = new HashSet<>(sourceTopics);
            allSourceTopic.addAll(repartitionSourceTopics);
            List<StreamsGroupHeartbeatResponseData.TopicPartition> topicPartitionList = topicPartitionListForTask(taskEntry.getValue(), allSourceTopic, groupTopicMetadata);
            topicPartitionsForTasks.addAll(topicPartitionList);
        }
        return topicPartitionsForTasks;
    }

    private static List<StreamsGroupHeartbeatResponseData.TopicPartition> topicPartitionListForTask(final Set<Integer> taskSet,
                                                                                                    final Set<String> topicNames,
                                                                                                    final Map<String, TopicMetadata> groupTopicMetadata) {
        return topicNames.stream().map(topic -> {
            int numPartitionsForTopic = groupTopicMetadata.get(topic).numPartitions();
            StreamsGroupHeartbeatResponseData.TopicPartition tp = new StreamsGroupHeartbeatResponseData.TopicPartition();
            tp.setTopic(topic);
            List<Integer> tpPartitions = new ArrayList<>(taskSet);
            if (numPartitionsForTopic < taskSet.size()) {
                Collections.sort(tpPartitions);
                tp.setPartitions(tpPartitions.subList(0, numPartitionsForTopic));
            } else {
                tp.setPartitions(tpPartitions);
            }
            return tp;
        }).toList();
    }
}