import { afterEach, beforeAll, describe, expect, it } from "vitest";
import { ObjectId } from "mongodb";
import { collections, ready } from "$lib/server/database";
import {
	prepareDetachedAgentPromptForSandbox,
	reseedAgentPromptForReplacement,
	type AgentPromptRecovery,
} from "$lib/server/agent/agentPromptRecovery";
import {
	AgentGenerationIncarnationMismatchError,
	AgentGenerationStore,
} from "$lib/server/agent/AgentGenerationStore";

beforeAll(async () => {
	await ready;
});

afterEach(async () => {
	await Promise.all([
		collections.generations.deleteMany({ generationId: /^prompt-recovery-/ }),
		collections.conversations.deleteMany({ sessionId: /^prompt-recovery-session-/ }),
		collections.workspaces.deleteMany({ sessionId: /^prompt-recovery-session-/ }),
	]);
});

describe.sequential("Agent replacement prompt recovery", () => {
	it("atomically seeds a ready-session prompt when reconnect returns a replacement sandbox", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const { conversationId, workspaceId } = await seedGeneration(generationId, {
			sandboxId: "job-before-dead-bridge",
			sandboxGeneration: "generation-before",
		});

		const message = await reseedAgentPromptForReplacement({
			generationId,
			conversationId,
			workspaceId,
			prompt: { id: "command-1", message: "continue" },
			recovery: {
				sandboxId: "job-before-dead-bridge",
				sandboxGeneration: "generation-before",
				seededMessage: "prior transcript\n\ncontinue",
			},
			sandbox: { id: "job-replacement", generation: "generation-after" },
		});

		expect(message).toBe("prior transcript\n\ncontinue");
		expect((await collections.generations.findOne({ generationId }))?.agent?.prompt?.message).toBe(
			"prior transcript\n\ncontinue"
		);
	});

	it("keeps an unseeded prompt when the ready sandbox incarnation reconnects", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const recovery: AgentPromptRecovery = {
			sandboxId: "job-stable",
			sandboxGeneration: "generation-stable",
			seededMessage: "prior transcript\n\ncontinue",
		};
		const { conversationId, workspaceId } = await seedGeneration(generationId, {
			sandboxId: "job-stable",
			sandboxGeneration: "generation-stable",
			recovery,
		});
		const message = await reseedAgentPromptForReplacement({
			generationId,
			conversationId,
			workspaceId,
			prompt: { id: "command-1", message: "continue" },
			recovery,
			sandbox: { id: "job-stable", generation: "generation-stable" },
		});
		expect(message).toBe("continue");
		const lease = await AgentGenerationStore.acquire({
			generationId,
			workspaceId,
			sandboxId: "job-stable",
			sandboxGeneration: "generation-stable",
		});
		expect(lease).not.toBeNull();
		expect(
			(await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery
		).toBeUndefined();
		await lease?.release();
	});

	it("refuses to rewrite a prompt after a relay may have dispatched it", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const { conversationId, workspaceId } = await seedGeneration(generationId, {
			leased: true,
			sandboxId: "job-before",
			sandboxGeneration: "generation-before",
		});
		await expect(
			reseedAgentPromptForReplacement({
				generationId,
				conversationId,
				workspaceId,
				prompt: { id: "command-1", message: "continue" },
				recovery: {
					sandboxId: "job-before",
					sandboxGeneration: "generation-before",
					seededMessage: "prior transcript\n\ncontinue",
				},
				sandbox: { id: "job-after", generation: "generation-after" },
			})
		).rejects.toThrow("changed before replacement context was restored");
		expect((await collections.generations.findOne({ generationId }))?.agent?.prompt?.message).toBe(
			"continue"
		);
	});

	it("fails replacement relay acquisition closed until prompt and binding move together", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const { conversationId, workspaceId } = await seedGeneration(generationId, {
			sandboxId: "job-before",
			sandboxGeneration: "generation-before",
		});

		await expect(
			AgentGenerationStore.acquire({
				generationId,
				workspaceId,
				sandboxId: "job-after",
				sandboxGeneration: "generation-after",
			})
		).rejects.toBeInstanceOf(AgentGenerationIncarnationMismatchError);

		await reseedAgentPromptForReplacement({
			generationId,
			conversationId,
			workspaceId,
			prompt: { id: "command-1", message: "continue" },
			recovery: {
				sandboxId: "job-before",
				sandboxGeneration: "generation-before",
				seededMessage: "prior transcript\n\ncontinue",
			},
			sandbox: { id: "job-after", generation: "generation-after" },
		});
		await collections.workspaces.updateOne(
			{ _id: workspaceId },
			{
				$set: {
					"sandbox.id": "job-after",
					"sandbox.generation": "generation-after",
				},
			}
		);

		const lease = await AgentGenerationStore.acquire({
			generationId,
			workspaceId,
			sandboxId: "job-after",
			sandboxGeneration: "generation-after",
		});
		expect(lease).not.toBeNull();
		expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({
			sandboxId: "job-after",
			sandboxGeneration: "generation-after",
			prompt: { message: "prior transcript\n\ncontinue" },
		});
		await lease?.release();
	});

	it("keeps a detached relay out while the replacement producer is still recovering", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const recovery: AgentPromptRecovery = {
			sandboxId: "job-before",
			sandboxGeneration: "generation-before",
			seededMessage: "prior transcript\n\ncontinue",
		};
		const { workspaceId } = await seedGeneration(generationId, {
			sandboxId: recovery.sandboxId,
			sandboxGeneration: recovery.sandboxGeneration,
			recovery,
			producer: { owner: "origin", expiresAt: new Date(Date.now() + 60_000) },
		});

		expect(
			await prepareDetachedAgentPromptForSandbox({
				generationId,
				workspaceId,
				sandbox: { id: "job-after", generation: "generation-after" },
			})
		).toBe("busy");
		expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({
			sandboxId: recovery.sandboxId,
			sandboxGeneration: recovery.sandboxGeneration,
			prompt: { message: "continue", recovery },
			producer: { owner: "origin" },
		});
	});

	it("lets the exact producer reseed a replacement while its lease is live", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const recovery: AgentPromptRecovery = {
			sandboxId: "job-before",
			sandboxGeneration: "generation-before",
			seededMessage: "prior transcript\n\ncontinue",
		};
		const { conversationId, workspaceId } = await seedGeneration(generationId, {
			sandboxId: recovery.sandboxId,
			sandboxGeneration: recovery.sandboxGeneration,
			recovery,
			producer: { owner: "origin", expiresAt: new Date(Date.now() + 60_000) },
		});

		expect(
			await reseedAgentPromptForReplacement({
				generationId,
				conversationId,
				workspaceId,
				prompt: { id: "command-1", message: "continue" },
				recovery,
				sandbox: { id: "job-after", generation: "generation-after" },
				producerOwner: "origin",
			})
		).toBe(recovery.seededMessage);
		expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({
			sandboxId: "job-after",
			sandboxGeneration: "generation-after",
			prompt: { message: recovery.seededMessage },
		});
		expect(
			(await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery
		).toBeUndefined();
	});

	it("durably reseeds after producer expiry before a detached relay acquires", async () => {
		const generationId = `prompt-recovery-${new ObjectId().toString()}`;
		const recovery: AgentPromptRecovery = {
			sandboxId: "job-before",
			sandboxGeneration: "generation-before",
			seededMessage: "prior transcript\n\ncontinue",
		};
		const { workspaceId } = await seedGeneration(generationId, {
			sandboxId: recovery.sandboxId,
			sandboxGeneration: recovery.sandboxGeneration,
			recovery,
			producer: { owner: "dead-origin", expiresAt: new Date(0) },
		});

		expect(
			await prepareDetachedAgentPromptForSandbox({
				generationId,
				workspaceId,
				sandbox: { id: "job-after", generation: "generation-after" },
			})
		).toBe("ready");
		await collections.workspaces.updateOne(
			{ _id: workspaceId },
			{
				$set: {
					"sandbox.id": "job-after",
					"sandbox.generation": "generation-after",
				},
			}
		);
		const lease = await AgentGenerationStore.acquire({
			generationId,
			workspaceId,
			sandboxId: "job-after",
			sandboxGeneration: "generation-after",
		});
		expect(lease).not.toBeNull();
		expect((await collections.generations.findOne({ generationId }))?.agent).toMatchObject({
			sandboxId: "job-after",
			sandboxGeneration: "generation-after",
			prompt: { message: recovery.seededMessage },
		});
		expect(
			(await collections.generations.findOne({ generationId }))?.agent?.prompt?.recovery
		).toBeUndefined();
		expect(
			(await collections.generations.findOne({ generationId }))?.agent?.producer
		).toBeUndefined();
		await lease?.release();
	});
});

