import { Injectable, Logger } from "@nestjs/common";
import { KafkaConsumer, Producer, ConsumerGlobalConfig } from "node-rdkafka";
import {
  SchemaRegistry,
  SchemaType,
  avdlToAVSCAsync,
} from "@kafkajs/confluent-schema-registry";
import { ConfigService } from "@nestjs/config";
import path from "path";
import { Observable } from "rxjs";
import { randomUUID } from "crypto";
import { readFileSync } from "fs";

let producerPromise: Promise<Producer>;

@Injectable()
export class KafkaService {
  constructor(private configService: ConfigService) {}

  schemaIds: Record<string, Promise<number>> = {};

  getKafkaConfig = () => {
    return {
      "metadata.broker.list": this.configService.get("kafka.host"),
    };
  };

  private removeBranchObjects(obj: Record<string, any>) {
    if (typeof obj !== "object" || !obj) return obj;
    if (
      Object.keys(obj)?.[0].startsWith(
        this.configService.get("kafka.avroNamespace")
      )
    ) {
      return Object.values(obj)[0];
    }

    return Object.fromEntries(
      Object.entries(obj).map(([k, v]) => {
        if (typeof v === "object") {
          return [k, this.removeBranchObjects(v)];
        } else {
          return [k, v];
        }
      })
    );
  }

  getKafkaProducer() {
    if (!producerPromise) {
      producerPromise = new Promise((resolve) => {
        const producer = new Producer(this.getKafkaConfig());
        producer.connect();
        producer.on("ready", () => {
          resolve(producer);
        });
      });
    }
    return producerPromise;
  }

  async getSchemaId(
    topicName: string,
    avdlPath: string,
    options?: { isKey?: boolean }
  ) {
    if (!this.schemaIds[topicName + (options?.isKey ? "-key" : "-value")]) {
      if (options?.isKey) {
        this.schemaIds[`${topicName}-key`] = this.getKafkaRegistry()
          .register(
            {
              type: SchemaType.AVRO,
              schema: JSON.stringify({
                namespace: this.configService.get("kafka.avroNamespace"),
                ...JSON.parse(
                  readFileSync(path.join(__dirname, "../key.json"), "utf-8")
                ),
              }),
            },
            { subject: topicName + (options?.isKey ? "-key" : "-value") }
          )
          .then((res) => {
            return res.id;
          });
        return this.schemaIds[`${topicName}-key`];
      }

      const schema = await avdlToAVSCAsync(avdlPath);
      const jsonSchema = JSON.stringify(schema);
      this.schemaIds[topicName + (options?.isKey ? "-key" : "-value")] =
        this.getKafkaRegistry()
          .register(
            {
              type: SchemaType.AVRO,
              schema: jsonSchema,
            },
            { subject: topicName + (options?.isKey ? "-key" : "-value") }
          )
          .then((res) => {
            return res.id;
          });
    }
    return this.schemaIds[topicName + (options?.isKey ? "-key" : "-value")];
  }

  schemaRegistry: SchemaRegistry;
  getKafkaRegistry() {
    if (!this.schemaRegistry) {
      this.schemaRegistry = new SchemaRegistry({
        host: this.configService.get("kafka.registryUrl"),
      });
    }
    return this.schemaRegistry;
  }

  async produceKafkaMessage<T extends { id: string }>(
    topic: string,
    messageValue: T,
    avdlPath: string,
    messageKey = { id: messageValue.id },
    avdlKeyPath?: string
  ): Promise<boolean> {
    const schemaId = await this.getSchemaId(topic, avdlPath);

    const value = messageValue
      ? await this.getKafkaRegistry().encode(schemaId, messageValue)
      : null;

    const keySchemaId = await this.getSchemaId(topic, avdlKeyPath, {
      isKey: true,
    });

    const key = await this.getKafkaRegistry().encode(keySchemaId, messageKey);

    Logger.log("produce", topic, 0, value, key);

    const result = await (
      await this.getKafkaProducer()
    ).produce(topic, 0, value, key);

    return result;
  }

  getReadStream<Value, Key = { id: string }>(
    topic: string,
    groupId: string = randomUUID(),
    autoOffsetRest:
      | "smallest"
      | "earliest"
      | "beginning"
      | "largest"
      | "latest"
      | "end"
      | "error" = "beginning"
  ) {
    return new Observable<{ value: Value; key: Key }>((observer) => {
      const readStream = KafkaConsumer.createReadStream(
        {
          ...this.getKafkaConfig(),
          "group.id": groupId,
          "socket.keepalive.enable": true,
          "enable.auto.commit": false,
        },
        {
          "auto.offset.reset": autoOffsetRest,
        },
        {
          topics: topic,
          waitInterval: 0,
          objectMode: true,
        }
      );
      readStream.on("data", async ({ key, value }) => {
        if (value) {
          value = await this.getKafkaRegistry().decode(value, {});
          value = this.removeBranchObjects(value);
        }
        if (key) {
          key = await this.getKafkaRegistry().decode(key);
        }
        observer.next({ key, value });
      });
      readStream.on("error", (e) => {
        observer.error(e);
      });
      return () => {
        readStream.close();
      };
    });
  }
}