import { AsyncLocalStorage } from 'node:async_hooks';
import type { Sql } from '../../db.js';
import type { PermissionMode } from './types.js';

/**
 * Per-run inference context for write tools.
 *
 * Write tools need ambient state (sql, sessionId, the permission gate) that the
 * BooChat tool-phase `execute(input, projectRoot, extraRoots?)` signature can't
 * carry. This used to be a single module-level `let current` — but the inference
 * runner's `enqueue()` is fire-and-forget, so two overlapping runs (a user
 * message racing a dispatcher-polled native task; two chat tabs streaming) would
 * clobber each other's context, and `cancel()` cleared it for ALL in-flight runs.
 *
 * AsyncLocalStorage gives each run its own context: `enqueue()` starts its async
 * loop synchronously inside `runWithInferenceContext`, so the store propagates
 * through every awaited tool execution in that run — and only that run.
 */

export interface InferenceContext {
  sql: Sql;
  sessionId: string;
  taskId: string | null;
  /** Native-BooCode permission gate, set per run from the request/task mode. */
  permissionMode?: PermissionMode;
}

const storage = new AsyncLocalStorage<InferenceContext>();

/**
 * Bind `ctx` for the duration of the (possibly detached) async chain `fn` starts.
 * The inference runner kicks off its loop synchronously within this call, so all
 * downstream `await`s — including write-tool `execute` via the adapter — read the
 * same store. Concurrent runs each get their own; nothing is shared or cleared
 * out from under an in-flight run.
 */
export function runWithInferenceContext<T>(ctx: InferenceContext, fn: () => T): T {
  return storage.run(ctx, fn);
}

export function getInferenceContext(): InferenceContext {
  const ctx = storage.getStore();
  if (!ctx) {
    throw new Error(
      'Write tool called outside inference context — runWithInferenceContext() did not wrap this run',
    );
  }
  return ctx;
}