import org.apache.kafka.clients.admin.*;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.config.ConfigResource;
import org.apache.kafka.common.serialization.*;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;

/**
 * Share groups (KIP-932): two consumers share ONE partition, which a classic consumer group cannot do.
 *   java -cp "kafka/libs/*" KafkaShareLab.java
 */
public class KafkaShareLab {
    public static void main(String[] args) throws Exception {
        String run = Long.toString(System.currentTimeMillis() % 100000), topic = "queue-" + run, group = "share-" + run;
        Properties base = new Properties();
        base.put("bootstrap.servers", "localhost:9092");
        try (Admin admin = Admin.create(base)) {
            admin.createTopics(List.of(new NewTopic(topic, 1, (short) 1))).all().get();
            ConfigResource gr = new ConfigResource(ConfigResource.Type.GROUP, group);
            admin.incrementalAlterConfigs(Map.of(gr, List.of(new AlterConfigOp(
                    new ConfigEntry("share.auto.offset.reset", "earliest"), AlterConfigOp.OpType.SET)))).all().get();
        }
        ConcurrentHashMap<String, List<String>> got = new ConcurrentHashMap<>();
        CountDownLatch done = new CountDownLatch(2);
        ExecutorService ex = Executors.newFixedThreadPool(2);
        for (String name : List.of("worker-A", "worker-B")) ex.submit(() -> consume(base, group, topic, name, got, done));

        Thread.sleep(8_000); // let both workers join the share group before producing
        Properties pp = new Properties();
        pp.putAll(base);
        pp.put("key.serializer", StringSerializer.class.getName());
        pp.put("value.serializer", StringSerializer.class.getName());
        try (var p = new KafkaProducer<String, String>(pp)) {
            for (int i = 1; i <= 20; i++) {
                p.send(new ProducerRecord<>(topic, "k", String.format("job-%02d", i))).get();
                Thread.sleep(60);
            }
        }
        done.await();
        ex.shutdown();
        int total = got.values().stream().mapToInt(List::size).sum();
        System.out.println("[6] 1 partition, 2 share consumers, " + total + " records delivered");
        System.out.println("    both workers got work: " + (got.size() == 2));
        System.out.println("    no record handled twice: "
                + (got.values().stream().flatMap(List::stream).distinct().count() == total));
    }

    static void consume(Properties base, String group, String topic, String name,
                        ConcurrentHashMap<String, List<String>> got, CountDownLatch done) {
        Properties cp = new Properties();
        cp.putAll(base);
        cp.put("group.id", group);
        cp.put("key.deserializer", StringDeserializer.class.getName());
        cp.put("value.deserializer", StringDeserializer.class.getName());
        cp.put("share.acknowledgement.mode", "explicit");
        cp.put("max.poll.records", "2");
        try (var c = new KafkaShareConsumer<String, String>(cp)) {
            c.subscribe(List.of(topic));
            long end = System.currentTimeMillis() + 40_000;
            int idle = 0;
            while (System.currentTimeMillis() < end && idle < 8) {
                var recs = c.poll(Duration.ofMillis(500));
                if (recs.isEmpty()) { if (!got.isEmpty()) idle++; continue; }
                idle = 0;
                for (var r : recs) {
                    got.computeIfAbsent(name, k -> new CopyOnWriteArrayList<>()).add(r.value());
                    Thread.sleep(150); // pretend to work
                    c.acknowledge(r, AcknowledgeType.ACCEPT);
                }
                c.commitSync();
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            done.countDown();
        }
    }
}