// (c) Copyright 2023 Supertenant Ltd. - all rights reserved.
// See LICENSE file in project root for license terms.
"use strict";const shimmer=require("shimmer"),requireHook=require("../../../util/requireHook"),tracingUtil=require("../../tracingUtil"),constants=require("../../constants"),cls=require("../../cls");let isActive=!1,clientHasBeenInstrumented=!1;exports.init=function(){requireHook.onModuleLoad("nats",instrumentNats)};let natsModule,connectionObject;function instrumentNats(_natsModule){natsModule=_natsModule,shimmer.wrap(natsModule,"connect",shimConnect)}function shimConnect(originalFunction){return function(){var client=originalFunction.apply(this,arguments);return clientHasBeenInstrumented||(client&&client.then&&client.catch?client.then(nc=>{let natsUrl="nats://";return(connectionObject=nc).protocol&&nc.protocol.server&&nc.protocol.server.listen&&(natsUrl="nats://"+nc.protocol.server.listen),shimmer.wrap(nc.constructor.prototype,"publish",shimPublish.bind(null,natsUrl,!0)),shimmer.wrap(nc.constructor.prototype,"request",shimRequest.bind(null,natsUrl)),shimmer.wrap(nc.constructor.prototype,"subscribe",shimSubscribe.bind(null,natsUrl,!0)),clientHasBeenInstrumented=!0,nc}):(shimmer.wrap(client.constructor.prototype,"publish",shimPublish.bind(null,client.options.url,!1)),shimmer.wrap(client.constructor.prototype,"request",shimRequest.bind(null,client.options.url)),shimmer.wrap(client.constructor.prototype,"subscribe",shimSubscribe.bind(null,client.options.url,!0)),clientHasBeenInstrumented=!0)),client}}function shimPublish(natsUrl,isLatest,originalFunction){return function(){var originalArgs=new Array(arguments.length);for(let i=0;i<arguments.length;i++)originalArgs[i]=arguments[i];return instrumentedPublish(this,originalFunction,originalArgs,natsUrl,isLatest)}}function addSuppressionHeaders(args,isLatest){if(isLatest&&natsModule&&natsModule.headers&&connectionObject&&connectionObject.info&&connectionObject.info.headers){var isLatest=args[2],Headers=natsModule.headers;if(isLatest&&isLatest.headers&&"object"==typeof isLatest.headers)isLatest.headers.append(constants.traceLevelHeaderName,"0");else if(isLatest){const h=Headers();h.append(constants.traceLevelHeaderName,"0"),isLatest.headers=h}else{isLatest=Headers();isLatest.append(constants.traceLevelHeaderName,"0"),args[2]={headers:isLatest}}}}function addTraceCorrelationHeaders(args,isLatest,span){if(isLatest&&natsModule&&natsModule.headers&&connectionObject&&connectionObject.info&&connectionObject.info.headers){var isLatest=args[2],Headers=natsModule.headers;if(isLatest&&isLatest.headers&&"object"==typeof isLatest.headers)isLatest.headers.append(constants.traceIdHeaderName,span.t),isLatest.headers.append(constants.spanIdHeaderName,span.s);else if(isLatest){const h=Headers();h.append(constants.traceIdHeaderName,span.t),h.append(constants.spanIdHeaderName,span.s),isLatest.headers=h}else{isLatest=Headers();isLatest.append(constants.traceIdHeaderName,span.t),isLatest.append(constants.spanIdHeaderName,span.s),args[2]={headers:isLatest}}}}function instrumentedPublish(ctx,originalPublish,originalArgs,natsUrl,isLatest){var skipTracingResult=cls.skipExitTracing({isActive:!0,extendedResponse:!0});if(skipTracingResult.skip)return skipTracingResult.suppressed&&addSuppressionHeaders(originalArgs,isLatest),originalPublish.apply(ctx,originalArgs);const subject="function"!=typeof originalArgs[0]?originalArgs[0]:"";let callbackIndex=-1;for(let i=3;0<=i;i--)"function"==typeof originalArgs[i]&&(callbackIndex=i);const originalCallback=0<=callbackIndex?originalArgs[callbackIndex]:null;return cls.ns.runAndReturn(()=>{const span=cls.startSpan("nats",constants.EXIT);span.ts=Date.now(),span.stack=tracingUtil.getStackTrace(instrumentedPublish),span.data.nats={sort:"publish",address:natsUrl,subject:subject},addTraceCorrelationHeaders(originalArgs,isLatest,span),originalCallback&&(originalArgs[callbackIndex]=cls.ns.bind(function(err){addErrorToSpan(err,span),finishSpan(span,"nats"),originalCallback.apply(this,arguments)}));try{var result=originalPublish.apply(ctx,originalArgs);return result&&"function"==typeof result.catch&&result.catch(err=>{addErrorToSpan(err,span)}),setImmediate(()=>{finishSpan(span,"nats")}),result}catch(e){throw addErrorToSpan(e,span),finishSpan(span,"nats"),e}})}function shimRequest(natsUrl,originalFunction){return function(){if(isActive&&cls.isTracing()){var opts;if(this.protocol&&this.protocol.transport)return(opts=arguments[2])&&"noMux"in opts&&!opts.noMux?originalFunction.apply(this,arguments):instrumentedPublish(this,originalFunction,arguments,natsUrl);for(let i=3;0<=i;i--)if("function"==typeof arguments[i]){arguments[i]=cls.ns.bind(arguments[i]);break}}return originalFunction.apply(this,arguments)}}function shimSubscribe(natsUrl,isLatest,originalFunction){return function(){var originalSubscribeArgs=new Array(arguments.length);for(let i=0;i<arguments.length;i++)originalSubscribeArgs[i]=arguments[i];return instrumentedSubscribe(this,originalFunction,originalSubscribeArgs,natsUrl,isLatest)}}function instrumentedSubscribe(ctx,originalSubscribe,originalSubscribeArgs,natsUrl,isLatest){const subject=originalSubscribeArgs[0];let callbackIndex=-1,isCallbackAttr=!1;for(let i=2;1<=i;i--)"function"==typeof originalSubscribeArgs[i]?callbackIndex=i:originalSubscribeArgs[i]&&originalSubscribeArgs[i].callback&&"function"==typeof originalSubscribeArgs[i].callback&&(callbackIndex=i,isCallbackAttr=!0);if(-1<callbackIndex){let originalCallback;return isCallbackAttr?(originalCallback=originalSubscribeArgs[callbackIndex].callback,originalSubscribeArgs[callbackIndex].callback=function(err,msg){return instrumentedSubscribeCallback(natsUrl,subject,originalCallback,null,isLatest).bind(this)(err,msg)}):(originalCallback=originalSubscribeArgs[callbackIndex],originalSubscribeArgs[callbackIndex]=instrumentedSubscribeCallback(natsUrl,subject,originalCallback)),originalSubscribe.apply(ctx,originalSubscribeArgs)}return cls.ns.runAndReturn(currentCtx=>{const sub=originalSubscribe.apply(ctx,originalSubscribeArgs);var instanaIterator=async function*(){cls.ns.enter(currentCtx);for await(const msg of sub)await new Promise(resolve=>{instrumentedSubscribeCallback(natsUrl,subject,resolve,currentCtx,isLatest)(null,msg)}),yield msg}();return Object.assign(instanaIterator,sub),instanaIterator})}function instrumentedSubscribeCallback(natsUrl,subject,originalSubscribeCallback,instanaCtx,isLatest){return function(err,msg){let suppressed=!1,traceId,parentSpanId;const originalCallbackThis=this,originalCallbackArgs=new Array(arguments.length);for(let i=0;i<arguments.length;i++)originalCallbackArgs[i]=arguments[i];if(isLatest&&msg&&msg.headers)for(var[key,value]of msg.headers)key===constants.traceLevelHeaderName&&"0"===value[0]?suppressed=!0:key===constants.traceIdHeaderName?traceId=value[0]:key===constants.spanIdHeaderName&&(parentSpanId=value[0]);return subject="string"==typeof originalCallbackArgs[2]?originalCallbackArgs[2]:subject,cls.ns.runAndReturn(()=>{if(suppressed)return cls.setTracingLevel("0"),originalSubscribeCallback.apply(originalCallbackThis,originalCallbackArgs);if(!isActive)return originalSubscribeCallback.apply(originalCallbackThis,originalCallbackArgs);{const span=cls.startSpan("nats",constants.ENTRY,traceId,parentSpanId);span.ts=Date.now(),span.stack=tracingUtil.getStackTrace(instrumentedSubscribeCallback),span.data.nats={sort:"consume",address:natsUrl,subject:subject};try{return originalSubscribeCallback.apply(originalCallbackThis,originalCallbackArgs)}catch(e){throw addErrorToSpan(e,span),e}finally{setImmediate(()=>{finishSpan(span,"nats")})}}},instanaCtx)}}function addErrorToSpan(err,span){if(err){span.ec=1;let errMsg=null;err.message?errMsg=err.message:"string"==typeof err&&(errMsg=err),errMsg&&span.data.nats.error?span.data.nats.error+=", "+errMsg:errMsg&&(span.data.nats.error=errMsg)}}function finishSpan(span,expectedName){span.transmitted||span.n!==expectedName||(span.d=Date.now()-span.ts,span.transmit())}exports.activate=function(){isActive=!0},exports.deactivate=function(){isActive=!1};