import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.AdminClientConfig;
import org.apache.kafka.clients.admin.NewPartitions;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.admin.TopicDescription;

public final class EnsureKafkaTopics {
  private static final short REPLICATION_FACTOR = 1;

  private record TopicSpec(String name, int partitions) {}

  private EnsureKafkaTopics() {}

  public static void main(String[] args) throws Exception {
    String brokers = args.length > 0 ? args[0] : "172.17.111.56:9092";
    List<TopicSpec> specs =
        List.of(
            topic("vehicle.event.gb32960.v1", 12),
            topic("vehicle.raw.gb32960.v1", 12),
            topic("vehicle.dlq.gb32960.v1", 3),
            topic("vehicle.event.jt808.v1", 12),
            topic("vehicle.raw.jt808.v1", 12),
            topic("vehicle.dlq.jt808.v1", 3),
            topic("vehicle.event.mqtt-yutong.v1", 12),
            topic("vehicle.raw.mqtt-yutong.v1", 12),
            topic("vehicle.dlq.mqtt-yutong.v1", 3));

    Properties props = new Properties();
    props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokers);
    props.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "15000");
    props.put(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "30000");

    try (AdminClient admin = AdminClient.create(props)) {
      Set<String> existing = admin.listTopics().names().get();
      List<NewTopic> missing =
          specs.stream()
              .filter(spec -> !existing.contains(spec.name()))
              .map(spec -> new NewTopic(spec.name(), spec.partitions(), REPLICATION_FACTOR))
              .toList();
      if (!missing.isEmpty()) {
        try {
          admin.createTopics(missing).all().get();
        } catch (ExecutionException e) {
          System.out.println("createTopics result=" + e.getCause());
        }
      }

      Map<String, TopicSpec> byName =
          specs.stream().collect(Collectors.toMap(TopicSpec::name, spec -> spec));
      Map<String, TopicDescription> descriptions =
          admin.describeTopics(byName.keySet()).allTopicNames().get();

      Map<String, NewPartitions> increases =
          descriptions.entrySet().stream()
              .filter(entry -> entry.getValue().partitions().size() < byName.get(entry.getKey()).partitions())
              .collect(
                  Collectors.toMap(
                      Map.Entry::getKey,
                      entry -> NewPartitions.increaseTo(byName.get(entry.getKey()).partitions())));
      if (!increases.isEmpty()) {
        admin.createPartitions(increases).all().get();
        Thread.sleep(Duration.ofSeconds(2).toMillis());
        descriptions = admin.describeTopics(byName.keySet()).allTopicNames().get();
      }

      List<String> lines = new ArrayList<>();
      descriptions.entrySet().stream()
          .sorted(Map.Entry.comparingByKey())
          .forEach(entry -> lines.add(entry.getKey() + " partitions=" + entry.getValue().partitions().size()));
      System.out.println("bootstrap=" + brokers);
      lines.forEach(System.out::println);
    }
  }

  private static TopicSpec topic(String name, int partitions) {
    return new TopicSpec(name, partitions);
  }
}