content: package io.aiven.klaw.clusterapi.services;
  
  import io.aiven.klaw.clusterapi.models.ApiResponse;
  import io.aiven.klaw.clusterapi.models.consumergroup.OffsetDetails;
  import io.aiven.klaw.clusterapi.models.consumergroup.OffsetResetType;
  import io.aiven.klaw.clusterapi.models.consumergroup.OffsetsTiming;
  import io.aiven.klaw.clusterapi.models.consumergroup.ResetConsumerGroupOffsetsRequest;
  import io.aiven.klaw.clusterapi.models.enums.KafkaSupportedProtocol;
  import io.aiven.klaw.clusterapi.utils.ClusterApiUtils;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.TreeMap;
  import java.util.concurrent.ExecutionException;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.kafka.clients.admin.AdminClient;
  import org.apache.kafka.clients.admin.AlterConsumerGroupOffsetsOptions;
  import org.apache.kafka.clients.admin.AlterConsumerGroupOffsetsResult;
  import org.apache.kafka.clients.admin.DescribeTopicsResult;
  import org.apache.kafka.clients.admin.ListOffsetsResult;
  import org.apache.kafka.clients.admin.OffsetSpec;
  import org.apache.kafka.clients.admin.TopicDescription;
  import org.apache.kafka.clients.consumer.OffsetAndMetadata;
  import org.apache.kafka.common.KafkaFuture;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.TopicPartitionInfo;
  import org.springframework.stereotype.Service;
  
  @Service
  @Slf4j
  public class ConsumerGroupService {
    public static final int TIMEOUT_MS = 2500;
    private final ClusterApiUtils clusterApiUtils;
  
    public ConsumerGroupService(ClusterApiUtils clusterApiUtils) {
      this.clusterApiUtils = clusterApiUtils;
    }
  
    public ApiResponse resetConsumerGroupOffsets(
        String environment,
        KafkaSupportedProtocol protocol,
        String clusterIdentification,
        ResetConsumerGroupOffsetsRequest consumerGroupOffsetsRequest)
        throws Exception {
      log.debug(
          "Reset consumer group offsets env {} request {}", environment, consumerGroupOffsetsRequest);
  
      if (OffsetResetType.TO_DATE_TIME.equals(consumerGroupOffsetsRequest.getOffsetResetType())
          && null == consumerGroupOffsetsRequest.getConsumerGroupResetTimestampMilliSecs()) {
        throw new Exception("Timestamp must be provided for reset type TO_DATE_TIME");
      }
      AdminClient adminClient =
          clusterApiUtils.getAdminClient(environment, protocol, clusterIdentification);
  
      TopicDescription topicDescription =
          describeTopic(adminClient, consumerGroupOffsetsRequest.getTopicName());
      if (topicDescription == null) {
        throw new Exception(
            "Topic " + consumerGroupOffsetsRequest.getTopicName() + " does not exist.");
      }
  
      OffsetSpec offsetSpec =
          switch (consumerGroupOffsetsRequest.getOffsetResetType()) {
            case EARLIEST -> OffsetSpec.earliest();
            case LATEST -> OffsetSpec.latest();
            case TO_DATE_TIME -> OffsetSpec.forTimestamp(
                consumerGroupOffsetsRequest.getConsumerGroupResetTimestampMilliSecs());
          };
      Map<OffsetsTiming, Map<String, Long>> offsetPositionsBeforeAndAfter = new HashMap<>();
  
      extractOffsetsBeforeAndAfter(
          consumerGroupOffsetsRequest,
          adminClient,
          topicDescription,
          offsetSpec,
          offsetPositionsBeforeAndAfter);
  
      if (offsetPositionsBeforeAndAfter.isEmpty()) {
        return ApiResponse.builder()
            .success(false)
            .message("Unable to reset/retrieve offsets")
            .build();
      } else {
        return ApiResponse.builder().success(true).data(offsetPositionsBeforeAndAfter).build();
      }
    }
  
    private void extractOffsetsBeforeAndAfter(
        ResetConsumerGroupOffsetsRequest consumerGroupOffsetsRequest,
        AdminClient adminClient,
        TopicDescription topicDescription,
        OffsetSpec offsetSpec,
        Map<OffsetsTiming, Map<String, Long>> offsetPositionsBeforeAndAfter)
        throws Exception {
      // Get offsets before consumer group offset update
      Map<String, Long> currentOffsets =
          getCurrentOffsetsPositions(consumerGroupOffsetsRequest.getConsumerGroup(), adminClient);
      if (!currentOffsets.isEmpty()) {
        offsetPositionsBeforeAndAfter.put(OffsetsTiming.BEFORE_OFFSET_RESET, currentOffsets);
  
        // reset offsets
        try {
          AlterConsumerGroupOffsetsOptions alterConsumerGroupOffsetsOptions =
              new AlterConsumerGroupOffsetsOptions();
          alterConsumerGroupOffsetsOptions.timeoutMs(TIMEOUT_MS);
          AlterConsumerGroupOffsetsResult futureResult =
              adminClient.alterConsumerGroupOffsets(
                  consumerGroupOffsetsRequest.getConsumerGroup(),
                  getTopicPartitionOffsetsAndMetadataMap(
                      consumerGroupOffsetsRequest.getTopicName(),
                      offsetSpec,
                      topicDescription,
                      adminClient),
                  alterConsumerGroupOffsetsOptions);
          futureResult.all().get();
        } catch (Exception e) {
          throw new Exception("Unable to reset consumer group offsets.", e);
        }
  
        // Get offsets after consumer group offset update
        offsetPositionsBeforeAndAfter.put(
            OffsetsTiming.AFTER_OFFSET_RESET,
            getCurrentOffsetsPositions(consumerGroupOffsetsRequest.getConsumerGroup(), adminClient));
      }
    }
  
    public Map<TopicPartition, OffsetAndMetadata> getTopicPartitionOffsetsAndMetadataMap(
        String topicName,
        OffsetSpec offsetSpec,
        TopicDescription topicDescription,
        AdminClient adminClient)
        throws ExecutionException, InterruptedException {
      Map<TopicPartition, OffsetAndMetadata> offsetAndMetadataMap = new HashMap<>();
      Map<TopicPartition, OffsetSpec> toOffsetSpec = new HashMap<>();
  
      List<TopicPartitionInfo> topicPartitionInfos = topicDescription.partitions();
      TopicPartition topicPartition;
  
      long partitionOffset;
      for (TopicPartitionInfo topicPartitionInfo : topicPartitionInfos) {
        topicPartition = new TopicPartition(topicName, topicPartitionInfo.partition());
        toOffsetSpec.put(topicPartition, offsetSpec);
        partitionOffset =
            adminClient.listOffsets(toOffsetSpec).partitionResult(topicPartition).get().offset();
        offsetAndMetadataMap.put(topicPartition, new OffsetAndMetadata(partitionOffset));
      }
  
      return offsetAndMetadataMap;
    }
  
    public Map<String, Long> getCurrentOffsetsPositions(String consumerGroup, AdminClient adminClient)
        throws Exception {
      Map<String, Long> currentOffsetPositionsMap = new TreeMap<>();
      KafkaFuture<Map<TopicPartition, OffsetAndMetadata>> offsetsBeforeResetFuture =
          adminClient.listConsumerGroupOffsets(consumerGroup).partitionsToOffsetAndMetadata();
      try {
        Map<TopicPartition, OffsetAndMetadata> topicPartitionOffsetAndMetadataMap =
            offsetsBeforeResetFuture.get();
        for (TopicPartition topicPartition : topicPartitionOffsetAndMetadataMap.keySet()) {
          currentOffsetPositionsMap.put(
              topicPartition.toString(),
              topicPartitionOffsetAndMetadataMap.get(topicPartition).offset());
        }
        return currentOffsetPositionsMap;
      } catch (InterruptedException | ExecutionException e) {
        // ignore error as there may not be any events or couldn't retrieve events
        return new HashMap<>();
      }
    }
  
    public TopicDescription describeTopic(AdminClient adminClient, String topicName)
        throws Exception {
      try {
        DescribeTopicsResult result = adminClient.describeTopics(Collections.singleton(topicName));
        if (result.values().containsKey(topicName)) {
          return result.values().get(topicName).get();
        } else {
          return null;
        }
      } catch (Exception ex) {
        if (ex.getMessage().contains("UnknownTopic")) {
          log.info("Topic {} does not exist on the cluster.", topicName);
          return null;
        } else {
          String errorMessage = "Unable to describe the topic " + topicName;
          log.error(errorMessage, ex);
          throw new Exception(errorMessage, ex);
        }
      }
    }
  
    public List<OffsetDetails> getConsumerGroupDetails(
        String consumerGroupId,
        String topicName,
        String environment,
        KafkaSupportedProtocol protocol,
        String clusterName)
        throws Exception {
      log.info(
          "getConsumerGroupDetails topicName {} consumerGroupId {} environment {} protocol {} clusterName {}",
          topicName,
          consumerGroupId,
          environment,
          protocol,
          clusterName);
  
      AdminClient adminClient = clusterApiUtils.getAdminClient(environment, protocol, clusterName);
  
      List<OffsetDetails> consumerGroupOffsetList = new ArrayList<>();
      OffsetDetails offsetDetails;
      try {
        DescribeTopicsResult describeTopicsResult =
            adminClient.describeTopics(Collections.singletonList(topicName));
        List<TopicPartitionInfo> topicPartitions =
            describeTopicsResult.values().get(topicName).get().partitions();
  
        TopicPartition topicPartition;
  
        ListOffsetsResult listOffsetsEarliestResult =
            adminClient.listOffsets(
                getTopicPartitionOffsetSpecMap(topicName, topicPartitions, OffsetSpec.earliest()));
  
        ListOffsetsResult listOffsetsLatestResult =
            adminClient.listOffsets(
                getTopicPartitionOffsetSpecMap(topicName, topicPartitions, OffsetSpec.latest()));
  
        for (TopicPartitionInfo topicPartitionInfo : topicPartitions) {
          topicPartition = new TopicPartition(topicName, topicPartitionInfo.partition());
  
          offsetDetails = new OffsetDetails();
          long earliestOffset =
              listOffsetsEarliestResult.partitionResult(topicPartition).get().offset();
          long latestOffset = listOffsetsLatestResult.partitionResult(topicPartition).get().offset();
          long lag = latestOffset - earliestOffset;
  
          offsetDetails.setTopicPartitionId(Long.toString(topicPartition.partition()));
          offsetDetails.setCurrentOffset(Long.toString(earliestOffset));
          offsetDetails.setEndOffset(Long.toString(latestOffset));
          offsetDetails.setLag(Long.toString(lag));
          consumerGroupOffsetList.add(offsetDetails);
        }
        return consumerGroupOffsetList;
      } catch (Exception exception) {
        log.error(
            "Cannot retrieve consumer offset details topicName: {} groupid: {} Error: {}",
            topicName,
            consumerGroupId,
            exception);
        return consumerGroupOffsetList;
      }
    }
  
    private Map<TopicPartition, OffsetSpec> getTopicPartitionOffsetSpecMap(
        String topicName, List<TopicPartitionInfo> topicPartitions, OffsetSpec offsetSpec) {
      TopicPartition topicPartition;
      Map<TopicPartition, OffsetSpec> topicPartitionOffsetSpecMap = new HashMap<>();
      for (TopicPartitionInfo topicPartitionInfo : topicPartitions) {
        topicPartition = new TopicPartition(topicName, topicPartitionInfo.partition());
        topicPartitionOffsetSpecMap.put(topicPartition, offsetSpec);
      }
      return topicPartitionOffsetSpecMap;
    }
  }
  