content: /*
   * 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;
  
  import org.apache.kafka.common.message.StreamsGroupHeartbeatRequestData;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
  
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Objects;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  /**
   * An immutable tuple containing active, standby and warm-up tasks.
   *
   * @param activeTasks           Active tasks.
   *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
   * @param standbyTasks          Standby tasks.
   *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
   * @param warmupTasks           Warm-up tasks.
   *                              The key of the map is the subtopology ID, and the value is the set of partition IDs.
   */
  public record TasksTuple(Map<String, Set<Integer>> activeTasks,
                           Map<String, Set<Integer>> standbyTasks,
                           Map<String, Set<Integer>> warmupTasks) {
  
      public TasksTuple {
          activeTasks = Collections.unmodifiableMap(Objects.requireNonNull(activeTasks));
          standbyTasks = Collections.unmodifiableMap(Objects.requireNonNull(standbyTasks));
          warmupTasks = Collections.unmodifiableMap(Objects.requireNonNull(warmupTasks));
      }
  
      /**
       * An empty task tuple.
       */
      public static final TasksTuple EMPTY = new TasksTuple(
          Map.of(),
          Map.of(),
          Map.of()
      );
  
      /**
       * @return true if all collections in the tuple are empty.
       */
      public boolean isEmpty() {
          return activeTasks.isEmpty() && standbyTasks.isEmpty() && warmupTasks.isEmpty();
      }
  
      /**
       * Merges this task tuple with another task tuple.
       *
       * @param other The other task tuple.
       * @return A new task tuple, containing all active tasks, standby tasks and warm-up tasks from both tuples.
       */
      public TasksTuple merge(TasksTuple other) {
          Map<String, Set<Integer>> mergedActiveTasks = merge(activeTasks, other.activeTasks);
          Map<String, Set<Integer>> mergedStandbyTasks = merge(standbyTasks, other.standbyTasks);
          Map<String, Set<Integer>> mergedWarmupTasks = merge(warmupTasks, other.warmupTasks);
          return new TasksTuple(mergedActiveTasks, mergedStandbyTasks, mergedWarmupTasks);
      }
  
      private static Map<String, Set<Integer>> merge(final Map<String, Set<Integer>> tasks1, final Map<String, Set<Integer>> tasks2) {
          HashMap<String, Set<Integer>> result = new HashMap<>();
          tasks1.forEach((subtopologyId, tasks) ->
              result.put(subtopologyId, new HashSet<>(tasks)));
          tasks2.forEach((subtopologyId, tasks) -> result
              .computeIfAbsent(subtopologyId, __ -> new HashSet<>())
              .addAll(tasks));
          return result;
      }
  
      /**
       * Checks if this task tuple contains any of the tasks in another task tuple.
       *
       * @param other Another task tuple.
       * @return true if there is at least one active, standby or warm-up task that is present in both tuples.
       */
      public boolean containsAny(TasksTuple other) {
          return activeTasks.entrySet().stream().anyMatch(
              entry -> other.activeTasks.containsKey(entry.getKey()) && !Collections.disjoint(entry.getValue(), other.activeTasks.get(entry.getKey()))
          ) || standbyTasks.entrySet().stream().anyMatch(
              entry -> other.standbyTasks.containsKey(entry.getKey()) && !Collections.disjoint(entry.getValue(), other.standbyTasks.get(entry.getKey()))
          ) || warmupTasks.entrySet().stream().anyMatch(
              entry -> other.warmupTasks.containsKey(entry.getKey()) && !Collections.disjoint(entry.getValue(), other.warmupTasks.get(entry.getKey()))
          );
      }
  
      /**
       * Creates a {{@link TasksTuple}} from a
       * {{@link org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue}}.
       *
       * @param record The record.
       * @return A {{@link TasksTuple}}.
       */
      public static TasksTuple fromTargetAssignmentRecord(StreamsGroupTargetAssignmentMemberValue record) {
          return new TasksTuple(
              record.activeTasks().stream()
                  .collect(Collectors.toMap(
                          StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  ),
              record.standbyTasks().stream()
                  .collect(Collectors.toMap(
                          StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  ),
              record.warmupTasks().stream()
                  .collect(Collectors.toMap(
                          StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  )
          );
      }
  
      public String toString() {
          return "(active=" + taskAssignmentToString(activeTasks) +
              ", standby=" + taskAssignmentToString(standbyTasks) +
              ", warmup=" + taskAssignmentToString(warmupTasks) +
              ')';
      }
  
      public static TasksTuple fromHeartbeatRequest(final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedActiveTasks,
                                                    final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedStandbyTasks,
                                                    final List<StreamsGroupHeartbeatRequestData.TaskIds> ownedWarmupTasks) {
          return new TasksTuple(
              ownedActiveTasks.stream()
                  .collect(Collectors.toMap(
                      StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  ),
              ownedStandbyTasks.stream()
                  .collect(Collectors.toMap(
                      StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  ),
              ownedWarmupTasks.stream()
                  .collect(Collectors.toMap(
                      StreamsGroupHeartbeatRequestData.TaskIds::subtopologyId,
                          taskId -> new HashSet<>(taskId.partitions())
                      )
                  )
          );
      }
  
      /**
       * @return The provided assignment as a String.
       *
       * Example:
       * [subtopologyID1-0, subtopologyID1-1, subtopologyID2-0, subtopologyID2-1]
       */
      private static String taskAssignmentToString(
          Map<String, Set<Integer>> assignment
      ) {
          StringBuilder builder = new StringBuilder("[");
          Iterator<Entry<String, Set<Integer>>> subtopologyIterator = assignment.entrySet().iterator();
          while (subtopologyIterator.hasNext()) {
              Map.Entry<String, Set<Integer>> entry = subtopologyIterator.next();
              Iterator<Integer> partitionsIterator = entry.getValue().iterator();
              while (partitionsIterator.hasNext()) {
                  builder.append(entry.getKey());
                  builder.append("-");
                  builder.append(partitionsIterator.next());
                  if (partitionsIterator.hasNext() || subtopologyIterator.hasNext()) {
                      builder.append(", ");
                  }
              }
          }
          builder.append("]");
          return builder.toString();
      }
  }
  