import { ROOT_CONTEXT, SpanKind, SpanStatusCode, context as otelCtx, propagation, trace } from '@opentelemetry/api';
import { SEMATTRS_MESSAGING_DESTINATION_KIND, SEMATTRS_MESSAGING_SYSTEM } from '@opentelemetry/semantic-conventions';
import { OBS_SOCKET_STATUS_CODE, OBS_TAG_API_ID, OBS_TAG_APPLICATION_ID, OBS_TAG_BRANCH, OBS_TAG_COMMIT_ID, OBS_TAG_HTTP_ROUTE, OBS_TAG_HTTP_STATUS_CODE } from '../observability/index.js';
import { sanitizeError } from '../tracing/errorSanitizer.js';
import { ISocket } from './socket.js';
function isPromise(obj) {
    return obj !== null && typeof obj === 'object' && typeof obj.then === 'function';
}
// Taken from https://github.com/gadget-inc/opentelemetry-instrumentations/blob/main/packages/opentelemetry-instrumentation-ws/src/index.ts#L21
const endSpan = (traced, span) => {
    try {
        const result = traced();
        if (isPromise(result)) {
            return Promise.resolve(result)
                .catch((err) => {
                setHttpStatusFromError(span, typeof err === 'string' ? new Error(err) : err);
                throw err;
            })
                .finally(() => span.end());
        }
        else {
            span.end();
            return result;
        }
    }
    catch (error) {
        setHttpStatusFromError(span, error);
        span.end();
        throw error;
    }
};
function setHttpStatusFromError(span, error) {
    // Simplified error handling without server-side dependencies
    span.setAttribute(OBS_TAG_HTTP_STATUS_CODE, 500);
    span.recordException(sanitizeError(error));
    span.setStatus({ code: SpanStatusCode.ERROR, message: error.message });
}
/**
 * @deprecated Use `TracedSocket` from `@superblocksteam/telemetry` instead.
 * This class will be removed in a future version.
 *
 * Migration:
 * ```typescript
 * // Before
 * import { TracedSocket } from '@superblocksteam/shared';
 *
 * // After
 * import { TracedSocket, getTelemetryInstance } from '@superblocksteam/telemetry';
 * const { policyEvaluator } = getTelemetryInstance();
 * new TracedSocket(ws, handlers, middlewares, { requestHeaders, policyEvaluator }, options);
 * ```
 */