async function seedGeneration(
	generationId: string,
	options: {
		leased?: boolean;
		sandboxId?: string;
		sandboxGeneration?: string;
		recovery?: AgentPromptRecovery;
		producer?: { owner: string; expiresAt: Date };
	} = {}
): Promise<{ conversationId: ObjectId; workspaceId: ObjectId }> {
	const now = new Date();
	const workspaceId = new ObjectId();
	const conversationId = new ObjectId();
	const messageId = new ObjectId().toString();
	const sessionId = `prompt-recovery-session-${generationId}`;
	await collections.conversations.insertOne({
		_id: conversationId,
		sessionId,
		mode: "agent",
		workspaceId,
		messages: [{ id: messageId, from: "assistant", content: "", updates: [], generationId }],
		createdAt: now,
		updatedAt: now,
	} as never);
	await collections.workspaces.insertOne({
		_id: workspaceId,
		conversationId,
		sessionId,
		flavor: "cpu-basic",
		model: "test-model",
		thinkingLevel: "medium",
		policy: "safe",
		sandbox: {
			status: "ready",
			id: options.sandboxId,
			generation: options.sandboxGeneration,
			nonce: "nonce",
		},
		...(options.producer
			? {
					execution: {
						generationId,
						phase: "producer",
						owner: options.producer.owner,
						expiresAt: options.producer.expiresAt,
					},
				}
			: {}),
		createdAt: now,
		updatedAt: now,
	} as never);
	await collections.generations.insertOne({
		_id: new ObjectId(),
		generationId,
		conversationId,
		messageId,
		sessionId,
		status: "running",
		seq: 0,
		lastHeartbeatAt: now,
		startedAt: now,
		createdAt: now,
		updatedAt: now,
		mode: "agent",
		workspaceId,
		agent: {
			sourceSequence: 0,
			prompt: {
				id: "command-1",
				message: "continue",
				...(options.recovery ? { recovery: options.recovery } : {}),
			},
			...(options.sandboxId ? { sandboxId: options.sandboxId } : {}),
			...(options.sandboxGeneration ? { sandboxGeneration: options.sandboxGeneration } : {}),
			...(options.producer ? { producer: options.producer } : {}),
			...(options.leased
				? {
						relay: {
							owner: "another-server",
							fence: 1,
							expiresAt: new Date(now.getTime() + 60_000),
						},
					}
				: {}),
		},
	});
	return { conversationId, workspaceId };
}