import { AsyncLocalStorage } from "node:async_hooks";
import path from "node:path";
import { DatabaseSync } from "node:sqlite";
import { expect, it, vi } from "vitest";
import type { ContextEngine } from "../../context-engine/types.js";
import { resetCommandQueueStateForTest } from "../../process/command-queue.test-support.js";
import {
  markGatewayRestartDraining,
  resetGatewayWorkAdmission,
} from "../../process/gateway-work-admission.js";
import {
  AsyncWorkScope,
  getAsyncWorkSignal,
  trackAsyncWork,
} from "../../shared/async-work-scope.js";
import { createDeferredCore } from "../../shared/deferred.js";
import {
  resetTaskFlowRegistryForTests,
  resetTaskRegistryForTests,
} from "../../tasks/task-runtime.test-helpers.js";
import { withStateDirEnv } from "../../test-helpers/state-dir-env.js";
import {
  runContextEngineMaintenance,
  waitForDeferredTurnMaintenanceForSession,
} from "./context-engine-maintenance.js";

const unchanged = { changed: false, bytesFreed: 0, rewrittenEntries: 0 };

function engine(maintain: NonNullable<ContextEngine["maintain"]>): ContextEngine {
  return {
    info: { id: "resource-test", name: "Resource test", turnMaintenanceMode: "background" },
    ingest: async () => ({ ingested: false }),
    assemble: async ({ messages }) => ({ messages, estimatedTokens: 0 }),
    compact: async () => ({ ok: true, compacted: false }),
    maintain,
  };
}

it("keeps SIGTERM cancellation attached through actual engine disposal", async () => {
  await withResources(async (db, schedule) => {
    const context = new AsyncLocalStorage<string>();
    const entered = createDeferredCore();
    const release = createDeferredCore();
    const cancelled = vi.fn();
    const disposed = vi.fn();
    const keepTestProcessAlive = () => {};
    process.on("SIGTERM", keepTestProcessAlive);
    let lifecycleSignal: AbortSignal | undefined;
    const contextEngine = engine(async ({ abortSignal }) => {
      lifecycleSignal = abortSignal;
      return unchanged;
    });
    contextEngine.dispose = async () => {
      const signal = getAsyncWorkSignal();
      if (!signal) {
        throw new Error("expected disposer work signal");
      }
      signal.addEventListener("abort", () => cancelled(context.getStore(), signal.reason), {
        once: true,
      });
      entered.resolve();
      await release.promise;
      disposed(db.prepare("SELECT value FROM probe").get()?.value);
    };
    try {
      await context.run("disposer", () => schedule(contextEngine));
      await entered.promise;
      context.run("unrelated", () => process.emit("SIGTERM", "SIGTERM"));
      expect(lifecycleSignal?.aborted).toBe(true);
      expect(cancelled).toHaveBeenCalledExactlyOnceWith("disposer", lifecycleSignal?.reason);
      release.resolve();
      await waitForDeferredTurnMaintenanceForSession("agent:main:maintenance-resources");
      expect(disposed).toHaveBeenCalledExactlyOnceWith(42);
    } finally {
      release.resolve();
      await waitForDeferredTurnMaintenanceForSession("agent:main:maintenance-resources");
      process.off("SIGTERM", keepTestProcessAlive);
    }
  });
});

async function withResources(
  run: (
    db: DatabaseSync,
    schedule: (contextEngine: ContextEngine) => Promise<void>,
  ) => Promise<void>,
) {
  await withStateDirEnv("openclaw-maintenance-resources-", async ({ stateDir }) => {
    resetCommandQueueStateForTest();
    resetTaskRegistryForTests({ persist: false });
    resetTaskFlowRegistryForTests({ persist: false });
    const db = new DatabaseSync(path.join(stateDir, "registration.sqlite"));
    db.exec("CREATE TABLE probe(value INTEGER); INSERT INTO probe VALUES (42)");
    const pending: Promise<void>[] = [];
    const schedule = async (contextEngine: ContextEngine) => {
      await runContextEngineMaintenance({
        contextEngine,
        sessionId: "resources",
        sessionKey: "agent:main:maintenance-resources",
        sessionFile: path.join(stateDir, "session.jsonl"),
        reason: "turn",
        disposeDeferredContextEngineAfterMaintenance: true,
        onDeferredMaintenance: (promise) => {
          pending.push(promise);
        },
      });
    };
    try {
      await run(db, schedule);
    } finally {
      await Promise.allSettled(pending);
      db.close();
      resetCommandQueueStateForTest();
      resetTaskRegistryForTests({ persist: false });
      resetTaskFlowRegistryForTests({ persist: false });
    }
  });
}

