package com.ankurm.kafkabasics;

import org.apache.kafka.common.utils.Utils;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.kafka.support.SendResult;
import org.springframework.kafka.test.context.EmbeddedKafka;
import org.springframework.test.context.ActiveProfiles;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;
import static org.awaitility.Awaitility.await;

/**
 * Where a record lands, and why the answer is worth knowing before you pick a key.
 *
 * <p>Ordering in Kafka is per partition, never per topic. So the key is not a label &mdash; it is
 * the ordering guarantee, and choosing it is the most consequential design decision in a
 * producer.
 *
 * @see <a href="../../../../../docs/04-keys-and-partitions.md">docs/04-keys-and-partitions.md</a>
 */
@SpringBootTest
@ActiveProfiles("test")
@EmbeddedKafka(topics = { "orders" }, partitions = 3)
class KeysAndPartitionsTest {

	@Autowired
	OrderProducer producer;

	@Autowired
	OrderConsumer consumer;

	@Test
	void theSameKeyAlwaysLandsOnTheSamePartition() throws Exception {
		Map<String, Integer> placement = new LinkedHashMap<>();
		for (String orderId : List.of("o-1", "o-2", "o-3", "o-4", "o-5", "o-6")) {
			SendResult<String, OrderEvent> result =
					this.producer.sendAndWait(OrderEvent.of(orderId, "c-1", "10.00"));
			placement.put(orderId, result.getRecordMetadata().partition());
		}

		System.out.println("=== key -> partition, 3 partitions ===");
		System.out.printf("%-6s %-11s %-24s %s%n", "key", "partition", "murmur2 & 0x7fffffff % 3",
				"Math.abs(murmur2) % 3");
		placement.forEach((key, partition) -> System.out.printf("%-6s %-11d %-24d %d%n", key, partition,
				partitionFor(key), Math.abs(Utils.murmur2(key.getBytes(StandardCharsets.UTF_8))) % 3));

		// The default partitioner is murmur2 of the serialized key, modulo the partition count.
		// It is deterministic and it is not a hash you can change your mind about later:
		// ADDING PARTITIONS REPARTITIONS EVERY KEY, which breaks per-key ordering across the
		// boundary for anything still in flight.
		placement.forEach((key, partition) -> assertThat(partition).isEqualTo(partitionFor(key)));

		// Sending the same key again lands in the same place.
		SendResult<String, OrderEvent> again = this.producer.sendAndWait(OrderEvent.of("o-1", "c-9", "99.00"));
		assertThat(again.getRecordMetadata().partition()).isEqualTo(placement.get("o-1"));
	}

	/**
	 * The real formula. Note {@code & 0x7fffffff} and not {@code Math.abs}: they disagree
	 * whenever murmur2 returns a negative value, because masking the sign bit is not the same
	 * number as negating. Writing it with {@code Math.abs} reproduces the broker's placement for
	 * roughly half of all keys, which is the worst possible failure mode for a test.
	 */
	private static int partitionFor(String key) {
		return (Utils.murmur2(key.getBytes(StandardCharsets.UTF_8)) & 0x7fffffff) % 3;
	}

	@Test
	void everythingProducedIsConsumedWithItsMetadata() {
		this.consumer.clear();
		this.producer.sendAndForget(OrderEvent.of("e-1", "c-1", "1.00"));
		this.producer.sendAndForget(OrderEvent.of("e-2", "c-2", "2.00"));

		await().atMost(Duration.ofSeconds(20))
				.until(() -> this.consumer.received().stream()
						.map((r) -> r.event().orderId()).toList().containsAll(List.of("e-1", "e-2")));

		assertThat(this.consumer.received())
				.extracting((r) -> r.event().orderId()).contains("e-1", "e-2");
		// The key arrives as a header, deserialized by the KEY deserializer, and it is the
		// producer's key - not anything derived from the payload.
		assertThat(this.consumer.received()).allSatisfy(
				(r) -> assertThat(r.key()).isEqualTo(r.event().orderId()));
		assertThat(this.consumer.received()).allSatisfy((r) -> assertThat(r.offset()).isGreaterThanOrEqualTo(0));
	}

}