import org.apache.kafka.clients.admin.*;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;

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

/**
 * Companion for "Top 40 Kafka Interview Questions". Needs a broker on localhost:9092
 * and kafka-clients 4.4.0 on the classpath:
 *   java -cp "kafka/libs/*" KafkaInterviewLab.java
 */
public class KafkaInterviewLab {
    static final String BS = "localhost:9092";

    static Properties base() { Properties p = new Properties(); p.put("bootstrap.servers", BS); return p; }

    static KafkaProducer<String, String> producer(Map<String, Object> extra) {
        Properties p = base();
        p.put("key.serializer", StringSerializer.class.getName());
        p.put("value.serializer", StringSerializer.class.getName());
        p.putAll(extra);
        return new KafkaProducer<>(p);
    }

    static KafkaConsumer<String, String> consumer(String group, Map<String, Object> extra) {
        Properties p = base();
        p.put("group.id", group);
        p.put("key.deserializer", StringDeserializer.class.getName());
        p.put("value.deserializer", StringDeserializer.class.getName());
        p.put("auto.offset.reset", "earliest");
        p.put("enable.auto.commit", "false");
        p.putAll(extra);
        return new KafkaConsumer<>(p);
    }

    public static void main(String[] args) throws Exception {
        String run = Long.toString(System.currentTimeMillis() % 100000);
        try (Admin admin = Admin.create(base())) {
            keysAndPartitions(admin, run);
            offsetsAndLag(admin, run);
            rebalance(admin, run);
            transactions(admin, run);
            deadLetter(admin, run);
        }
    }

    // 1. same key -> same partition, order preserved only inside a partition
    static void keysAndPartitions(Admin admin, String run) throws Exception {
        String topic = "orders-" + run;
        admin.createTopics(List.of(new NewTopic(topic, 3, (short) 1))).all().get();
        System.out.println("[1] key -> partition (topic has 3 partitions)");
        try (var p = producer(Map.of())) {
            for (String key : List.of("order-1", "order-2", "order-3", "order-1", "order-2", "order-1")) {
                RecordMetadata m = p.send(new ProducerRecord<>(topic, key, "v")).get();
                System.out.printf("    key=%s partition=%d offset=%d%n", key, m.partition(), m.offset());
            }
        }
    }

    // 2. committed offset vs log end offset = lag
    static void offsetsAndLag(Admin admin, String run) throws Exception {
        String topic = "lag-" + run, group = "g-lag-" + run;
        admin.createTopics(List.of(new NewTopic(topic, 1, (short) 1))).all().get();
        try (var p = producer(Map.of())) {
            for (int i = 0; i < 10; i++) p.send(new ProducerRecord<>(topic, "k", "m" + i)).get();
        }
        try (var c = consumer(group, Map.of("max.poll.records", 4))) {
            c.subscribe(List.of(topic));
            ConsumerRecords<String, String> recs = ConsumerRecords.empty();
            while (recs.isEmpty()) recs = c.poll(Duration.ofMillis(500));
            c.commitSync();
            TopicPartition tp = new TopicPartition(topic, 0);
            long committed = c.committed(Set.of(tp)).get(tp).offset();
            long end = c.endOffsets(Set.of(tp)).get(tp);
            System.out.println("[2] polled=" + recs.count() + " committed=" + committed
                    + " logEnd=" + end + " lag=" + (end - committed));
        }
    }

    // 3. two consumers in one group: partitions are split, then rebalanced when one leaves
    static void rebalance(Admin admin, String run) throws Exception {
        String topic = "rb-" + run, group = "g-rb-" + run;
        admin.createTopics(List.of(new NewTopic(topic, 4, (short) 1))).all().get();
        var c1 = consumer(group, Map.of());
        c1.subscribe(List.of(topic));
        awaitAssignment(c1);
        System.out.println("[3] c1 alone      : " + parts(c1));
        var c2 = consumer(group, Map.of());
        c2.subscribe(List.of(topic));
        for (int i = 0; i < 20 && (c1.assignment().size() != 2 || c2.assignment().size() != 2); i++) {
            c1.poll(Duration.ofMillis(300));
            c2.poll(Duration.ofMillis(300));
        }
        System.out.println("[3] c1 + c2 joined: c1=" + parts(c1) + " c2=" + parts(c2));
        c2.close();
        for (int i = 0; i < 20 && c1.assignment().size() != 4; i++) c1.poll(Duration.ofMillis(300));
        System.out.println("[3] c2 left       : c1=" + parts(c1));
        c1.close();
    }
    static void awaitAssignment(KafkaConsumer<String, String> c) {
        for (int i = 0; i < 20 && c.assignment().isEmpty(); i++) c.poll(Duration.ofMillis(300));
    }