it.each(["maintenance", "disposal"] as const)(
  "joins actual %s descendants before maintenance completion",
  async (phase) => {
    await withStateDirEnv("openclaw-maintenance-tail-", async ({ stateDir }) => {
      resetCommandQueueStateForTest();
      resetTaskRegistryForTests({ persist: false });
      resetTaskFlowRegistryForTests({ persist: false });
      const db = new DatabaseSync(path.join(stateDir, "registration.sqlite"));
      db.exec("CREATE TABLE probe(value INTEGER); INSERT INTO probe VALUES (42)");
      const release = createDeferredCore();
      const entered = createDeferredCore();
      let tail: Promise<unknown> | undefined;
      let completion: Promise<void> | undefined;
      let returned = false;
      const startTail = () => {
        tail = trackAsyncWork(async () => {
          entered.resolve();
          await release.promise;
          return db.prepare("SELECT value FROM probe").get()?.value;
        });
        void tail.catch(() => {});
      };
      const contextEngine = engine(async () => {
        if (phase === "maintenance") {
          startTail();
        }
        return unchanged;
      });
      const dispose = vi.fn(async () => {
        if (phase === "disposal") {
          startTail();
        }
      });
      contextEngine.dispose = dispose;
      try {
        await runContextEngineMaintenance({
          contextEngine,
          sessionId: "tail",
          sessionKey: "agent:main:maintenance-tail",
          sessionFile: path.join(stateDir, "session.jsonl"),
          reason: "turn",
          disposeDeferredContextEngineAfterMaintenance: true,
          onDeferredMaintenance: (promise) => {
            completion = promise.then(() => {
              returned = true;
              db.close();
            });
          },
        });
        await entered.promise;
        await new Promise<void>((resolve) => {
          setImmediate(resolve);
        });
        expect.soft(returned).toBe(false);
        if (phase === "maintenance") {
          expect.soft(dispose).not.toHaveBeenCalled();
        }
        release.resolve();
        await expect(tail).resolves.toBe(42);
        await completion;
        expect(dispose).toHaveBeenCalledOnce();
      } finally {
        release.resolve();
        await Promise.allSettled([...(tail ? [tail] : []), ...(completion ? [completion] : [])]);
        if (!returned) {
          db.close();
        }
        resetCommandQueueStateForTest();
        resetTaskRegistryForTests({ persist: false });
        resetTaskFlowRegistryForTests({ persist: false });
      }
    });
  },
);

it("joins disposal of a superseded queued engine before reporting the active maintenance settled", async () => {
  await withResources(async (db, schedule) => {
    const context = new AsyncLocalStorage<string>();
    const contexts: Array<string | undefined> = [];
    const releaseActive = createDeferredCore();
    const activeStarted = createDeferredCore();
    const releaseDisposal = createDeferredCore();
    const disposalStarted = createDeferredCore();
    const finalRan = createDeferredCore();
    const disposed = vi.fn();
    const active = engine(async () => {
      contexts.push(context.getStore());
      activeStarted.resolve();
      await releaseActive.promise;
      return unchanged;
    });
    const superseded = engine(async () => unchanged);
    superseded.dispose = async () => {
      contexts.push(context.getStore());
      disposalStarted.resolve();
      await releaseDisposal.promise;
      disposed(db.prepare("SELECT value FROM probe").get()?.value);
    };
    const latest = engine(async () => {
      contexts.push(context.getStore());
      finalRan.resolve();
      return unchanged;
    });
    let allSettled = false;
    let joined: Promise<void> | undefined;
    try {
      await context.run("active", () => schedule(active));
      await activeStarted.promise;
      await context.run("superseded", () => schedule(superseded));
      await context.run("latest", () => schedule(latest));
      await disposalStarted.promise;
      joined = waitForDeferredTurnMaintenanceForSession("agent:main:maintenance-resources").then(
        () => {
          allSettled = true;
        },
      );
      releaseActive.resolve();
      await finalRan.promise;
      await new Promise<void>((resolve) => {
        setImmediate(resolve);
      });
      expect.soft(allSettled).toBe(false);
      releaseDisposal.resolve();
      await joined;
      expect(disposed).toHaveBeenCalledExactlyOnceWith(42);
      expect(contexts).toEqual(["active", "superseded", "latest"]);
    } finally {
      releaseActive.resolve();
      releaseDisposal.resolve();
      await joined;
    }
  });
});

it.each(["parent completion", "gateway restart"] as const)(
  "keeps accepted maintenance independent of %s",
  async (mode) => {
    await withResources(async (db, schedule) => {
      const parent = new AsyncWorkScope();
      const context = new AsyncLocalStorage<string>();
      const entered = createDeferredCore();
      const release = createDeferredCore();
      const cancelled = vi.fn();
      let tail: Promise<unknown> | undefined;
      let maintenanceSignal: AbortSignal | undefined;
      let parentClosed = false;
      let closing: Promise<void> | undefined;
      let cancellationCheckpoint: Promise<void> | undefined;
      let checkpointPassed = false;
      try {
        await context.run("admitted", () =>
          parent.track(() =>
            schedule(
              engine(async ({ abortSignal }) => {
                maintenanceSignal = abortSignal;
                const signal = getAsyncWorkSignal();
                if (!signal) {
                  throw new Error("expected maintenance work cancellation");
                }
                signal.addEventListener(
                  "abort",
                  () => {
                    cancelled(context.getStore(), signal.reason);
                    cancellationCheckpoint = waitForDeferredTurnMaintenanceForSession(
                      "agent:main:maintenance-resources",
                    ).then(() => {
                      checkpointPassed = true;
                    });
                  },
                  { once: true },
                );
                tail = trackAsyncWork(async () => {
                  entered.resolve();
                  await release.promise;
                  return db.prepare("SELECT value FROM probe").get()?.value;
                });
                return unchanged;
              }),
            ),
          ),
        );
        await entered.promise;
        closing = context
          .run("unrelated", () => parent.drain())
          .then(() => {
            parentClosed = true;
          });
        await vi.waitFor(() => expect(parentClosed).toBe(true));
        expect.soft(maintenanceSignal?.aborted).toBe(false);
        expect.soft(cancelled).not.toHaveBeenCalled();
        if (mode === "gateway restart") {
          context.run("unrelated", () => markGatewayRestartDraining());
          expect(maintenanceSignal?.aborted).toBe(true);
          expect(cancelled).toHaveBeenCalledExactlyOnceWith("admitted", maintenanceSignal?.reason);
          await vi.waitFor(() => expect(checkpointPassed).toBe(true));
        }
        release.resolve();
        await expect(tail).resolves.toBe(42);
      } finally {
        release.resolve();
        await tail;
        await closing;
        await parent.drain();
        await cancellationCheckpoint;
        resetGatewayWorkAdmission();
      }
    });
  },
);