import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import {
  type ConversationId,
  defineDoc,
  GenerationTask,
  type Harness,
  InboxDoc,
  LiveDoc,
} from "@earendil-works/pi-durable";

/** The last time the conversation handed control back to the user. */
export const SettledDoc = defineDoc<{ updated?: number }>({
  kind: "berry.settled",
  version: 1,
  scope: "conversation",
  history: "rewindable",
  fork: "asOf",
  initial: () => ({}),
  checkpointWhen: () => true,
});

const pending = new WeakMap<Harness, Promise<void>>();

/** Record when a run yields control instead of handing it to another generation. */
export function recordSettled(harness: Harness): () => Promise<void> {
  const queued = new Map<ConversationId, number>();
  const unsubscribe = harness.subscribeCommits((publication) => {
    for (const change of publication.changes) {
      if (
        change.type === "document" && change.record.kind === InboxDoc.definition.kind &&
        change.conversationId !== undefined && Array.isArray(change.value?.items)
      ) queued.set(change.conversationId, change.value.items.length);
    }
    for (const change of publication.changes) {
      if (
        change.type !== "task" || change.value.kind !== GenerationTask.definition.name ||
        change.value.state.status !== "terminal"
      ) continue;
      const task = change.value;
      const live = publication.changes.find((change) =>
        change.type === "document" && change.record.kind === LiveDoc.definition.kind &&
        change.conversationId === task.conversationId
      );
      if (live?.type !== "document" || live.value === null || live.value.run !== undefined) {
        continue;
      }
      const count = queued.get(task.conversationId);
      const next = (pending.get(harness) ?? Promise.resolve()).then(async () => {
        const remaining = count ??
          (await harness.snapshot(InboxDoc, task.conversationId, context))?.items.length ?? 0;
        if (remaining > 0) return;
        await harness.commit(async (tx) => {
          (await tx.doc(SettledDoc, task.conversationId)).updated = task.endedAt ?? Date.now();
        }, context);
      }).catch((error) => console.error("[berry] recording settled turn:", error));
      pending.set(harness, next);
    }
  });
  return async () => {
    unsubscribe();
    await settledThrough(harness);
  };
}

/** Wait for the timestamps queued by committed run endings. */
export async function settledThrough(harness: Harness): Promise<void> {
  let next;
  do {
    next = pending.get(harness);
    await next;
  } while (next !== pending.get(harness));
}