content: package io.aiven.klaw.clusterapi.services;
  
  import io.aiven.klaw.clusterapi.models.ApiResponse;
  import io.aiven.klaw.clusterapi.models.ClusterKeyIdentifier;
  import io.aiven.klaw.clusterapi.models.ClusterTopicRequest;
  import io.aiven.klaw.clusterapi.models.LoadTopicsResponse;
  import io.aiven.klaw.clusterapi.models.TopicConfig;
  import io.aiven.klaw.clusterapi.models.enums.ApiResultStatus;
  import io.aiven.klaw.clusterapi.models.enums.KafkaSupportedProtocol;
  import io.aiven.klaw.clusterapi.utils.ClusterApiUtils;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.kafka.clients.admin.AdminClient;
  import org.apache.kafka.clients.admin.Config;
  import org.apache.kafka.clients.admin.ConfigEntry;
  import org.apache.kafka.clients.admin.CreateTopicsResult;
  import org.apache.kafka.clients.admin.DeleteTopicsResult;
  import org.apache.kafka.clients.admin.DescribeTopicsResult;
  import org.apache.kafka.clients.admin.ListTopicsOptions;
  import org.apache.kafka.clients.admin.ListTopicsResult;
  import org.apache.kafka.clients.admin.NewPartitions;
  import org.apache.kafka.clients.admin.NewTopic;
  import org.apache.kafka.clients.admin.TopicDescription;
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.config.ConfigResource;
  import org.springframework.scheduling.annotation.Async;
  import org.springframework.scheduling.annotation.Scheduled;
  import org.springframework.stereotype.Service;
  
  @Slf4j
  @Service
  public class ApacheKafkaTopicService {
    private final ClusterApiUtils clusterApiUtils;
  
    private final SchemaService schemaService;
  
    private static Map<ClusterKeyIdentifier, Set<TopicConfig>> cachedTopics = new HashMap<>();
  
    private static Set<ClusterKeyIdentifier> topicCacheKeySets = new HashSet<>();
  
    private static Map<ClusterKeyIdentifier, Boolean> topicsLoadingStatusOfClusters = new HashMap<>();
  
    public ApacheKafkaTopicService(ClusterApiUtils clusterApiUtils, SchemaService schemaService) {
      this.clusterApiUtils = clusterApiUtils;
      this.schemaService = schemaService;
    }
  
    public synchronized LoadTopicsResponse loadTopics(
        String environment,
        KafkaSupportedProtocol protocol,
        String clusterIdentification,
        boolean resetCache)
        throws Exception {
      log.info("loadTopics {} {}", environment, protocol);
      AdminClient client =
          clusterApiUtils.getAdminClient(environment, protocol, clusterIdentification);
      Set<TopicConfig> topics = new HashSet<>();
      if (client == null) {
        throw new Exception("Cannot connect to cluster.");
      }
  
      ClusterKeyIdentifier clusterKeyIdentifier =
          new ClusterKeyIdentifier(environment, protocol, clusterIdentification);
  
      boolean topicsLoadingStatus = false;
      if (topicsLoadingStatusOfClusters.containsKey(clusterKeyIdentifier)) {
        topicsLoadingStatus = topicsLoadingStatusOfClusters.get(clusterKeyIdentifier);
      } else {
        topicsLoadingStatusOfClusters.put(clusterKeyIdentifier, false);
      }
  
      if (topicsLoadingStatus) {
        return LoadTopicsResponse.builder()
            .loadingInProgress(true)
            .topicConfigSet(new HashSet<>())
            .build();
      } else {
        topicsLoadingStatusOfClusters.put(clusterKeyIdentifier, true);
      }
  
      if (resetCache || !cachedTopics.containsKey(clusterKeyIdentifier)) {
        loadTopicsForCache(client, topics, clusterKeyIdentifier);
      } else {
        topics = cachedTopics.get(clusterKeyIdentifier);
      }
  
      topicsLoadingStatusOfClusters.put(clusterKeyIdentifier, false);
      return LoadTopicsResponse.builder().loadingInProgress(false).topicConfigSet(topics).build();
    }
  
    private void loadTopicsForCache(
        AdminClient client, Set<TopicConfig> topics, ClusterKeyIdentifier clusterKeyIdentifier) {
      try {
        Map<String, TopicDescription> topicDescriptionsPerAdminClient =
            loadTopicDescriptionsMap(client);
  
        Set<String> keySet = topicDescriptionsPerAdminClient.keySet();
        keySet.remove("_schemas");
        List<String> lstK = new ArrayList<>(keySet);
        TopicConfig topicConfig;
        for (String topicName : lstK) {
          if (topicName.startsWith("_confluent") || topicName.startsWith("__connect")) {
            continue;
          }
          topicConfig = new TopicConfig();
          topicConfig.setTopicName(topicName);
          TopicDescription topicDescription = topicDescriptionsPerAdminClient.get(topicName);
          topicConfig.setReplicationFactor(
              "" + topicDescription.partitions().get(0).replicas().size());
          topicConfig.setPartitions("" + topicDescription.partitions().size());
          topics.add(topicConfig);
        }
        updateCache(clusterKeyIdentifier, topics);
      } catch (InterruptedException | ExecutionException | TimeoutException e) {
        log.error("Exception:", e);
      }
    }
  
    private void updateCache(
        ClusterKeyIdentifier clusterKeyIdentifier, Set<TopicConfig> topicConfigSet) {
      cachedTopics.put(clusterKeyIdentifier, topicConfigSet);
      topicCacheKeySets.add(clusterKeyIdentifier);
    }
  
    private Map<String, TopicDescription> loadTopicDescriptionsMap(AdminClient client)
        throws InterruptedException, ExecutionException, TimeoutException {
      ListTopicsOptions listTopicsOptions = new ListTopicsOptions();
      listTopicsOptions = listTopicsOptions.listInternal(false);
  
      ListTopicsResult topicsResult = client.listTopics(listTopicsOptions);
      DescribeTopicsResult describeTopicsResult =
          client.describeTopics(new ArrayList<>(topicsResult.names().get()));
  
      return describeTopicsResult
          .allTopicNames()
          .get(clusterApiUtils.getAdminClientProperties().getTopicsTimeoutSecs(), TimeUnit.SECONDS);
    }
  
    public synchronized ApiResponse createTopic(ClusterTopicRequest clusterTopicRequest)
        throws Exception {
      log.info("createTopic {}", clusterTopicRequest);
      AdminClient client =
          clusterApiUtils.getAdminClient(
              clusterTopicRequest.getEnv(),
              clusterTopicRequest.getProtocol(),
              clusterTopicRequest.getClusterName());
      if (client == null) {
        throw new Exception("Cannot connect to cluster.");
      }
      try {
        NewTopic topic =
            new NewTopic(
                    clusterTopicRequest.getTopicName(),
                    clusterTopicRequest.getPartitions(),
                    clusterTopicRequest.getReplicationFactor())
                .configs(clusterTopicRequest.getAdvancedTopicConfiguration());
  
        CreateTopicsResult result = client.createTopics(Collections.singletonList(topic));
        result
            .values()
            .get(clusterTopicRequest.getTopicName())
            .get(clusterApiUtils.getAdminClientProperties().getTopicsTimeoutSecs(), TimeUnit.SECONDS);
      } catch (KafkaException e) {
        log.error("Invalid properties: ", e);
        throw e;
      } catch (NumberFormatException e) {
        log.error("Invalid replica assignment string", e);
        throw e;
      } catch (ExecutionException e) {
        log.error(
            "Unable to create topic {}, {}",
            clusterTopicRequest.getTopicName(),
            e.getCause().getMessage());
        // TopicExistsException is wrapped in ExecutionException so we have to dig into the lower
        // exception.
        if (e.getMessage().contains("TopicExistsException")) {
          log.warn(
              "Topic: {} already exists in {}",
              clusterTopicRequest.getTopicName(),
              clusterTopicRequest.getEnv());
  
          if (checkIfTopicExistsWithSameConfig(clusterTopicRequest, client)) {
            return ApiResponse.builder().success(true).message(ApiResultStatus.SUCCESS.value).build();
          }
  
          return ApiResponse.builder().success(false).message(e.getMessage()).build();
        }
        throw e;
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        log.error(
            "Unable to create topic {}, {}", clusterTopicRequest.getTopicName(), e.getMessage());
        throw e;
      } catch (Exception e) {
        log.error("Exception:", e);
        throw e;
      }
  
      return ApiResponse.builder().success(true).message(ApiResultStatus.SUCCESS.value).build();
    }
  
    // check if topic exists with same configuration as request
    private boolean checkIfTopicExistsWithSameConfig(
        ClusterTopicRequest clusterTopicRequest, AdminClient adminClient)
        throws ExecutionException, InterruptedException, TimeoutException {
      DescribeTopicsResult describeTopicsResult =
          adminClient.describeTopics(Collections.singletonList(clusterTopicRequest.getTopicName()));
  
      TopicDescription topicDescription =
          describeTopicsResult.all().get().get(clusterTopicRequest.getTopicName());
      return topicDescription.partitions().size() == clusterTopicRequest.getPartitions()
          && topicDescription.partitions().get(0).replicas().size()
              == clusterTopicRequest.getReplicationFactor();
    }
  
    public synchronized ApiResponse updateTopic(ClusterTopicRequest clusterTopicRequest)
        throws Exception {
      log.info("updateTopic Name: {}", clusterTopicRequest);
  
      AdminClient client =
          clusterApiUtils.getAdminClient(
              clusterTopicRequest.getEnv(),
              clusterTopicRequest.getProtocol(),
              clusterTopicRequest.getClusterName());
  
      if (client == null) {
        throw new Exception("Cannot connect to cluster.");
      }
  
      DescribeTopicsResult describeTopicsResult =
          client.describeTopics(Collections.singleton(clusterTopicRequest.getTopicName()));
      TopicDescription result =
          describeTopicsResult
              .all()
              .get(
                  clusterApiUtils.getAdminClientProperties().getTopicsTimeoutSecs(), TimeUnit.SECONDS)
              .get(clusterTopicRequest.getTopicName());
  
      if (result.partitions().size() > clusterTopicRequest.getPartitions()) {
        // delete topic and recreate
        deleteTopic(clusterTopicRequest);
        createTopic(clusterTopicRequest);
      } else {
        // Update partitions
        Map<String, NewPartitions> newPartitionSet = new HashMap<>();
        newPartitionSet.put(
            clusterTopicRequest.getTopicName(),
            NewPartitions.increaseTo(clusterTopicRequest.getPartitions()));
        if (result.partitions().size() != clusterTopicRequest.getPartitions()) {
          client.createPartitions(newPartitionSet);
        }
  
        // Update advanced config
        ConfigResource configResource =
            new ConfigResource(ConfigResource.Type.TOPIC, clusterTopicRequest.getTopicName());
        Map<ConfigResource, Config> updateConfig = new HashMap<>();
  
        Map<String, String> advancedConfig = clusterTopicRequest.getAdvancedTopicConfiguration();
        Collection<ConfigEntry> entries = new ArrayList<>();
        for (String key : advancedConfig.keySet()) {
          ConfigEntry configEntry = new ConfigEntry(key, advancedConfig.get(key));
          entries.add(configEntry);
        }
  
        if (!advancedConfig.isEmpty()) {
          updateConfig.put(configResource, new Config(entries));
          client.alterConfigs(updateConfig);
        }
      }
  
      return ApiResponse.builder().success(true).message(ApiResultStatus.SUCCESS.value).build();
    }
  
    public synchronized ApiResponse deleteTopic(ClusterTopicRequest clusterTopicRequest)
        throws Exception {
      log.info("deleteTopic Topic {}", clusterTopicRequest);
  
      AdminClient client;
      try {
        client =
            clusterApiUtils.getAdminClient(
                clusterTopicRequest.getEnv(),
                clusterTopicRequest.getProtocol(),
                clusterTopicRequest.getClusterName());
        if (client == null) {
          throw new Exception("Cannot connect to cluster.");
        }
  
        DeleteTopicsResult result =
            client.deleteTopics(Collections.singletonList(clusterTopicRequest.getTopicName()));
        result
            .values()
            .get(clusterTopicRequest.getTopicName())
            .get(clusterApiUtils.getAdminClientProperties().getTopicsTimeoutSecs(), TimeUnit.SECONDS);
  
        // delete associated schema if requested
        String schemaDeletionStatus = "";
        if (clusterTopicRequest.getDeleteAssociatedSchema()) {
          schemaDeletionStatus = schemaService.deleteSchema(clusterTopicRequest).getMessage();
          log.info("Schema deletion status : {}", schemaDeletionStatus);
        }
        return ApiResponse.builder().success(true).message(ApiResultStatus.SUCCESS.value).build();
      } catch (KafkaException e) {
        log.error("Invalid properties: ", e);
        throw e;
      } catch (ExecutionException | InterruptedException e) {
        String errorMessage;
        if (e instanceof ExecutionException) {
          errorMessage = e.getCause().getMessage();
        } else {
          Thread.currentThread().interrupt();
          errorMessage = e.getMessage();
        }
        if ((e.getMessage().contains("UnknownTopicOrPartition"))) {
          return ApiResponse.builder().success(true).message(ApiResultStatus.SUCCESS.value).build();
        }
        log.error("Unable to delete topic {}, {}", clusterTopicRequest.getTopicName(), errorMessage);
        throw e;
      } catch (Exception e) {
        log.error("Exception:", e);
        throw e;
      }
    }
  
    @Async("resetTopicsCacheTaskExecutor")
    @Scheduled(
        cron = "${klaw.topics.cron.expression:0 0 0 * * ?}",
        zone = "${klaw.topics.cron.expression.timezone:UTC}")
    public void resetTopicsCacheScheduler() {
      topicCacheKeySets.forEach(
          clusterKeyIdentifier -> {
            try {
              log.info("Loading topics {}", clusterKeyIdentifier);
              loadTopics(
                  clusterKeyIdentifier.getBootstrapServers(),
                  clusterKeyIdentifier.getProtocol(),
                  clusterKeyIdentifier.getClusterIdentification(),
                  true);
            } catch (Exception e) {
              log.error("Error while loading topics {}", clusterKeyIdentifier);
            }
          });
    }
  }
  