/*
/* Copyright 2018-2026 contributors to the OpenLineage project
/* SPDX-License-Identifier: Apache-2.0
*/

/*
 * 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 io.openlineage.flink.testutils;

import static org.apache.commons.lang3.RandomStringUtils.randomAlphanumeric;

import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.regex.Pattern;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.source.Boundedness;
import org.apache.flink.api.connector.source.Source;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.KafkaSourceBuilder;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.connector.kafka.source.reader.deserializer.KafkaRecordDeserializationSchema;
import org.apache.flink.connector.testframe.external.ExternalSystemSplitDataWriter;
import org.apache.flink.connector.testframe.external.source.DataStreamSourceExternalContext;
import org.apache.flink.connector.testframe.external.source.TestingSourceSettings;
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;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * External context for testing {@link KafkaSource}.
 *
 * <p>This code has been copied from flink-kafka-connector
 * https://github.com/apache/flink-connector-kafka/tree/main/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/testutils
 */
public class KafkaSourceExternalContext implements DataStreamSourceExternalContext<String> {

  private static final Logger LOG = LoggerFactory.getLogger(KafkaSourceExternalContext.class);

  private static final String TOPIC_NAME_PREFIX = "kafka-test-topic-";
  private static final Pattern TOPIC_NAME_PATTERN = Pattern.compile(TOPIC_NAME_PREFIX + ".*");
  private static final String GROUP_ID_PREFIX = "kafka-source-external-context-";
  private static final int NUM_RECORDS_UPPER_BOUND = 500;
  private static final int NUM_RECORDS_LOWER_BOUND = 100;

  private final List<URL> connectorJarPaths;
  private final String bootstrapServers;
  private final String topicName;
  private final SplitMappingMode splitMappingMode;
  private final AdminClient adminClient;
  private final List<KafkaPartitionDataWriter> writers = new ArrayList<>();

  protected KafkaSourceExternalContext(
      String bootstrapServers, SplitMappingMode splitMappingMode, List<URL> connectorJarPaths) {
    this.connectorJarPaths = connectorJarPaths;
    this.bootstrapServers = bootstrapServers;
    this.topicName = randomize(TOPIC_NAME_PREFIX);
    this.splitMappingMode = splitMappingMode;
    this.adminClient = createAdminClient();
  }

  @Override
  public List<URL> getConnectorJarPaths() {
    return this.connectorJarPaths;
  }

  @Override
  public Source<String, ?, ?> createSource(TestingSourceSettings sourceSettings) {
    final KafkaSourceBuilder<String> builder = KafkaSource.builder();

    builder
        .setBootstrapServers(bootstrapServers)
        .setTopicPattern(TOPIC_NAME_PATTERN)
        .setGroupId(randomize(GROUP_ID_PREFIX))
        .setDeserializer(KafkaRecordDeserializationSchema.valueOnly(StringDeserializer.class));

    if (sourceSettings.getBoundedness().equals(Boundedness.BOUNDED)) {
      builder.setBounded(OffsetsInitializer.latest());
    }

    return builder.build();
  }

  @Override
  public ExternalSystemSplitDataWriter<String> createSourceSplitDataWriter(
      TestingSourceSettings sourceSettings) {
    KafkaPartitionDataWriter writer;
    try {
      switch (splitMappingMode) {
        case TOPIC:
          writer = createSinglePartitionTopic(writers.size());
          break;
        case PARTITION:
          writer = scaleOutTopic(this.topicName);
          break;
        default:
          throw new IllegalArgumentException("Split mode should be either TOPIC or PARTITION");
      }
    } catch (Exception e) {
      throw new RuntimeException("Failed to create new splits", e);
    }
    writers.add(writer);
    return writer;
  }

  @Override
  public List<String> generateTestData(
      TestingSourceSettings sourceSettings, int splitIndex, long seed) {
    Random random = new Random(seed);
    int recordNum =
        random.nextInt(NUM_RECORDS_UPPER_BOUND - NUM_RECORDS_LOWER_BOUND) + NUM_RECORDS_LOWER_BOUND;
    List<String> records = new ArrayList<>(recordNum);

    for (int i = 0; i < recordNum; i++) {
      int stringLength = random.nextInt(50) + 1;
      records.add(splitIndex + "-" + randomAlphanumeric(stringLength));
    }

    return records;
  }

  @Override
  public TypeInformation<String> getProducedType() {
    return TypeInformation.of(String.class);
  }

  @Override
  public void close() throws Exception {
    final List<String> topics = new ArrayList<>();
    writers.forEach(
        writer -> {
          topics.add(writer.getTopicPartition().topic());
          writer.close();
        });
    adminClient.deleteTopics(topics).all().get();
  }

  @Override
  public String toString() {
    return "KafkaSource-" + splitMappingMode.toString();
  }

  private String randomize(String prefix) {
    return prefix + ThreadLocalRandom.current().nextLong(0, Long.MAX_VALUE);
  }

  private AdminClient createAdminClient() {
    Properties config = new Properties();
    config.setProperty(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
    return AdminClient.create(config);
  }

  private KafkaPartitionDataWriter createSinglePartitionTopic(int topicIndex) throws Exception {
    String newTopicName = topicName + "-" + topicIndex;
    LOG.info("Creating topic '{}'", newTopicName);
    adminClient
        .createTopics(Collections.singletonList(new NewTopic(newTopicName, 1, (short) 1)))
        .all()
        .get();
    return new KafkaPartitionDataWriter(
        getKafkaProducerProperties(topicIndex), new TopicPartition(newTopicName, 0));
  }

  private KafkaPartitionDataWriter scaleOutTopic(String topicName) throws Exception {
    final Set<String> topics = adminClient.listTopics().names().get();
    if (topics.contains(topicName)) {
      final Map<String, TopicDescription> topicDescriptions =
          adminClient.describeTopics(Collections.singletonList(topicName)).allTopicNames().get();
      final int numPartitions = topicDescriptions.get(topicName).partitions().size();
      LOG.info("Creating partition {} for topic '{}'", numPartitions + 1, topicName);
      adminClient
          .createPartitions(
              Collections.singletonMap(topicName, NewPartitions.increaseTo(numPartitions + 1)))
          .all()
          .get();
      return new KafkaPartitionDataWriter(
          getKafkaProducerProperties(numPartitions), new TopicPartition(topicName, numPartitions));
    } else {
      LOG.info("Creating topic '{}'", topicName);
      adminClient
          .createTopics(Collections.singletonList(new NewTopic(topicName, 1, (short) 1)))
          .all()
          .get();
      return new KafkaPartitionDataWriter(
          getKafkaProducerProperties(0), new TopicPartition(topicName, 0));
    }
  }

  private Properties getKafkaProducerProperties(int producerId) {
    Properties kafkaProducerProperties = new Properties();
    kafkaProducerProperties.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
    kafkaProducerProperties.setProperty(
        ProducerConfig.CLIENT_ID_CONFIG,
        String.join(
            "-",
            "flink-kafka-split-writer",
            Integer.toString(producerId),
            Long.toString(ThreadLocalRandom.current().nextLong(Long.MAX_VALUE))));
    kafkaProducerProperties.setProperty(
        ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
    kafkaProducerProperties.setProperty(
        ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
    return kafkaProducerProperties;
  }

  /** Mode of mapping split to Kafka components. */
  public enum SplitMappingMode {
    /** Use a single-partitioned topic as a split. */
    TOPIC,

    /** Use a partition in topic as a split. */
    PARTITION
  }
}
