import { AsyncLocalStorage } from "node:async_hooks";
import { traceParent } from "./tracecontext";

/**
 * The context for the current request.
 */
export interface Context {
  traceID: Uint8Array; // Fixed 16-byte array for the 128-bit trace ID.
  spanID: Uint8Array; // Fixed 8-byte array for the 64-bit span ID.
  parentSpanID?: Uint8Array;

  extRequestID?: string; // The external request ID, if any.
  correlationID?: string; // A string that can be used to correlate logs and metrics across services, if any.
}

/**
 * Used to track the current request context.
 */
export default class ReqTrack {
  private readonly localStorage = new AsyncLocalStorage<Context>();

  /**
   * run the given request handler with the given context, such that any calls to `current()` within the request handler
   * will return the given context.
   */
  run<R>(ctx: Context, requestHandler: () => Promise<R>): Promise<R> {
    return this.localStorage.run(ctx, requestHandler);
  }

  /**
   * Returns the current context for the current request, or undefined if there is no current request.
   */
  current(): Context | undefined {
    return this.localStorage.getStore();
  }

  /**
   * Returns a traceparent value for the current span.
   */
  currentTraceParent(): string | undefined {
    const ctx = this.current();
    if (ctx === undefined) {
      return undefined;
    }
    return traceParent(ctx.traceID, ctx.spanID);
  }
}