import { AsyncLocalStorage } from "node:async_hooks";
import type { ContextEngine } from "../../context-engine/types.js";
import { formatErrorMessage } from "../../infra/errors.js";
import {
  AsyncWorkScope,
  captureAsyncWorkTracker,
  getAsyncWorkSignal,
} from "../../shared/async-work-scope.js";
import { createDeferredCore } from "../../shared/deferred.js";
import type { PreparedModelRuntimeLease } from "../prepared-model-runtime.js";
import type { ContextEngineMaintenanceResources } from "./context-engine-maintenance-work.js";
import { log } from "./logger.js";
import type { EmbeddedAgentCompactResult } from "./types.js";

export type ForegroundCompactionOwner = {
  adoptLease: (lease: PreparedModelRuntimeLease) => ContextEngineMaintenanceResources;
  resolveEngine: (factory: () => Promise<ContextEngine>) => Promise<ContextEngine>;
  captureContext: () => void;
  transferEngine: () => void;
};

/** Keeps physical cleanup separate from the compaction result and transcript fences. */
export async function runForegroundCompactionWork(
  run: (owner: ForegroundCompactionOwner) => Promise<EmbeddedAgentCompactResult>,
): Promise<EmbeddedAgentCompactResult> {
  const result = createDeferredCore<EmbeddedAgentCompactResult>();
  const trackOwner = captureAsyncWorkTracker();
  const parentSignal = getAsyncWorkSignal();
  void trackOwner(async () => {
    const work = new AsyncWorkScope();
    const factoryWork = new AsyncWorkScope();
    const cleanupWork = new AsyncWorkScope();
    let runInContext = work.run(() => AsyncLocalStorage.snapshot());
    let factoryContext = factoryWork.run(() => AsyncLocalStorage.snapshot());
    let cleanupContext = cleanupWork.run(() => AsyncLocalStorage.snapshot());
    let foregroundClosed: Promise<void> | undefined;
    const closeForegroundWork = () =>
      (foregroundClosed ??= AsyncWorkScope.runWhenAllIdle(
        () => [work],
        () => runInContext(() => work.drain()),
      ));
    let factoryClosed: Promise<void> | undefined;
    const closeFactoryWork = () => (factoryClosed ??= factoryContext(() => factoryWork.drain()));
    let lease: PreparedModelRuntimeLease | undefined;
    let engine: ContextEngine | undefined;
    let engineTransferred = false;
    const closeFromParent = () => {
      runInContext(() => work.beginClose(parentSignal?.reason));
      if (!engineTransferred) {
        factoryContext(() => factoryWork.beginClose(parentSignal?.reason));
      }
      cleanupContext(() => cleanupWork.beginClose(parentSignal?.reason));
    };
    parentSignal?.addEventListener("abort", closeFromParent, { once: true });
    if (parentSignal?.aborted) {
      closeFromParent();
    }
    try {
      result.resolve(
        await work.track(() =>
          run({
            adoptLease: (acquired) => {
              lease = acquired;
              return {
                closeFactoryWork,
                release: async () => {
                  // Fast maintenance still shares any admitted foreground preparation tails.
                  await closeForegroundWork();
                  await acquired[Symbol.asyncDispose]();
                },
              };
            },
            resolveEngine: (factory) =>
              factoryWork.track(async () => {
                factoryContext = AsyncLocalStorage.snapshot();
                engine = await factory();
                return engine;
              }),
            captureContext: () => {
              runInContext = AsyncLocalStorage.snapshot();
            },
            transferEngine: () => {
              engine = undefined;
              engineTransferred = true;
            },
          }),
        ),
      );
    } catch (error) {
      result.reject(error);
    } finally {
      try {
        await closeForegroundWork();
        if (!engineTransferred) {
          cleanupContext = factoryContext(() =>
            cleanupWork.run(() => AsyncLocalStorage.snapshot()),
          );
          // Factory services can need their disposer to start before their signal closes.
          const disposal = cleanupContext(() =>
            cleanupWork.track(async () => {
              try {
                await engine?.dispose?.();
              } catch (error) {
                log.warn("context engine dispose failed", {
                  errorMessage: formatErrorMessage(error),
                });
              }
            }),
          );
          await Promise.all([disposal, closeFactoryWork()]);
          await AsyncWorkScope.runWhenAllIdle(
            () => [cleanupWork],
            () => cleanupContext(() => cleanupWork.drain()),
          );
        }
      } finally {
        parentSignal?.removeEventListener("abort", closeFromParent);
        if (!engineTransferred) {
          await lease?.[Symbol.asyncDispose]();
        }
      }
    }
  }).catch(result.reject);
  return await result.promise;
}