import { randomUUID } from "node:crypto";
import { ObjectId } from "mongodb";
import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import { collections, ready } from "$lib/server/database";
import type { Generation } from "$lib/types/Generation";
import { MessageUpdateType } from "$lib/types/MessageUpdate";
import type { Workspace } from "$lib/types/Workspace";
import { enqueueAgentControl, type AgentControlTarget } from "./AgentControlMailbox";
import { AgentControlPump } from "./AgentControlPump";
import { AgentGenerationStore, AgentRelayLeaseLostError } from "./AgentGenerationStore";
import type { AgentSessionManager } from "./AgentSessionManager";
import { PiRpcCommandRejectedError } from "./PiRpcClient";

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

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

describe.sequential("AgentControlPump", () => {
	it("dispatches, receipts, and seals a command before terminalization", async () => {
		const seeded = await seedPumpRun();
		const queued = await enqueueAgentControl({
			target: seeded.target,
			id: randomUUID(),
			payload: { kind: "steer", message: "take the safe path" },
		});
		expect(queued.status).toBe("queued");
		const commands: Record<string, unknown>[] = [];
		const manager = fakeManager({
			command: async (command) => {
				commands.push(command);
				return { type: "response", success: true };
			},
		});
		const pump = new AgentControlPump({
			store: seeded.store,
			manager,
			workspace: seeded.workspace,
			pollIntervalMs: 1,
		});

		await pump.sealAndStop();

		expect(commands).toHaveLength(1);
		expect(commands[0]).toMatchObject({
			type: "steer",
			_chatUiExpectedPromptId: seeded.promptId,
		});
		const generation = await collections.generations.findOne({
			generationId: seeded.generationId,
		});
		expect(generation?.agent?.control).toBeUndefined();
		expect(generation?.agent?.controlsOpen).toBe(false);
		expect(generation?.agent?.controlReceipts).toHaveLength(1);
		expect(await seeded.store.finish("completed")).toBe(true);
	});

	it("persists always policy only after the guarded bridge acknowledgement", async () => {
		const seeded = await seedPumpRun({ approvalRequestId: "approval-1", policy: "ask" });
		await enqueueAgentControl({
			target: seeded.target,
			id: randomUUID(),
			payload: {
				kind: "approval",
				requestId: "approval-1",
				response: { confirmed: true, always: true },
			},
		});
		const order: string[] = [];
		const manager = fakeManager({
			command: async () => {
				order.push("ack");
				return { type: "response", success: true };
			},
			writeFile: async () => {
				expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe(
					"auto"
				);
				order.push("policy-file");
			},
		});
		const pump = new AgentControlPump({
			store: seeded.store,
			manager,
			workspace: seeded.workspace,
			pollIntervalMs: 1,
		});

		await pump.sealAndStop();

		expect(order).toEqual(["ack", "policy-file"]);
		expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe(
			"auto"
		);
	});

	it("commits durable auto before a runtime-file write can outlive the exact relay lease", async () => {
		const seeded = await seedPumpRun({ approvalRequestId: "approval-race", policy: "ask" });
		await enqueueAgentControl({
			target: seeded.target,
			id: randomUUID(),
			payload: {
				kind: "approval",
				requestId: "approval-race",
				response: { confirmed: true, always: true },
			},
		});
		let policyAtRuntimeWrite: Workspace["policy"] | undefined;
		let reportFailure!: (cause: Error) => void;
		const failed = new Promise<Error>((resolve) => (reportFailure = resolve));
		const pump = new AgentControlPump({
			store: seeded.store,
			manager: fakeManager({
				command: async () => ({ type: "response", success: true }),
				writeFile: async () => {
					policyAtRuntimeWrite = (
						await collections.workspaces.findOne({ _id: seeded.workspace._id })
					)?.policy;
					const expired = new Date(0);
					await Promise.all([
						collections.workspaces.updateOne(
							{ _id: seeded.workspace._id },
							{ $set: { "execution.expiresAt": expired } }
						),
						collections.generations.updateOne(
							{ generationId: seeded.generationId },
							{ $set: { "agent.relay.expiresAt": expired } }
						),
					]);
				},
			}),
			workspace: seeded.workspace,
			pollIntervalMs: 1,
			onFailure: reportFailure,
		});

		expect(await failed).toBeInstanceOf(AgentRelayLeaseLostError);
		await expect(pump.waitForStop()).rejects.toBeInstanceOf(AgentRelayLeaseLostError);
		expect(policyAtRuntimeWrite).toBe("auto");
		expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe(
			"auto"
		);
	});

	it("receipts semantic rejection but leaves ambiguous transport failure for a replacement", async () => {
		const rejected = await seedPumpRun();
		await enqueueAgentControl({
			target: rejected.target,
			id: randomUUID(),
			payload: { kind: "compact" },
		});
		const rejectedPump = new AgentControlPump({
			store: rejected.store,
			manager: fakeManager({
				command: async () => {
					throw new PiRpcCommandRejectedError({
						type: "response",
						success: false,
						error: "Agent control no longer belongs to the active prompt",
					});
				},
			}),
			workspace: rejected.workspace,
			pollIntervalMs: 1,
		});
		await rejectedPump.sealAndStop();
		const rejectedGeneration = await collections.generations.findOne({
			generationId: rejected.generationId,
		});
		expect(rejectedGeneration?.agent?.controlReceipts?.[0]?.outcome).toBe("rejected");

		const ambiguous = await seedPumpRun();
		await enqueueAgentControl({
			target: ambiguous.target,
			id: randomUUID(),
			payload: { kind: "steer", message: "may have crossed the socket" },
		});
		let reportFailure!: (cause: Error) => void;
		const failed = new Promise<Error>((resolve) => (reportFailure = resolve));
		const ambiguousPump = new AgentControlPump({
			store: ambiguous.store,
			manager: fakeManager({
				command: async () => {
					throw new Error("socket closed after write");
				},
			}),
			workspace: ambiguous.workspace,
			pollIntervalMs: 1,
			onFailure: reportFailure,
		});
		expect((await failed).message).toContain("socket closed");
		await expect(ambiguousPump.waitForStop()).rejects.toThrow("socket closed");
		const ambiguousGeneration = await collections.generations.findOne({
			generationId: ambiguous.generationId,
		});
		expect(ambiguousGeneration?.agent?.control?.state).toBe("dispatching");
		expect(ambiguousGeneration?.agent?.controlsOpen).not.toBe(false);
		expect(ambiguousGeneration?.agent?.controlReceipts).toBeUndefined();
	});

	it("never persists always policy when Pi says the native approval is no longer pending", async () => {
		const seeded = await seedPumpRun({ approvalRequestId: "approval-expired", policy: "ask" });
		await enqueueAgentControl({
			target: seeded.target,
			id: randomUUID(),
			payload: {
				kind: "approval",
				requestId: "approval-expired",
				response: { confirmed: true, always: true },
			},
		});
		const writeFile = vi.fn(async () => undefined);
		const pump = new AgentControlPump({
			store: seeded.store,
			manager: fakeManager({
				command: async (command) => {
					expect(command).toMatchObject({
						type: "extension_ui_response",
						requestId: "approval-expired",
					});
					throw new PiRpcCommandRejectedError({
						type: "response",
						success: false,
						error: "Extension UI request is no longer pending",
					});
				},
				writeFile,
			}),
			workspace: seeded.workspace,
			pollIntervalMs: 1,
		});

		await pump.sealAndStop();
		expect(writeFile).not.toHaveBeenCalled();
		expect((await collections.workspaces.findOne({ _id: seeded.workspace._id }))?.policy).toBe(
			"ask"
		);
		expect(
			(await collections.generations.findOne({ generationId: seeded.generationId }))?.agent
				?.controlReceipts?.[0]?.outcome
		).toBe("rejected");
	});
});

