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.common.runtime.CoordinatorRecord;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupCurrentMemberAssignmentKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupCurrentMemberAssignmentValue;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupMemberMetadataKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupMemberMetadataValue;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupMetadataKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupMetadataValue;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMemberValue;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMetadataKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTargetAssignmentMetadataValue;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTopologyKey;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupTopologyValue;
  import org.apache.kafka.server.common.ApiMessageAndVersion;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  
  /**
   * This class contains helper methods to create records stored in the __consumer_offsets topic.
   */
  public class StreamsCoordinatorRecordHelpers {
  
      public static CoordinatorRecord newStreamsGroupMemberRecord(
          String groupId,
          StreamsGroupMember member
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(member, "member should not be null here");
  
          return CoordinatorRecord.record(
              new StreamsGroupMemberMetadataKey()
                  .setGroupId(groupId)
                  .setMemberId(member.memberId()),
              new ApiMessageAndVersion(
                  new StreamsGroupMemberMetadataValue()
                      .setRackId(member.rackId().orElse(null))
                      .setInstanceId(member.instanceId().orElse(null))
                      .setClientId(member.clientId())
                      .setClientHost(member.clientHost())
                      .setRebalanceTimeoutMs(member.rebalanceTimeoutMs())
                      .setTopologyEpoch(member.topologyEpoch())
                      .setProcessId(member.processId())
                      .setUserEndpoint(member.userEndpoint().orElse(null))
                      .setClientTags(member.clientTags().entrySet().stream().map(e ->
                          new StreamsGroupMemberMetadataValue.KeyValue()
                              .setKey(e.getKey())
                              .setValue(e.getValue())
                      ).sorted(Comparator.comparing(StreamsGroupMemberMetadataValue.KeyValue::key)).toList()),
                  (short) 0
              )
          );
      }
  
      /**
       * Creates a StreamsGroupMemberMetadata tombstone.
       *
       * @param groupId  The streams group id.
       * @param memberId The streams group member id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupMemberTombstoneRecord(
          String groupId,
          String memberId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(memberId, "memberId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupMemberMetadataKey()
                  .setGroupId(groupId)
                  .setMemberId(memberId)
          );
      }
  
      public static CoordinatorRecord newStreamsGroupEpochRecord(
          String groupId,
          int newGroupEpoch,
          long metadataHash
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
  
          return CoordinatorRecord.record(
              new StreamsGroupMetadataKey()
                  .setGroupId(groupId),
              new ApiMessageAndVersion(
                  new StreamsGroupMetadataValue()
                      .setEpoch(newGroupEpoch)
                      .setMetadataHash(metadataHash),
                  (short) 0
              )
          );
      }
  
      /**
       * Creates a StreamsGroupMetadata tombstone.
       *
       * @param groupId The streams group id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupEpochTombstoneRecord(
          String groupId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupMetadataKey()
                  .setGroupId(groupId)
          );
      }
  
      public static CoordinatorRecord newStreamsGroupTargetAssignmentRecord(
          String groupId,
          String memberId,
          TasksTuple assignment
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(memberId, "memberId should not be null here");
          Objects.requireNonNull(assignment, "assignment should not be null here");
  
          List<StreamsGroupTargetAssignmentMemberValue.TaskIds> activeTaskIds = new ArrayList<>(assignment.activeTasks().size());
          for (Map.Entry<String, Set<Integer>> entry : assignment.activeTasks().entrySet()) {
              activeTaskIds.add(
                  new StreamsGroupTargetAssignmentMemberValue.TaskIds()
                      .setSubtopologyId(entry.getKey())
                      .setPartitions(entry.getValue().stream().sorted().toList())
              );
          }
          activeTaskIds.sort(Comparator.comparing(StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId));
          List<StreamsGroupTargetAssignmentMemberValue.TaskIds> standbyTaskIds = new ArrayList<>(assignment.standbyTasks().size());
          for (Map.Entry<String, Set<Integer>> entry : assignment.standbyTasks().entrySet()) {
              standbyTaskIds.add(
                  new StreamsGroupTargetAssignmentMemberValue.TaskIds()
                      .setSubtopologyId(entry.getKey())
                      .setPartitions(entry.getValue().stream().sorted().toList())
              );
          }
          standbyTaskIds.sort(Comparator.comparing(StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId));
          List<StreamsGroupTargetAssignmentMemberValue.TaskIds> warmupTaskIds = new ArrayList<>(assignment.warmupTasks().size());
          for (Map.Entry<String, Set<Integer>> entry : assignment.warmupTasks().entrySet()) {
              warmupTaskIds.add(
                  new StreamsGroupTargetAssignmentMemberValue.TaskIds()
                      .setSubtopologyId(entry.getKey())
                      .setPartitions(entry.getValue().stream().sorted().toList())
              );
          }
          warmupTaskIds.sort(Comparator.comparing(StreamsGroupTargetAssignmentMemberValue.TaskIds::subtopologyId));
  
          return CoordinatorRecord.record(
              new StreamsGroupTargetAssignmentMemberKey()
                  .setGroupId(groupId)
                  .setMemberId(memberId),
              new ApiMessageAndVersion(
                  new StreamsGroupTargetAssignmentMemberValue()
                      .setActiveTasks(activeTaskIds)
                      .setStandbyTasks(standbyTaskIds)
                      .setWarmupTasks(warmupTaskIds),
                  (short) 0
              )
          );
      }
  
      /**
       * Creates a StreamsGroupTargetAssignmentMember tombstone.
       *
       * @param groupId  The streams group id.
       * @param memberId The streams group member id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupTargetAssignmentTombstoneRecord(
          String groupId,
          String memberId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(memberId, "memberId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupTargetAssignmentMemberKey()
                  .setGroupId(groupId)
                  .setMemberId(memberId)
          );
      }
  
  
      public static CoordinatorRecord newStreamsGroupTargetAssignmentEpochRecord(
          String groupId,
          int assignmentEpoch
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
  
          return CoordinatorRecord.record(
              new StreamsGroupTargetAssignmentMetadataKey()
                  .setGroupId(groupId),
              new ApiMessageAndVersion(
                  new StreamsGroupTargetAssignmentMetadataValue()
                      .setAssignmentEpoch(assignmentEpoch),
                  (short) 0
              )
          );
      }
  
      /**
       * Creates a StreamsGroupTargetAssignmentMetadata tombstone.
       *
       * @param groupId The streams group id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupTargetAssignmentEpochTombstoneRecord(
          String groupId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupTargetAssignmentMetadataKey()
                  .setGroupId(groupId)
          );
      }
  
      public static CoordinatorRecord newStreamsGroupCurrentAssignmentRecord(
          String groupId,
          StreamsGroupMember member
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(member, "member should not be null here");
  
          return CoordinatorRecord.record(
              new StreamsGroupCurrentMemberAssignmentKey()
                  .setGroupId(groupId)
                  .setMemberId(member.memberId()),
              new ApiMessageAndVersion(
                  new StreamsGroupCurrentMemberAssignmentValue()
                      .setMemberEpoch(member.memberEpoch())
                      .setPreviousMemberEpoch(member.previousMemberEpoch())
                      .setState(member.state().value())
                      .setActiveTasks(toTaskIds(member.assignedTasks().activeTasks()))
                      .setStandbyTasks(toTaskIds(member.assignedTasks().standbyTasks()))
                      .setWarmupTasks(toTaskIds(member.assignedTasks().warmupTasks()))
                      .setActiveTasksPendingRevocation(toTaskIds(member.tasksPendingRevocation().activeTasks()))
                      .setStandbyTasksPendingRevocation(toTaskIds(member.tasksPendingRevocation().standbyTasks()))
                      .setWarmupTasksPendingRevocation(toTaskIds(member.tasksPendingRevocation().warmupTasks())),
                  (short) 0
              )
          );
      }
  
      /**
       * Creates a StreamsGroupCurrentMemberAssignment tombstone.
       *
       * @param groupId  The streams group id.
       * @param memberId The streams group member id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupCurrentAssignmentTombstoneRecord(
          String groupId,
          String memberId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(memberId, "memberId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupCurrentMemberAssignmentKey()
                  .setGroupId(groupId)
                  .setMemberId(memberId)
          );
      }
  
      private static List<StreamsGroupCurrentMemberAssignmentValue.TaskIds> toTaskIds(
          Map<String, Set<Integer>> tasks
      ) {
          List<StreamsGroupCurrentMemberAssignmentValue.TaskIds> taskIds = new ArrayList<>(tasks.size());
          tasks.forEach((subtopologyId, partitions) ->
              taskIds.add(new StreamsGroupCurrentMemberAssignmentValue.TaskIds()
                  .setSubtopologyId(subtopologyId)
                  .setPartitions(partitions.stream().sorted().toList()))
          );
          taskIds.sort(Comparator.comparing(StreamsGroupCurrentMemberAssignmentValue.TaskIds::subtopologyId));
          return taskIds;
      }
  
      /**
       * Creates a StreamsTopology record.
       *
       * @param groupId  The consumer group id.
       * @param topology The new topology.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupTopologyRecord(String groupId,
                                                                    StreamsGroupHeartbeatRequestData.Topology topology) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(topology, "topology should not be null here");
  
          return newStreamsGroupTopologyRecord(groupId, convertToStreamsGroupTopologyRecord(topology));
      }
  
      /**
       * Creates a StreamsTopology record.
       *
       * @param groupId The consumer group id.
       * @param value   The encoded topology record value.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupTopologyRecord(String groupId, StreamsGroupTopologyValue value) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
          Objects.requireNonNull(value, "value should not be null here");
  
          return CoordinatorRecord.record(
              new StreamsGroupTopologyKey()
                  .setGroupId(groupId),
              new ApiMessageAndVersion(value, (short) 0)
          );
      }
  
      /**
       * Encodes subtopologies from the Heartbeat RPC to a StreamsTopology record value.
       *
       * @param topology The new topology
       * @return The record value.
       */
      public static StreamsGroupTopologyValue convertToStreamsGroupTopologyRecord(StreamsGroupHeartbeatRequestData.Topology topology) {
          Objects.requireNonNull(topology, "topology should not be null here");
  
          StreamsGroupTopologyValue value = new StreamsGroupTopologyValue();
          value.setEpoch(topology.epoch());
          topology.subtopologies().forEach(subtopology -> {
              List<StreamsGroupTopologyValue.TopicInfo> repartitionSourceTopics =
                  subtopology.repartitionSourceTopics().stream()
                      .map(StreamsCoordinatorRecordHelpers::convertToTopicInfo)
                      .toList();
  
              List<StreamsGroupTopologyValue.TopicInfo> stateChangelogTopics =
                  subtopology.stateChangelogTopics().stream()
                      .map(StreamsCoordinatorRecordHelpers::convertToTopicInfo)
                      .toList();
  
              List<StreamsGroupTopologyValue.CopartitionGroup> copartitionGroups =
                  subtopology.copartitionGroups().stream()
                      .map(copartitionGroup -> new StreamsGroupTopologyValue.CopartitionGroup()
                          .setSourceTopics(copartitionGroup.sourceTopics())
                          .setSourceTopicRegex(copartitionGroup.sourceTopicRegex())
                          .setRepartitionSourceTopics(copartitionGroup.repartitionSourceTopics())
                      )
                      .toList();
  
              value.subtopologies().add(
                  new StreamsGroupTopologyValue.Subtopology()
                      .setSubtopologyId(subtopology.subtopologyId())
                      .setSourceTopics(subtopology.sourceTopics())
                      .setSourceTopicRegex(subtopology.sourceTopicRegex())
                      .setRepartitionSinkTopics(subtopology.repartitionSinkTopics())
                      .setRepartitionSourceTopics(repartitionSourceTopics)
                      .setStateChangelogTopics(stateChangelogTopics)
                      .setCopartitionGroups(copartitionGroups)
              );
          });
          return value;
      }
  
      private static StreamsGroupTopologyValue.TopicInfo convertToTopicInfo(StreamsGroupHeartbeatRequestData.TopicInfo topicInfo) {
          List<StreamsGroupTopologyValue.TopicConfig> topicConfigs = topicInfo.topicConfigs() != null ? topicInfo.topicConfigs().stream()
              .map(config -> new StreamsGroupTopologyValue.TopicConfig().setKey(config.key()).setValue(config.value()))
              .toList() : null;
          return new StreamsGroupTopologyValue.TopicInfo()
              .setName(topicInfo.name())
              .setTopicConfigs(topicConfigs)
              .setPartitions(topicInfo.partitions())
              .setReplicationFactor(topicInfo.replicationFactor());
      }
  
      /**
       * Creates a StreamsGroupTopology tombstone.
       *
       * @param groupId The streams group id.
       * @return The record.
       */
      public static CoordinatorRecord newStreamsGroupTopologyRecordTombstone(
          String groupId
      ) {
          Objects.requireNonNull(groupId, "groupId should not be null here");
  
          return CoordinatorRecord.tombstone(
              new StreamsGroupTopologyKey()
                  .setGroupId(groupId)
          );
      }
  }
  