package com.lingniu.ingest.sink.kafka;

import com.lingniu.ingest.api.consumer.EnvelopeConsumerRecord;
import com.lingniu.ingest.api.consumer.EnvelopeBatchIngestor;
import com.lingniu.ingest.api.consumer.EnvelopeConsumerProcessor;
import com.lingniu.ingest.api.consumer.EnvelopeDeadLetterRecord;
import com.lingniu.ingest.api.consumer.EnvelopeIngestResult;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.MockConsumer;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.common.TopicPartition;
import org.junit.jupiter.api.Test;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;

class KafkaEnvelopeConsumerWorkerTest {

    @Test
    void pollsKafkaRecordsAndDispatchesThemToConfiguredTopicProcessor() {
        MockConsumer<String, byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
        TopicPartition partition = new TopicPartition("vehicle.realtime", 0);
        consumer.assign(List.of(partition));
        consumer.updateBeginningOffsets(Map.of(partition, 0L));
        consumer.addRecord(new ConsumerRecord<>("vehicle.realtime", 0, 12L, "VIN001", new byte[]{0x01, 0x02}));
        AtomicReference<EnvelopeDeadLetterRecord> captured = new AtomicReference<>();
        EnvelopeConsumerProcessor processor = new EnvelopeConsumerProcessor(
                "vehicle-state",
                bytes -> EnvelopeIngestResult.invalid("bad envelope"),
                captured::set);

        int processed = new KafkaEnvelopeConsumerWorker(
                consumer, Map.of("vehicle.realtime", processor)).pollOnce(Duration.ZERO);

        assertThat(processed).isEqualTo(1);
        assertThat(captured.get()).satisfies(record -> {
            assertThat(record.topic()).isEqualTo("vehicle.realtime");
            assertThat(record.partition()).isEqualTo(0);
            assertThat(record.offset()).isEqualTo(12L);
            assertThat(record.key()).isEqualTo("VIN001");
            assertThat(record.payload()).containsExactly(0x01, 0x02);
        });
    }

    @Test
    void pollsKafkaRecordsAndDispatchesBatchWhenProcessorSupportsIt() {
        MockConsumer<String, byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
        TopicPartition partition = new TopicPartition("vehicle.raw", 0);
        consumer.assign(List.of(partition));
        consumer.updateBeginningOffsets(Map.of(partition, 0L));
        consumer.addRecord(new ConsumerRecord<>("vehicle.raw", 0, 12L, "VIN001", new byte[]{0x01}));
        consumer.addRecord(new ConsumerRecord<>("vehicle.raw", 0, 13L, "VIN002", new byte[]{0x02}));
        AtomicInteger batchCalls = new AtomicInteger();
        EnvelopeConsumerProcessor processor = new EnvelopeConsumerProcessor(
                "event-history",
                new EnvelopeBatchIngestor() {
                    @Override
                    public EnvelopeIngestResult tryIngest(byte[] kafkaValue) {
                        throw new AssertionError("single-record ingest should not be used for a batch-capable processor");
                    }

                    @Override
                    public List<EnvelopeIngestResult> tryIngestAll(List<byte[]> kafkaValues) {
                        batchCalls.incrementAndGet();
                        assertThat(kafkaValues).hasSize(2);
                        return List.of(
                                EnvelopeIngestResult.stored("event-1", "VIN001"),
                                EnvelopeIngestResult.invalid("bad envelope"));
                    }
                },
                record -> {});

        int processed = new KafkaEnvelopeConsumerWorker(
                consumer, Map.of("vehicle.raw", processor)).pollOnce(Duration.ZERO);

        assertThat(processed).isEqualTo(2);
        assertThat(batchCalls).hasValue(1);
    }
}