export class TracedSocket extends ISocket {
    /**
     * The TracedSocket class extends the ISocket class to add tracing capabilities for each connection. It manages
     * spans for handler calls to provide detailed tracing information.
     *
     * Unlike Express, the ISocket handler call chain is fully recursive. This recursion makes it challenging to determine
     * when to create and end spans for specific handlers, as handlers may call other handlers before returning a result.
     *
     * An example span for recursive call
     * Main span ----------------------------------------------------------------------------------
     * Handler1     ----------------------------------------------------------------------------
     * Handler2           ----------------------------------------------------------------------
     *
     * To manage this, we keep a reference to the most recently created span and complete that span when we go one step deeper into the recursion.
     * This approach works for now since we don't support much parallelism, but it will need to be updated when we start handling requests concurrently.
     *
     * An example span for recursive call after the above fixes
     * Main span ----------------------------------------------------------------------------------
     * Handler1     -----------------------------
     * Handler2                                  ------------------------------------------------
     *
     */
    activeSpanByRequestId = new Map();
    middlewareSpanByReqId = new Map();
    lastActiveSpan;
    tracer;
    constructor(ws, requestHandlers, globalMiddlewares, tracer, options) {
        super(ws, requestHandlers, globalMiddlewares, options);
        this.tracer = tracer;
    }
    async callHandler(handler, params, ctx, client, next) {
        let currentContext = otelCtx.active();
        const middleWareSpan = this.middlewareSpanByReqId.get(ctx.requestId);
        const activeSpan = this.activeSpanByRequestId.get(ctx.requestId);
        // We need to end the middleware span since callHandler is called recursively
        // Otherwise spans will be nested which is not representing actual time taken to execute a middleware
        // Since we are ending the span here, it won't capture error state of nested handlers
        if (middleWareSpan) {
            middleWareSpan.end();
            if (activeSpan) {
                currentContext = trace.setSpan(otelCtx.active(), activeSpan);
            }
        }
        let result;
        await otelCtx.with(currentContext, async () => {
            await this.tracer.startActiveSpan(`WS HANDLER ${handler.name === '' ? ctx.method : handler.name}`, {
                attributes: {
                    [SEMATTRS_MESSAGING_SYSTEM]: 'ws',
                    [SEMATTRS_MESSAGING_DESTINATION_KIND]: 'websocket'
                },
                kind: SpanKind.SERVER
            }, async (span) => {
                this.middlewareSpanByReqId.set(ctx.requestId, span);
                result = await endSpan(() => super.callHandler(handler, params, ctx, client, next), span);
                this.middlewareSpanByReqId.delete(ctx.requestId);
            });
        });
        return result;
    }
    async handleMessage(message) {
        if (message.request) {
            const spanName = message.request.method;
            const requestId = message.request.id;
            const payload = message.request.payload;
            await otelCtx.with(propagation.extract(ROOT_CONTEXT, message.request), async () => await this.tracer.startActiveSpan(`WS SERVER ${spanName}`, {
                attributes: {
                    [SEMATTRS_MESSAGING_SYSTEM]: 'ws',
                    [SEMATTRS_MESSAGING_DESTINATION_KIND]: 'websocket',
                    [OBS_TAG_HTTP_ROUTE]: spanName,
                    [OBS_TAG_APPLICATION_ID]: payload?.['applicationId'],
                    [OBS_TAG_API_ID]: payload?.['apiId'],
                    [OBS_TAG_BRANCH]: payload?.['branch'] ?? payload?.['branchName'],
                    [OBS_TAG_COMMIT_ID]: payload?.['commitId']
                },
                kind: SpanKind.SERVER
            }, async (span) => {
                this.lastActiveSpan = span;
                this.activeSpanByRequestId.set(requestId, span);
                const result = await endSpan(() => super.handleMessage(message), span);
                this.activeSpanByRequestId.delete(requestId);
                return result;
            }));
        }
        else if (message.response) {
            this.addEvent('ws.received-response', { ['ws.response.id']: message.response.id });
            return await super.handleMessage(message);
        }
    }
    request(method, params, authorization) {
        return this.tracer.startActiveSpan(`WS CLIENT ${method}`, (span) => {
            this.addEvent('ws.send-request', { ['ws.request.method']: method, ['ws.request.id']: this.nxtRequestId });
            const result = super.request(method, params, authorization);
            span.end();
            return result;
        });
    }
    decorateToSend(message) {
        // TODO testing this out
        // Use explicit tracing headers if already provided in the request
        if (message.request?.traceparent || message.request?.tracestate) {
            // Headers already provided - no need to inject from context
            return message;
        }
        // Fallback to context.active() injection for backwards compatibility
        // TODO end
        propagation.inject(otelCtx.active(), message.request);
        return message;
    }
    respond(requestId, result) {
        this.addEvent('ws.send-response', { ['ws.response.id']: requestId });
        return super.respond(requestId, result);
    }
    respondError(requestId, error, exception) {
        this.addEvent('ws.send-error', { ['ws.response.id']: requestId });
        if (this.lastActiveSpan) {
            setHttpStatusFromError(this.lastActiveSpan, exception ?? new Error(error.message));
            this.lastActiveSpan.setAttribute(OBS_SOCKET_STATUS_CODE, error.code);
        }
        return super.respondError(requestId, error, exception);
    }
    addEvent(eventName, attributes) {
        // Since we do not have parent requestId in all the methods, we are assuming latest created span is the only active span
        // This assumption can be broken if we start sending multiple request through the same connection.
        if (this.lastActiveSpan) {
            this.lastActiveSpan.addEvent(eventName, attributes, new Date());
        }
    }
    close(reason) {
        for (const span of this.activeSpanByRequestId.values()) {
            span.end();
        }
        this.activeSpanByRequestId = new Map();
        this.lastActiveSpan = undefined;
        super.close(reason);
    }
}
//# sourceMappingURL=tracedSocket.js.map