import ddTrace from "dd-trace";
import LLMObsEvalMetricsWriter from "dd-trace/packages/dd-trace/src/llmobs/writers/evaluations.js";
import LLMObsSpanWriter from "dd-trace/packages/dd-trace/src/llmobs/writers/spans.js";
import { getLogger } from "../../util/logger.js";
class SpanInterceptor {
    constructor(shouldDropSpans = true) {
        this.shouldDrop = shouldDropSpans;
    }
    init() {
        this.reset();
        this.restore();
        getLogger().debug("SpanInterceptor: Patching LLMObsSpanWriter.flush");
        this.patchSpanWriterFlush();
        getLogger().debug("SpanInterceptor: Patching LLMObsEvalMetricsWriter.flush");
        this.patchEvaluationWriterFlush();
    }
    patchSpanWriterFlush() {
        if (!this.originalSpanFlushMethod) {
            this.originalSpanFlushMethod = LLMObsSpanWriter.prototype.flush;
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            const self = this;
            LLMObsSpanWriter.prototype.flush = function () {
                getLogger().debug("LLMObs span writer flush intercepted");
                getLogger().debug(`Capturing ${this._buffer.length} spans from flush`);
                if (this._buffer.length > 0) {
                    self.captureSpansFromFlush(this._buffer);
                }
                if (self.shouldDrop) {
                    this._buffer = [];
                    this._bufferSize = 0;
                    getLogger().debug("Spans dropped to prevent DataDog submission");
                }
                else {
                    getLogger().debug(`Flushing ${this._buffer.length} spans to DataDog`);
                    // Call the original method with the correct 'this' context
                    self.originalSpanFlushMethod.call(this);
                }
            };
            getLogger().debug("Successfully patched LLMObsSpanWriter.flush method for span capture");
        }
    }
    shouldDrop;
    capturedSpans = [];
    capturedEvaluations = [];
    originalSpanFlushMethod;
    originalEvalFlushMethod;
    patchEvaluationWriterFlush() {
        if (!this.originalEvalFlushMethod) {
            this.originalEvalFlushMethod = LLMObsEvalMetricsWriter.prototype.flush;
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            const self = this;
            LLMObsEvalMetricsWriter.prototype.flush = function () {
                getLogger().debug("LLMObs evaluation writer flush intercepted");
                getLogger().debug(`Capturing ${this._buffer.length} evaluation metrics from flush`);
                if (this._buffer.length > 0) {
                    self.captureEvaluationsFromFlush(this._buffer);
                }
                if (self.shouldDrop) {
                    this._buffer = [];
                    this._bufferSize = 0;
                    getLogger().debug("Evaluations dropped to prevent DataDog submission");
                }
                else {
                    getLogger().debug(`Flushing ${this._buffer.length} evaluations to DataDog`);
                    // Call the original method with the correct 'this' context
                    self.originalEvalFlushMethod.call(this);
                }
            };
            getLogger().debug("Successfully patched LLMObsEvalMetricsWriter.flush method for evaluation capture");
        }
    }
    collectSpansFromFlush() {
        getLogger().debug(`Found ${this.capturedSpans.length} spans captured from flush`);
        return this.capturedSpans.map((spanEvent) => this.convertFlushSpanToIntercepted(spanEvent));
    }
    captureSpansFromFlush(spanEvents) {
        getLogger().debug(`Capturing ${spanEvents.length} span events from flush`);
        spanEvents.forEach((spanEvent) => {
            getLogger().debug(`Captured span: ${spanEvent.name} (${spanEvent.meta?.["span.kind"]})`);
            this.capturedSpans.push(spanEvent);
        });
        getLogger().debug(`Total captured spans: ${this.capturedSpans.length}`);
    }
    collectEvaluationsFromFlush() {
        getLogger().debug(`Found ${this.capturedEvaluations.length} evaluation metrics captured from flush`);
        return this.capturedEvaluations.map((metric) => this.convertEvaluationMetric(metric));
    }
    captureEvaluationsFromFlush(metricEvents) {
        getLogger().debug(`Capturing ${metricEvents.length} evaluation metrics from flush`);
        metricEvents.forEach((metricEvent) => {
            getLogger().debug(`Captured evaluation metric: ${metricEvent.label} (${metricEvent.metric_type})`);
            this.capturedEvaluations.push(metricEvent);
        });
        getLogger().debug(`Total captured evaluation metrics: ${this.capturedEvaluations.length}`);
    }
    convertFlushSpanToIntercepted(spanEvent) {
        return {
            traceId: spanEvent.trace_id || "unknown",
            spanId: spanEvent.span_id?.toString() || "unknown",
            parentId: spanEvent.parent_id &&
                spanEvent.parent_id !== "undefined" &&
                spanEvent.parent_id !== "0000000000000000"
                ? spanEvent.parent_id
                : undefined,
            operationName: spanEvent.name || "unknown",
            startTime: spanEvent.start_ns || Date.now() * 1000,
            finishTime: spanEvent.start_ns && spanEvent.duration
                ? spanEvent.start_ns + spanEvent.duration
                : undefined,
            duration: spanEvent.duration,
            tags: spanEvent.tags || {},
            logs: [],
            meta: spanEvent.meta || {},
            metrics: spanEvent.metrics || {},
            llmobs: {
                kind: spanEvent.meta?.["span.kind"] || "unknown",
                sessionId: spanEvent.session_id,
                inputData: spanEvent.meta?.input,
                outputData: spanEvent.meta?.output,
                metadata: spanEvent.meta?.metadata,
                metrics: spanEvent.metrics,
            },
        };
    }
    convertEvaluationMetric(metricEvent) {
        const valueKey = `${metricEvent.metric_type}_value`;
        return {
            traceId: metricEvent.trace_id || "unknown",
            spanId: metricEvent.span_id || "unknown",
            label: metricEvent.label || "unknown",
            metricType: metricEvent.metric_type || "unknown",
            value: metricEvent[valueKey],
            mlApp: metricEvent.ml_app,
            timestampMs: metricEvent.timestamp_ms,
            tags: Array.isArray(metricEvent.tags) ? metricEvent.tags : [],
            raw: metricEvent,
        };
    }
    getSpans() {
        ddTrace.llmobs.flush();
        return this.collectSpansFromFlush();
    }
    getEvaluations() {
        ddTrace.llmobs.flush();
        return this.collectEvaluationsFromFlush();
    }
    getTraces() {
        const traces = {};
        const spans = this.getSpans();
        spans.forEach((span) => {
            if (!traces[span.traceId]) {
                traces[span.traceId] = [];
            }
            traces[span.traceId].push(span);
        });
        return traces;
    }
    getSpanHierarchy(traceId) {
        const spans = this.getSpans().filter((s) => s.traceId === traceId);
        const spanMap = new Map();
        const children = new Map();
        spans.forEach((span) => {
            spanMap.set(span.spanId, span);
            if (span.parentId) {
                if (!children.has(span.parentId)) {
                    children.set(span.parentId, []);
                }
                children.get(span.parentId).push(span);
            }
        });
        const buildTree = (span) => ({
            ...span,
            children: (children.get(span.spanId) || []).map(buildTree),
        });
        const rootSpans = spans.filter((s) => !s.parentId);
        return rootSpans.map(buildTree);
    }
    verifySpanStructure(expected) {
        const errors = [];
        const spans = this.getSpans();
        const llmobsSpans = spans.filter((s) => s.llmobs);
        if (expected.totalSpans !== undefined &&
            spans.length !== expected.totalSpans) {
            errors.push(`Expected ${expected.totalSpans} total spans, got ${spans.length}`);
        }
        if (expected.workflowSpans !== undefined) {
            const workflowSpans = llmobsSpans.filter((s) => s.llmobs?.kind === "workflow");
            if (workflowSpans.length !== expected.workflowSpans) {
                errors.push(`Expected ${expected.workflowSpans} workflow spans, got ${workflowSpans.length}`);
            }
        }
        if (expected.agentSpans !== undefined) {
            const agentSpans = llmobsSpans.filter((s) => s.llmobs?.kind === "agent");
            if (agentSpans.length !== expected.agentSpans) {
                errors.push(`Expected ${expected.agentSpans} agent spans, got ${agentSpans.length}`);
            }
        }
        if (expected.llmSpans !== undefined) {
            const llmSpans = llmobsSpans.filter((s) => s.llmobs?.kind === "llm");
            if (llmSpans.length !== expected.llmSpans) {
                errors.push(`Expected ${expected.llmSpans} LLM spans, got ${llmSpans.length}`);
            }
        }
        if (expected.toolSpans !== undefined) {
            const toolSpans = llmobsSpans.filter((s) => s.llmobs?.kind === "tool");
            if (toolSpans.length !== expected.toolSpans) {
                errors.push(`Expected ${expected.toolSpans} tool spans, got ${toolSpans.length}`);
            }
        }
        if (expected.hasParentChildRelationship) {
            const hasParentChild = spans.some((s) => s.parentId !== undefined);
            if (!hasParentChild) {
                errors.push("Expected parent-child relationships but found none");
            }
        }
        const traces = this.getTraces();
        const details = {
            totalSpans: spans.length,
            llmobsSpans: llmobsSpans.length,
            spanKinds: llmobsSpans.reduce((acc, s) => {
                const kind = s.llmobs?.kind || "unknown";
                acc[kind] = (acc[kind] || 0) + 1;
                return acc;
            }, {}),
            traceCount: Object.keys(traces).length,
            hierarchies: Object.keys(traces).map((traceId) => this.getSpanHierarchy(traceId)),
        };
        return { success: errors.length === 0, errors, details };
    }
    printSpanTree() {
        const traces = this.getTraces();
        getLogger().debug("\n=== Intercepted Span Tree ===");
        Object.keys(traces).forEach((traceId) => {
            getLogger().debug(`\nTrace: ${traceId.substring(0, 8)}...`);
            const hierarchy = this.getSpanHierarchy(traceId);
            const printSpan = (span, indent = 0) => {
                const prefix = "  ".repeat(indent);
                const duration = span.duration
                    ? `${(span.duration / 1000).toFixed(2)}ms`
                    : "ongoing";
                const kind = span.llmobs?.kind ? ` [${span.llmobs.kind}]` : "";
                getLogger().debug(`${prefix}${span.operationName}${kind} (${duration})`);
                if (span.llmobs?.sessionId) {
                    getLogger().debug(`${prefix}  session: ${span.llmobs.sessionId}`);
                }
                if (span.children && span.children.length > 0) {
                    span.children.forEach((child) => printSpan(child, indent + 1));
                }
            };
            hierarchy.forEach((span) => printSpan(span));
        });
        getLogger().debug("=============================\n");
    }
    reset() {
        this.capturedSpans = [];
        this.capturedEvaluations = [];
    }
    restore() {
        if (this.originalSpanFlushMethod) {
            LLMObsSpanWriter.prototype.flush = this.originalSpanFlushMethod;
            this.originalSpanFlushMethod = null;
            getLogger().debug("Restored original LLMObsSpanWriter.flush method");
        }
        if (this.originalEvalFlushMethod) {
            LLMObsEvalMetricsWriter.prototype.flush = this.originalEvalFlushMethod;
            this.originalEvalFlushMethod = null;
            getLogger().debug("Restored original LLMObsEvalMetricsWriter.flush method");
        }
    }
}
export const spanInterceptor = new SpanInterceptor(process.env.DD_FORWARD == undefined);
//# sourceMappingURL=span-interceptor.js.map