"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.ConfluentKafkaInstrumentation = void 0;
const instrumentation_1 = require("@opentelemetry/instrumentation");
const api_1 = require("@opentelemetry/api");
const types_1 = require("node:util/types");
const incubating_1 = require("@opentelemetry/semantic-conventions/incubating");
// @ts-ignore
const package_json_1 = require("../package.json");
const KAFKAJS_PROMISE_API_CONTEXT_KEY = Symbol.for("kafkajs-promise-api");
class ConfluentKafkaInstrumentation extends instrumentation_1.InstrumentationBase {
    constructor(config = {}) {
        super(package_json_1.name, package_json_1.version, config);
    }
    init() {
        const module = new instrumentation_1.InstrumentationNodeModuleDefinition("@confluentinc/kafka-javascript", [">=1.3.0"], this.patch.bind(this), this.unpatch.bind(this));
        return module;
    }
    patch(moduleExports) {
        api_1.diag.debug("confluent-kafka instrumentation: applying patch");
        moduleExports = moduleExports?.default ?? moduleExports;
        this.unpatch(moduleExports);
        this._wrap(moduleExports?.Producer.prototype, "produce", this._patchProduce());
        this._wrap(moduleExports?.Producer.prototype, "on", this._patchProducerOn());
        this._wrap(moduleExports?.KafkaConsumer.prototype, "on", this._patchOn());
        if (moduleExports?.KafkaJS?.Kafka) {
            this._wrap(moduleExports.KafkaJS.Kafka.prototype, "producer", this._patchKafkaProducer());
            this._wrap(moduleExports.KafkaJS.Kafka.prototype, "consumer", this._patchKafkaConsumer());
        }
        return moduleExports;
    }
    unpatch(moduleExports) {
        api_1.diag.debug("confluent-kafka instrumentation: un-patching");
        moduleExports = moduleExports?.default ?? moduleExports;
        if ((0, instrumentation_1.isWrapped)(moduleExports?.Producer?.prototype.produce)) {
            this._unwrap(moduleExports.Producer.prototype, "produce");
        }
        if ((0, instrumentation_1.isWrapped)(moduleExports?.Producer?.prototype.on)) {
            this._unwrap(moduleExports.Producer.prototype, "on");
        }
        if ((0, instrumentation_1.isWrapped)(moduleExports?.KafkaConsumer?.prototype.on)) {
            this._unwrap(moduleExports.KafkaConsumer.prototype, "on");
        }
        if (moduleExports?.KafkaJS?.Kafka) {
            if ((0, instrumentation_1.isWrapped)(moduleExports.KafkaJS.Kafka.prototype.producer)) {
                this._unwrap(moduleExports.KafkaJS.Kafka.prototype, "producer");
            }
            if ((0, instrumentation_1.isWrapped)(moduleExports.KafkaJS.Kafka.prototype.consumer)) {
                this._unwrap(moduleExports.KafkaJS.Kafka.prototype, "consumer");
            }
        }
    }
    _patchOn() {
        const self = this;
        return (original) => {
            return function (ev, originalListener) {
                const eventName = ev;
                if (eventName !== "data") {
                    return original.apply(this, [ev, originalListener]);
                }
                const wrappedListener = function (data) {
                    const headerCarrier = {};
                    if (data?.headers) {
                        data.headers.forEach((header) => {
                            Object.entries(header).forEach(([key, value]) => {
                                if (value !== undefined && value !== null) {
                                    headerCarrier[key] = Buffer.isBuffer(value)
                                        ? value.toString("utf8")
                                        : typeof value === "string"
                                            ? value
                                            : String(value);
                                }
                            });
                        });
                    }
                    const ctx = Object.keys(headerCarrier).length > 0
                        ? api_1.propagation.extract(api_1.ROOT_CONTEXT, headerCarrier, api_1.defaultTextMapGetter)
                        : api_1.ROOT_CONTEXT;
                    const span = self.tracer.startSpan(data.topic ?? eventName, {
                        kind: api_1.SpanKind.CONSUMER,
                        attributes: {
                            [incubating_1.ATTR_MESSAGING_SYSTEM]: incubating_1.MESSAGING_SYSTEM_VALUE_KAFKA,
                            [incubating_1.ATTR_MESSAGING_DESTINATION_NAME]: data.topic,
                            [incubating_1.ATTR_MESSAGING_OPERATION_NAME]: "received",
                            [incubating_1.ATTR_MESSAGING_OPERATION_TYPE]: incubating_1.MESSAGING_OPERATION_TYPE_VALUE_RECEIVE,
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_KEY]: data.key?.toString(),
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_TOMBSTONE]: data?.key && data.value === null ? true : undefined,
                            [incubating_1.ATTR_MESSAGING_KAFKA_OFFSET]: data.offset.toString(),
                            [incubating_1.ATTR_MESSAGING_CONSUMER_GROUP_NAME]: data.headers
                                ?.find((h) => Object.keys(h).some(key => key === "group-id"))?.["group-id"]?.toString(),
                            [incubating_1.ATTR_MESSAGING_DESTINATION_PARTITION_ID]: data.partition.toString(),
                            [incubating_1.ATTR_MESSAGING_CLIENT_ID]: data.headers
                                ?.find((h) => Object.keys(h).some(key => key === "client-id"))?.["client-id"]?.toString(),
                        },
                    }, ctx);
                    if (self._config.consumerHook) {
                        self._config.consumerHook(span, {
                            topic: data?.topic,
                            message: data,
                        });
                    }
                    return api_1.context.with(api_1.trace.setSpan(api_1.context.active(), span), () => self.endSpan(() => originalListener.apply(this, [data]), span));
                };
                return original.apply(this, [ev, wrappedListener]);
            };
        };
    }
    _patchProducerOn() {
        return (original) => {
            return function (ev, originalListener) {
                return original.apply(this, [ev, originalListener]);
            };
        };
    }
    _patchProduce() {
        const self = this;
        return (original) => {
            return function (topic, partition, message, key, timestamp, opaque, headers) {
                const activeContext = api_1.context.active();
                const isPromiseApiCall = activeContext.getValue(KAFKAJS_PROMISE_API_CONTEXT_KEY) === true;
                if (isPromiseApiCall) {
                    return original.apply(this, [
                        topic,
                        partition,
                        message,
                        key,
                        timestamp,
                        opaque,
                        headers,
                    ]);
                }
                const span = self.tracer.startSpan(topic, {
                    kind: api_1.SpanKind.PRODUCER,
                    attributes: {
                        [incubating_1.ATTR_MESSAGING_SYSTEM]: incubating_1.MESSAGING_SYSTEM_VALUE_KAFKA,
                        [incubating_1.ATTR_MESSAGING_DESTINATION_NAME]: topic,
                        [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_KEY]: key?.toString(),
                        [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_TOMBSTONE]: key && message === null ? true : undefined,
                        [incubating_1.ATTR_MESSAGING_DESTINATION_PARTITION_ID]: partition?.toString(),
                        [incubating_1.ATTR_MESSAGING_OPERATION_NAME]: "produce",
                        [incubating_1.ATTR_MESSAGING_OPERATION_TYPE]: incubating_1.MESSAGING_OPERATION_TYPE_VALUE_SEND,
                        [incubating_1.ATTR_MESSAGING_CLIENT_ID]: (headers ?? [])
                            .find((h) => Object.keys(h).some(key => key === "client-id"))?.["client-id"]?.toString(),
                    },
                });
                const ctx = api_1.trace.setSpan(api_1.context.active(), span);
                const carrier = {};
                api_1.propagation.inject(ctx, carrier);
                const kafkaHeaders = headers ? headers.map(h => ({ ...h })) : [];
                Object.entries(carrier).forEach(([k, v]) => {
                    if (!v)
                        return;
                    const headerValue = typeof v === "string" ? v : String(v);
                    const existingIndex = kafkaHeaders.findIndex((h) => {
                        return Object.keys(h).some(key => key === k);
                    });
                    if (existingIndex >= 0) {
                        kafkaHeaders[existingIndex] = { ...kafkaHeaders[existingIndex], [k]: Buffer.from(headerValue, "utf8") };
                    }
                    else {
                        kafkaHeaders.push({ [k]: Buffer.from(headerValue, "utf8") });
                    }
                });
                api_1.propagation.inject(ctx, opaque ?? {});
                if (self._config.producerHook) {
                    self._config.producerHook(span, {
                        topic,
                        message: {
                            topic,
                            partition: partition ?? 0,
                            offset: -1,
                            size: 0,
                            value: message,
                            key: key,
                            timestamp: timestamp ?? 0,
                            opaque: opaque,
                            headers: kafkaHeaders,
                        },
                    });
                }
                return api_1.context.with(ctx, () => self.endSpan(() => original.apply(this, [
                    topic,
                    partition,
                    message,
                    key,
                    timestamp,
                    opaque,
                    kafkaHeaders,
                ]), span));
            };
        };
    }
    _patchKafkaProducer() {
        const self = this;
        return (original) => {
            return function (config) {
                const producer = original.apply(this, [config]);
                if (producer && typeof producer.send === "function") {
                    self._wrap(producer, "send", self._patchProducerSend());
                }
                return producer;
            };
        };
    }
    _patchKafkaConsumer() {
        const self = this;
        return (original) => {
            return function (config) {
                const consumer = original.apply(this, [config]);
                if (consumer && typeof consumer.run === "function") {
                    self._wrap(consumer, "run", self._patchConsumerRun());
                }
                return consumer;
            };
        };
    }
    _patchProducerSend() {
        const self = this;
        return (original) => {
            return async function (record) {
                const topic = record.topic;
                const messages = record.messages || [];
                const spans = [];
                const ctxs = [];
                const modifiedMessages = [];
                const promiseApiContext = api_1.context.active().setValue(KAFKAJS_PROMISE_API_CONTEXT_KEY, true);
                for (const message of messages) {
                    const span = self.tracer.startSpan(topic, {
                        kind: api_1.SpanKind.PRODUCER,
                        attributes: {
                            [incubating_1.ATTR_MESSAGING_SYSTEM]: incubating_1.MESSAGING_SYSTEM_VALUE_KAFKA,
                            [incubating_1.ATTR_MESSAGING_DESTINATION_NAME]: topic,
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_KEY]: message.key?.toString(),
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_TOMBSTONE]: message.key && !message.value ? true : undefined,
                            [incubating_1.ATTR_MESSAGING_DESTINATION_PARTITION_ID]: message.partition?.toString(),
                            [incubating_1.ATTR_MESSAGING_OPERATION_NAME]: "send",
                            [incubating_1.ATTR_MESSAGING_OPERATION_TYPE]: incubating_1.MESSAGING_OPERATION_TYPE_VALUE_SEND,
                        },
                    });
                    const ctx = api_1.trace.setSpan(promiseApiContext, span);
                    ctxs.push(ctx);
                    spans.push(span);
                    const carrier = {};
                    api_1.propagation.inject(ctx, carrier);
                    const headers = message.headers ? { ...message.headers } : {};
                    Object.entries(carrier).forEach(([k, v]) => {
                        if (!v)
                            return;
                        const headerValue = typeof v === "string" ? v : String(v);
                        const headerBuffer = Buffer.from(headerValue, "utf8");
                        if (headers[k] !== undefined) {
                            const existing = headers[k];
                            if (Array.isArray(existing)) {
                                headers[k] = [...existing, headerBuffer];
                            }
                            else if (existing !== undefined) {
                                headers[k] = [existing, headerBuffer];
                            }
                            else {
                                headers[k] = headerBuffer;
                            }
                        }
                        else {
                            headers[k] = headerBuffer;
                        }
                    });
                    const modifiedMessage = {
                        ...message,
                        headers: headers,
                    };
                    modifiedMessages.push(modifiedMessage);
                    if (self._config.producerHook) {
                        const convertedHeaders = Object.entries(headers).map(([k, v]) => {
                            if (Array.isArray(v)) {
                                return { [k]: v[0] };
                            }
                            else if (typeof v === "string") {
                                return { [k]: Buffer.from(v, "utf8") };
                            }
                            else {
                                return { [k]: v };
                            }
                        });
                        const messageValue = message.value
                            ? typeof message.value === "string"
                                ? Buffer.from(message.value, "utf8")
                                : message.value
                            : null;
                        const messageKey = message.key
                            ? typeof message.key === "string"
                                ? Buffer.from(message.key, "utf8")
                                : message.key
                            : null;
                        self._config.producerHook(span, {
                            topic,
                            message: {
                                topic,
                                partition: message.partition ?? 0,
                                offset: -1,
                                size: 0,
                                value: messageValue,
                                key: messageKey,
                                timestamp: message.timestamp ? (typeof message.timestamp === "string" ? parseInt(message.timestamp) : message.timestamp) : 0,
                                headers: convertedHeaders,
                            },
                        });
                    }
                }
                const modifiedRecord = {
                    topic,
                    messages: modifiedMessages,
                };
                try {
                    const result = await api_1.context.with(promiseApiContext, async () => {
                        return await original.apply(this, [modifiedRecord]);
                    });
                    spans.forEach((span) => span.end());
                    return result;
                }
                catch (error) {
                    spans.forEach((span) => {
                        if (error instanceof Error) {
                            span.recordException(error);
                        }
                        span.setStatus({
                            code: api_1.SpanStatusCode.ERROR,
                            message: error instanceof Error ? error.message : String(error),
                        });
                        span.end();
                    });
                    throw error;
                }
            };
        };
    }
    _patchConsumerRun() {
        const self = this;
        return (original) => {
            return async function (config) {
                const originalEachMessage = config?.eachMessage;
                const wrappedEachMessage = async (payload) => {
                    const params = {
                        topic: payload.topic,
                        partition: payload.partition,
                        message: payload.message,
                    };
                    const headerCarrier = {};
                    if (params.message.headers) {
                        const headers = params.message.headers;
                        Object.entries(headers).forEach(([key, value]) => {
                            if (value !== undefined && value !== null) {
                                const actualValue = Array.isArray(value) ? value[0] : value;
                                headerCarrier[key] = Buffer.isBuffer(actualValue)
                                    ? actualValue.toString("utf8")
                                    : typeof actualValue === "string"
                                        ? actualValue
                                        : String(actualValue);
                            }
                        });
                    }
                    const ctx = Object.keys(headerCarrier).length > 0
                        ? api_1.propagation.extract(api_1.ROOT_CONTEXT, headerCarrier, api_1.defaultTextMapGetter)
                        : api_1.ROOT_CONTEXT;
                    const span = self.tracer.startSpan(payload.topic, {
                        kind: api_1.SpanKind.CONSUMER,
                        attributes: {
                            [incubating_1.ATTR_MESSAGING_SYSTEM]: incubating_1.MESSAGING_SYSTEM_VALUE_KAFKA,
                            [incubating_1.ATTR_MESSAGING_DESTINATION_NAME]: payload.topic,
                            [incubating_1.ATTR_MESSAGING_OPERATION_NAME]: "received",
                            [incubating_1.ATTR_MESSAGING_OPERATION_TYPE]: incubating_1.MESSAGING_OPERATION_TYPE_VALUE_RECEIVE,
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_KEY]: payload.message.key?.toString() ?? undefined,
                            [incubating_1.ATTR_MESSAGING_KAFKA_MESSAGE_TOMBSTONE]: payload.message.key && !payload.message.value ? true : undefined,
                            [incubating_1.ATTR_MESSAGING_KAFKA_OFFSET]: payload.message.offset,
                            [incubating_1.ATTR_MESSAGING_DESTINATION_PARTITION_ID]: payload.partition.toString(),
                        },
                    }, ctx);
                    if (self._config.consumerHook) {
                        let convertedHeaders = undefined;
                        if (payload.message.headers) {
                            const headers = payload.message.headers;
                            convertedHeaders = Object.entries(headers).map(([key, value]) => {
                                if (Array.isArray(value)) {
                                    return { [key]: Buffer.isBuffer(value[0]) ? value[0] : Buffer.from(value[0], "utf8") };
                                }
                                else if (typeof value === "string") {
                                    return { [key]: Buffer.from(value, "utf8") };
                                }
                                else {
                                    return { [key]: value };
                                }
                            });
                        }
                        self._config.consumerHook(span, {
                            topic: payload.topic,
                            message: {
                                topic: payload.topic,
                                partition: payload.partition,
                                offset: parseInt(payload.message.offset),
                                size: payload.message.value ? payload.message.value.length : 0,
                                value: payload.message.value ?? null,
                                key: payload.message.key ?? null,
                                headers: convertedHeaders,
                            },
                        });
                    }
                    return api_1.context.with(api_1.trace.setSpan(ctx, span), async () => {
                        try {
                            if (originalEachMessage) {
                                await originalEachMessage(payload);
                            }
                            span.end();
                        }
                        catch (error) {
                            if (error instanceof Error) {
                                span.recordException(error);
                            }
                            span.setStatus({
                                code: api_1.SpanStatusCode.ERROR,
                                message: error instanceof Error ? error.message : String(error),
                            });
                            span.end();
                            throw error;
                        }
                    });
                };
                return original.apply(this, [{ ...config, eachMessage: wrappedEachMessage }]);
            };
        };
    }
    endSpan(traced, span) {
        try {
            const result = traced();
            if ((0, types_1.isPromise)(result)) {
                return result.then((value) => {
                    span.end();
                    return value;
                }, (err) => {
                    span.recordException(err);
                    span.setStatus({
                        code: api_1.SpanStatusCode.ERROR,
                        message: err?.message,
                    });
                    span.end();
                    throw err;
                });
            }
            else {
                span.end();
                return result;
            }
        }
        catch (error) {
            span.recordException(error);
            span.setStatus({
                code: api_1.SpanStatusCode.ERROR,
                message: error instanceof Error ? error.message : String(error),
            });
            span.end();
            throw error;
        }
    }
}
exports.ConfluentKafkaInstrumentation = ConfluentKafkaInstrumentation;
//# sourceMappingURL=instrumentation.js.map