import type { ObjectId } from "mongodb";
import { AbortRegistry } from "$lib/server/abortRegistry";
import { authCondition } from "$lib/server/auth";
import { collections } from "$lib/server/database";
import type { Conversation } from "$lib/types/Conversation";
import type { Generation } from "$lib/types/Generation";

export interface GenerationStopMarker {
	generationId: string;
	seenContentLength?: number;
}

/** Return the durable stop point carried by a registered generation. */
export function generationStopMarker(
	generation: Pick<Generation, "generationId" | "stopRequested">
): GenerationStopMarker | null {
	if (!generation.stopRequested) return null;
	return {
		generationId: generation.generationId,
		...(generation.stopRequested.seenContentLength === undefined
			? {}
			: { seenContentLength: generation.stopRequested.seenContentLength }),
	};
}

/**
 * Resolve a stop intent without depending on the short-lived wake-up marker.
 *
 * Once a Generation exists, its stopRequested field is authoritative. The
 * aborted-generations collection is only a pre-registration/cross-pod wake-up
 * handshake and may expire while a detached Agent run is still recoverable.
 */
export async function resolveGenerationStop(options: {
	conversationId: ObjectId;
	generationId: string;
}): Promise<GenerationStopMarker | null> {
	const generation = await collections.generations.findOne(
		{
			conversationId: options.conversationId,
			generationId: options.generationId,
		},
		{ projection: { generationId: 1, stopRequested: 1 } }
	);
	if (generation) return generationStopMarker(generation);

	const marker = await collections.abortedGenerations.findOne(
		{
			conversationId: options.conversationId,
			generationId: options.generationId,
		},
		{ projection: { seenContentLength: 1 } }
	);
	if (!marker) return null;
	return {
		generationId: options.generationId,
		...(marker.seenContentLength === undefined
			? {}
			: { seenContentLength: marker.seenContentLength }),
	};
}

interface ResolveStopTargetParams {
	conversation: Conversation;
	locals: App.Locals;
	requestedGenerationId?: string;
}

/**
 * Resolve a stop request to one exact, still-running generation owned by this
 * caller. Agent legacy requests have no generation id, so their only safe
 * target is the generation named by the workspace's current live execution
 * claim. Chat legacy requests retain their historical newest-running fallback.
 */
export async function resolveStopGeneration({
	conversation,
	locals,
	requestedGenerationId,
}: ResolveStopTargetParams): Promise<string | undefined> {
	const owner = authCondition(locals);
	const base = {
		conversationId: conversation._id,
		status: "running" as const,
		...owner,
	};

	if (requestedGenerationId) {
		const generation = await collections.generations.findOne(
			{ generationId: requestedGenerationId },
			{
				projection: {
					generationId: 1,
					conversationId: 1,
					userId: 1,
					sessionId: 1,
					status: 1,
				},
			}
		);
		// The client chooses the UUID before the producer registers its writer.
		// Preserve an exact pre-registration stop so the eventual owner sees the
		// tombstone before dispatch. An existing row, however, must be the caller's
		// still-running row; terminal or foreign IDs are safe no-ops.
		if (!generation) return requestedGenerationId;
		const sameOwner = locals.user
			? generation.userId?.equals(locals.user._id) === true
			: generation.sessionId === locals.sessionId && generation.userId === undefined;
		if (
			generation.status !== "running" ||
			!generation.conversationId.equals(conversation._id) ||
			!sameOwner
		) {
			return undefined;
		}
		return generation.generationId;
	}

	if (conversation.mode === "agent") {
		if (!conversation.workspaceId) return undefined;

		const workspace = await collections.workspaces.findOne(
			{
				_id: conversation.workspaceId,
				conversationId: conversation._id,
				...owner,
				retirement: { $exists: false },
				"execution.expiresAt": { $gt: new Date() },
			},
			{ projection: { execution: 1 } }
		);
		if (!workspace?.execution) return undefined;

		const generation = await collections.generations.findOne(
			{
				...base,
				generationId: workspace.execution.generationId,
				mode: "agent",
				workspaceId: conversation.workspaceId,
			},
			{ projection: { generationId: 1 } }
		);
		return generation?.generationId;
	}

	const generation = await collections.generations.findOne(base, {
		projection: { generationId: 1 },
		sort: { startedAt: -1, _id: -1 },
	});
	return generation?.generationId;
}

