import { Injectable } from "@nestjs/common";
import { context, trace, type Span, type Context } from "@opentelemetry/api";

/**
 * Service for managing OpenTelemetry span context propagation.
 * Provides utilities for accessing and manipulating the active trace context.
 *
 * @example
 * ```typescript
 * const activeSpan = spanContextService.getActiveSpan();
 * if (activeSpan) {
 *   activeSpan.setAttribute('custom.key', 'value');
 * }
 * ```
 */
@Injectable()
export class SpanContextService {
  /**
   * Get the currently active span from the context.
   *
   * @returns The active span, or undefined if no span is active
   */
  getActiveSpan(): Span | undefined {
    return trace.getActiveSpan();
  }

  /**
   * Get the current trace context.
   *
   * @returns The current context
   */
  getContext(): Context {
    return context.active();
  }

  /**
   * Execute a function within a specific context.
   *
   * @param ctx - The context to run the function in
   * @param fn - The function to execute
   * @returns The result of the function
   *
   * @example
   * ```typescript
   * const result = spanContextService.with(customContext, () => {
   *   // This code runs with customContext active
   *   return doSomething();
   * });
   * ```
   */
  with<T>(ctx: Context, fn: () => T): T {
    return context.with(ctx, fn);
  }

  /**
   * Set a span as active for the duration of a function execution.
   *
   * @param span - The span to make active
   * @param fn - The function to execute with the active span
   * @returns The result of the function
   *
   * @example
   * ```typescript
   * const result = await spanContextService.withActiveSpan(span, async () => {
   *   // This code runs with span active
   *   return await doAsyncWork();
   * });
   * ```
   */
  withActiveSpan<T>(span: Span, fn: () => T): T {
    return context.with(trace.setSpan(context.active(), span), fn);
  }
}