import type { ObjectId } from "mongodb";
import { collections } from "$lib/server/database";
import { logger } from "$lib/server/logger";
import { markGenerationMessageInterrupted } from "./terminalClaim";

export interface ActiveRun {
	conversationId: ObjectId;
	messageId: string;
	mode?: "chat" | "agent";
}

// This pod's in-flight runs, so a graceful shutdown can finalize them at once
// instead of leaving each to the reaper's ~90s sweep.
const activeRuns = new Map<string, ActiveRun>();

export function registerActiveRun(generationId: string, run: ActiveRun): void {
	activeRuns.set(generationId, run);
}

export function unregisterActiveRun(generationId: string): void {
	activeRuns.delete(generationId);
}

// The message flag — not just the generation status — is what every existing reader
// treats as terminal, so both must move together. Flip the generation first, guarded on
// `status: "running"`, so only one caller (a reaper on any pod, or the owner's shutdown)
// wins the claim; mark the message only on that win.
export async function markGenerationInterrupted(
	generationId: string,
	_run: ActiveRun
): Promise<void> {
	const now = new Date();
	const generation = await collections.generations.findOneAndUpdate(
		{ generationId, status: "running" },
		{ $set: { status: "interrupted", endedAt: now, updatedAt: now } },
		{ returnDocument: "before", includeResultMetadata: false }
	);
	// Lost the claim: the run completed or errored elsewhere first. Marking the message now
	// would flag a finished answer as stopped, so leave it untouched.
	if (!generation) return;
	await collections.workspaces.updateOne(
		{ "execution.generationId": generationId },
		{ $unset: { execution: "" } }
	);
	await markGenerationMessageInterrupted(generation, now);
}

export async function finalizeActiveRunsOnExit(): Promise<void> {
	// A pi process runs in its HF Sandbox, not in this Node process.  Leaving an
	// Agent generation running is intentional: another pod will acquire its
	// durable relay lease and continue it. Chat generations still need the old
	// interrupted finalization because their upstream request dies with this pod.
	const runs = [...activeRuns.entries()].filter(([, run]) => run.mode !== "agent");
	if (runs.length === 0) return;
	logger.info({ count: runs.length }, "[generation] finalizing in-flight runs on shutdown");
	await Promise.all(
		runs.map(([generationId, run]) =>
			markGenerationInterrupted(generationId, run).catch((err) =>
				logger.error({ err, generationId }, "[generation] failed to finalize run on shutdown")
			)
		)
	);
}