interface RecordStopParams {
	conversationId: ObjectId;
	generationId: string;
	seenContentLength?: number;
}

function stopIntentUpdate(now: Date, seenContentLength?: number) {
	return {
		$set: {
			"stopRequested.requestedAt": now,
			updatedAt: now,
		},
		...(seenContentLength === undefined
			? {}
			: { $min: { "stopRequested.seenContentLength": seenContentLength } }),
	};
}

async function upsertStopMarker({
	conversationId,
	generationId,
	seenContentLength,
	now,
}: RecordStopParams & { now: Date }): Promise<void> {
	await collections.abortedGenerations.updateOne(
		{ conversationId, generationId },
		{
			$set: { updatedAt: now },
			...(seenContentLength === undefined ? {} : { $min: { seenContentLength } }),
			$setOnInsert: { createdAt: now },
		},
		{ upsert: true }
	);
}

/** Complete the registration side of the pre-registration stop handshake. */
export async function reconcileRegisteredGenerationStop(options: {
	conversationId: ObjectId;
	generationId: string;
}): Promise<boolean> {
	const marker = await collections.abortedGenerations.findOne(
		{ conversationId: options.conversationId, generationId: options.generationId },
		{ projection: { seenContentLength: 1, updatedAt: 1 } }
	);
	if (!marker) return false;
	const requestedAt = marker.updatedAt ?? new Date();
	const result = await collections.generations.updateOne(
		{
			conversationId: options.conversationId,
			generationId: options.generationId,
			status: "running",
		},
		stopIntentUpdate(requestedAt, marker.seenContentLength)
	);
	return result.matchedCount === 1;
}

/**
 * Atomically compete with terminalization on the exact Generation row, then
 * write the cross-pod wake-up marker and signal this pod's exact owner.
 *
 * A client can stop before its POST has inserted the Generation. The marker and
 * the post-marker retry form a two-sided handshake with writer registration:
 * whichever side inserts second observes the other, so a pre-registration stop
 * cannot fall into the check/insert gap.
 */
export async function recordGenerationStop({
	conversationId,
	generationId,
	seenContentLength,
}: RecordStopParams): Promise<boolean> {
	const now = new Date();
	const claimed = await collections.generations.updateOne(
		{ conversationId, generationId, status: "running" },
		stopIntentUpdate(now, seenContentLength)
	);
	if (claimed.matchedCount === 0) {
		// An existing row means terminalization (or a foreign/conflicting reuse)
		// already won. Do not leave a marker that can mislabel a terminal answer.
		const existing = await collections.generations.findOne(
			{ generationId },
			{ projection: { _id: 1 } }
		);
		if (existing) return false;
	}

	await upsertStopMarker({ conversationId, generationId, seenContentLength, now });

	if (claimed.matchedCount === 0) {
		// Close the no-row -> insert race. If registration appeared after our first
		// CAS, this retry installs the authoritative intent before returning.
		const raced = await collections.generations.updateOne(
			{ conversationId, generationId, status: "running" },
			stopIntentUpdate(now, seenContentLength)
		);
		if (raced.matchedCount === 0) {
			const terminal = await collections.generations.findOne(
				{ generationId },
				{ projection: { _id: 1 } }
			);
			if (terminal) {
				await collections.abortedGenerations.deleteOne({ conversationId, generationId });
				return false;
			}
		}
	}

	AbortRegistry.getInstance().abort(conversationId.toString(), generationId);
	return true;
}