package io.vertx.kafka.client.tests;

import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.MockConsumer;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.junit.Test;
import org.junit.runner.RunWith;

import io.vertx.core.Vertx;
import io.vertx.ext.unit.Async;
import io.vertx.ext.unit.TestContext;
import io.vertx.ext.unit.junit.VertxUnitRunner;
import io.vertx.kafka.client.common.TopicPartition;
import io.vertx.kafka.client.consumer.KafkaConsumer;
import io.vertx.kafka.client.consumer.KafkaReadStream;
import io.vertx.kafka.client.consumer.impl.KafkaConsumerImpl;

@RunWith(VertxUnitRunner.class)
public class KafkaReadStreamMockTest extends KafkaTestBase {

    private LinkedList<ConsumerRecord<String,String>> recordsMock = new LinkedList<>();

    private int SEND_BATCH = 5;
    private int TOTAL_MESSAGES = 400;
    private final String TOPIC = "topic";
    private Long timer = null;

    private void initRecords(){
        int numMessages = TOTAL_MESSAGES;
        for (int i = 0;i < numMessages;i++) {
            String key = "key-" + i;
            String value = "value-" + i;
            recordsMock.add(new ConsumerRecord<String,String>(TOPIC, 0, i, key, value));
        }
    }

    private MockConsumer<String, String> createMockConsumer(){
        MockConsumer<String, String> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);

        Map<org.apache.kafka.common.TopicPartition, Long> beginningOffsets = new HashMap<>();
        beginningOffsets.put(new org.apache.kafka.common.TopicPartition(TOPIC, 0), 0L);
        consumer.updateBeginningOffsets(beginningOffsets);
        return consumer;
    }

    private void sendNextBatch(MockConsumer<String, String> consumer){
        for(int i=0;i<SEND_BATCH && recordsMock.size()>0;i++)
        consumer.addRecord(recordsMock.pop());

    }

    @Test
    public void shouldNotLoseMessages(TestContext ctx){
        Vertx vertx = Vertx.vertx();

        Async done = ctx.async();

        initRecords();

        MockConsumer<String, String> consumer = createMockConsumer();
        KafkaReadStream<String, String> readStream =  KafkaReadStream.create(vertx, consumer);
        KafkaConsumer<String, String> consumerVertx = new KafkaConsumerImpl<>(readStream);


        AtomicLong partitionOffset = new AtomicLong(-1);


        consumerVertx.handler((r)->{
            long offset = r.offset();

            partitionOffset.addAndGet(1);
            ctx.assertEquals(partitionOffset.get(), offset);

            if(offset == TOTAL_MESSAGES-1){
                consumerVertx.close();
                done.complete();
            } else {

                if(timer!=null) vertx.cancelTimer(timer);
                timer = vertx.setTimer(5, (t)->{
                    consumerVertx.pause();
                    vertx.getOrCreateContext().runOnContext((t1)->{
                        consumerVertx.commit();
                        consumerVertx.resume();
                        sendNextBatch(consumer);
                        // sends two batches of messages
                        vertx.getOrCreateContext().runOnContext((t2)->{
                            sendNextBatch(consumer);
                        });
                    });
                });

            }

        });

        consumerVertx.exceptionHandler(t->ctx.fail(t));


        Set<TopicPartition> partitions = new LinkedHashSet<>();
        partitions.add(new TopicPartition(TOPIC, 0));

        consumerVertx.assign(partitions, (h)->{
            sendNextBatch(consumer);
        });

    }

}