    /** Poll until the deadline; group join can take a few seconds on a fresh group. */
    static List<ConsumerRecord<String, String>> drain(KafkaConsumer<String, String> c, int seconds, int expectAtLeast) {
        List<ConsumerRecord<String, String>> out = new ArrayList<>();
        long end = System.currentTimeMillis() + seconds * 1000L;
        while (System.currentTimeMillis() < end && out.size() < expectAtLeast) {
            c.poll(Duration.ofMillis(300)).forEach(out::add);
        }
        // one more short poll so a late extra record (e.g. an aborted one) would still show up
        c.poll(Duration.ofMillis(500)).forEach(out::add);
        return out;
    }
    static String parts(KafkaConsumer<String, String> c) {
        return c.assignment().stream().map(TopicPartition::partition).sorted().collect(Collectors.toList()).toString();
    }

    // 4. transactions: aborted records are invisible to read_committed
    static void transactions(Admin admin, String run) throws Exception {
        String topic = "tx-" + run;
        admin.createTopics(List.of(new NewTopic(topic, 1, (short) 1))).all().get();
        try (var p = producer(Map.of("transactional.id", "tx-" + run, "enable.idempotence", true))) {
            p.initTransactions();
            p.beginTransaction();
            p.send(new ProducerRecord<>(topic, "k", "committed-1")).get();
            p.commitTransaction();
            p.beginTransaction();
            p.send(new ProducerRecord<>(topic, "k", "aborted-2")).get();
            p.abortTransaction();
        }
        for (String level : List.of("read_uncommitted", "read_committed")) {
            try (var c = consumer("g-tx-" + level + "-" + run, Map.of("isolation.level", level))) {
                c.subscribe(List.of(topic));
                List<String> values = drain(c, 20, level.equals("read_committed") ? 1 : 2).stream()
                        .map(ConsumerRecord::value).collect(Collectors.toList());
                System.out.println("[4] " + level + " sees " + values);
            }
        }
    }

    // 5. poison pill -> dead-letter topic, then keep going
    static void deadLetter(Admin admin, String run) throws Exception {
        String in = "in-" + run, dlt = "in-" + run + ".DLT";
        admin.createTopics(List.of(new NewTopic(in, 1, (short) 1), new NewTopic(dlt, 1, (short) 1))).all().get();
        try (var p = producer(Map.of())) {
            for (String v : List.of("10", "20", "not-a-number", "40")) p.send(new ProducerRecord<>(in, "k", v)).get();
        }
        int sum = 0, dead = 0;
        try (var c = consumer("g-dlt-" + run, Map.of()); var p = producer(Map.of())) {
            c.subscribe(List.of(in));
            for (ConsumerRecord<String, String> r : drain(c, 20, 4)) {
                try {
                    sum += Integer.parseInt(r.value());
                } catch (NumberFormatException e) {
                    ProducerRecord<String, String> d = new ProducerRecord<>(dlt, r.key(), r.value());
                    d.headers().add("x-original-offset", Long.toString(r.offset()).getBytes());
                    d.headers().add("x-exception", e.getClass().getSimpleName().getBytes());
                    p.send(d).get();
                    dead++;
                }
            }
            c.commitSync();
        }
        System.out.println("[5] processed sum=" + sum + ", sent to DLT=" + dead);
        try (var c = consumer("g-dlt-read-" + run, Map.of())) {
            c.subscribe(List.of(dlt));
            for (ConsumerRecord<String, String> r : drain(c, 20, 1)) {
                System.out.println("    DLT record value=" + r.value()
                        + " original-offset=" + new String(r.headers().lastHeader("x-original-offset").value())
                        + " exception=" + new String(r.headers().lastHeader("x-exception").value()));
            }
        }
    }
}