import {KafkaJS as Confluent, RdKafka} from "@confluentinc/kafka-javascript";
import {Admin, Consumer, Kafka, Producer} from "kafkajs";

const KAFKA_JS_TOPIC = `test-confluent-topic-${Date.now()}`;
const KAFKA_JS_GROUP_ID = `test-confluent-group-${Date.now()}`;

const CONFLUENT_TOPIC = `test-confluent-topic-${Date.now()}`;
const CONFLUENT_GROUP_ID = `test-confluent-group-${Date.now()}`;

describe("partitionLag", () => {
  let kafkaJSKafka: Kafka;
  let kafkaJSAdmin: Admin;
  let kafkaJSConsumer: Consumer;
  let kafkaJSProducer: Producer;

  let confluentKafka: Confluent.Kafka;
  let confluentAdmin: Confluent.Admin;
  let confluentConsumer: Confluent.Consumer;
  let confluentProducer: Confluent.Producer;

  before(async () => {
    kafkaJSKafka = new Kafka({brokers: ["localhost:9092"]});
    kafkaJSAdmin = kafkaJSKafka.admin();
    await kafkaJSAdmin.connect();

    confluentKafka = new Confluent.Kafka({kafkaJS: {brokers: ["localhost:9092"]}});
    confluentAdmin = confluentKafka.admin();
    await confluentAdmin.connect();
  });

  beforeEach(async () => {
    await kafkaJSAdmin.createTopics({topics: [{topic: KAFKA_JS_TOPIC}]});
    kafkaJSProducer = kafkaJSKafka.producer();
    await kafkaJSProducer.connect();

    await confluentAdmin.createTopics({topics: [{topic: CONFLUENT_TOPIC}]});
    confluentProducer = confluentKafka.producer();
    await confluentProducer.connect();
  });

  afterEach(async () => {
    await kafkaJSProducer.disconnect();
    await kafkaJSConsumer?.disconnect();

    await confluentProducer.disconnect();
    await confluentConsumer?.disconnect();
  });

  after(async () => {
    await confluentAdmin.disconnect();
    await confluentProducer.disconnect();
    await confluentConsumer?.disconnect();

    await kafkaJSAdmin.disconnect();
    await kafkaJSProducer.disconnect();
    await kafkaJSConsumer?.disconnect();
  });

  it("reports lag with KafkaJS", async () => {
    for (let i = 0; i < 10; i++) {
      await kafkaJSProducer.send({
        topic: CONFLUENT_TOPIC,
        messages: [{value: "one"}]
      });
    }

    let ready = false;
    let receivedMessages: number = 0;
    kafkaJSConsumer = kafkaJSKafka.consumer({groupId: KAFKA_JS_GROUP_ID});
    kafkaJSConsumer.on(kafkaJSConsumer.events.GROUP_JOIN, (event: any) => {
      ready = true;
    });
    await kafkaJSConsumer.connect();
    await kafkaJSConsumer.subscribe({topic: KAFKA_JS_TOPIC, fromBeginning: true});
    await kafkaJSConsumer.run({
      eachBatch: async ({batch}) => {
        console.log(`Received batch with ${batch.messages.length} messages with highWatermark ${batch.highWatermark} on partition ${batch.partition}`);
        const highWatermark = parseInt(batch.highWatermark);
        for (const message of batch.messages) {
          const offset = parseInt(message.offset);
          console.log(`  Processing offset ${message.offset} which has a lag of ${highWatermark - offset}`);
          receivedMessages++;
        }
      }
    });

    await until(() => ready);

    for (let i = 0; i < 10; i++) {
      await kafkaJSProducer.send({
        topic: CONFLUENT_TOPIC,
        messages: [{value: "one"}]
      });
    }

    await until(() => receivedMessages == 20);
  });

  it("reports lag with Confluent", async () => {
    for (let i = 0; i < 10; i++) {
      await confluentProducer.send({
        topic: CONFLUENT_TOPIC,
        messages: [{value: "one"}]
      });
    }

    let ready = false;
    let receivedMessages: number = 0;
    confluentConsumer = confluentKafka.consumer({
      kafkaJS: {groupId: CONFLUENT_GROUP_ID, fromBeginning: true},
      rebalance_cb: (err: any, assignment: any, consumer: any) => {
        if (err.code !== RdKafka.CODES.ERRORS.ERR__ASSIGN_PARTITIONS) return;
        if (!ready) {
          ready = true;
        }
      }
    });
    await confluentConsumer.connect();
    await confluentConsumer.subscribe({topic: CONFLUENT_TOPIC});
    await confluentConsumer.run({
      eachBatch: async ({batch}) => {
        console.log(`Received batch with ${batch.messages.length} messages with highWatermark ${batch.highWatermark} on partition ${batch.partition}`);
        const highWatermark = parseInt(batch.highWatermark);
        for (const message of batch.messages) {
          const offset = parseInt(message.offset);
          console.log(`  Processing offset ${message.offset} which has a lag of ${highWatermark - offset}`);
          receivedMessages++;
        }
      }
    });

    await until(() => ready);

    for (let i = 0; i < 10; i++) {
      await kafkaJSProducer.send({
        topic: CONFLUENT_TOPIC,
        messages: [{value: "one"}]
      });
    }

    await until(() => receivedMessages == 20);
  });

  async function until(condition: () => boolean) {
    const timeout = 10000;
    const finish = Date.now() + timeout;
    while (Date.now() <= finish) {
      const result = condition();
      if (result) return;
      await new Promise(resolve => setTimeout(resolve, 50));
    }
    throw new Error(`Failed within ${timeout!}ms`);
  }
});