function fakeManager(options: {
	command: (command: Record<string, unknown> & { type: string }) => Promise<unknown>;
	writeFile?: () => Promise<void>;
}): AgentSessionManager {
	return {
		rpc: { command: options.command },
		sandbox: { writeFile: options.writeFile ?? (async () => undefined) },
	} as unknown as AgentSessionManager;
}

async function seedPumpRun(
	options: { approvalRequestId?: string; policy?: Workspace["policy"] } = {}
) {
	const now = new Date();
	const conversationId = new ObjectId();
	const workspaceId = new ObjectId();
	const generationId = randomUUID();
	const messageId = randomUUID();
	const promptId = randomUUID();
	const sessionId = `agent-pump-${randomUUID()}`;
	const sandboxId = `sandbox-${randomUUID()}`;
	const sandboxGeneration = `sandbox-generation-${randomUUID()}`;
	const workspace: Workspace = {
		_id: workspaceId,
		conversationId,
		sessionId,
		repo: { kind: "blank" },
		flavor: "cpu-basic",
		model: "fake/model",
		thinkingLevel: "medium",
		policy: options.policy ?? "auto",
		sandbox: {
			status: "ready",
			id: sandboxId,
			generation: sandboxGeneration,
			namespace: "tester",
			baseUrl: `https://${sandboxId}.example.test`,
			nonce: "nonce",
			image: "node:22",
			flavor: "cpu-basic",
			createdAt: now,
		},
		spend: { accruedUsd: 0 },
		createdAt: now,
		updatedAt: now,
	};
	await collections.conversations.insertOne({
		_id: conversationId,
		sessionId,
		mode: "agent",
		workspaceId,
		messages: [
			{
				id: messageId,
				from: "assistant",
				content: "",
				generationId,
				updates: options.approvalRequestId
					? [
							{
								type: MessageUpdateType.AgentApproval,
								requestId: options.approvalRequestId,
								method: "confirm",
								title: "Allow command?",
							},
						]
					: [],
				createdAt: now,
				updatedAt: now,
			},
		],
		createdAt: now,
		updatedAt: now,
	} as never);
	await collections.workspaces.insertOne(workspace);
	await collections.generations.insertOne({
		_id: new ObjectId(),
		generationId,
		conversationId,
		workspaceId,
		messageId,
		sessionId,
		mode: "agent",
		status: "running",
		seq: 0,
		lastHeartbeatAt: now,
		startedAt: now,
		createdAt: now,
		updatedAt: now,
		agent: {
			sourceSequence: 0,
			controlsOpen: true,
			prompt: { id: promptId, message: "do work" },
			...(options.approvalRequestId ? { pendingApprovalIds: [options.approvalRequestId] } : {}),
		},
	} satisfies Generation);
	const store = await AgentGenerationStore.acquire({
		generationId,
		workspaceId,
		sandboxId,
		sandboxGeneration,
		sandboxNonce: workspace.sandbox.nonce,
	});
	if (!store) throw new Error("Expected Agent relay lease");
	const target: AgentControlTarget = {
		generationId,
		conversationId,
		workspaceId,
		messageId,
		owner: { sessionId },
	};
	return { generationId, promptId, workspace, store, target };
}