import { AsyncLocalStorage } from "node:async_hooks";
import { traceIDToString } from "./encoding";
import reqtrack from "./index";

export type logCtx = Record<string, string | number | boolean>;

const logCtxStorage = new AsyncLocalStorage<logCtx>();

export function withLogCtx<R>(obj: logCtx, f: () => R): R {
  return logCtxStorage.run(obj, f);
}

export function logCtxHook(): logCtx {
  const ctx = reqtrack.current();

  if (ctx === undefined) {
    return {};
  }

  let obj: logCtx | undefined = logCtxStorage.getStore();
  if (obj === undefined) {
    obj = {};
  }

  obj["trace_id"] = traceIDToString(ctx.traceID);

  if (ctx.extRequestID !== undefined) {
    obj["x_request_id"] = ctx.extRequestID;
  }

  if (ctx.correlationID !== undefined) {
    obj["x_correlation_id"] = ctx.correlationID;
  }

  return obj;
}