import { AsyncLocalStorage } from "node:async_hooks";

const replyRequestContext = new AsyncLocalStorage<string>();
const replyBuffers = new Map<string, string[]>();

export function runWithReplyRequestContext<T>(requestId: string, fn: () => Promise<T>): Promise<T> {
  return replyRequestContext.run(requestId, fn);
}

export function startReplyCapture(requestId: string) {
  replyBuffers.set(requestId, []);
}

export function clearReplyCapture(requestId: string) {
  replyBuffers.delete(requestId);
}

export function collectReplyChunk(text: string) {
  const requestId = replyRequestContext.getStore();
  if (!requestId) {
    return false;
  }

  const chunks = replyBuffers.get(requestId);
  if (!chunks) {
    return false;
  }

  chunks.push(text);
  return true;
}

export function consumeReplyCapture(requestId: string) {
  return (replyBuffers.get(requestId) ?? []).join("");
}