import { error } from "@sveltejs/kit";
import { ObjectId } from "mongodb";
import { authCondition } from "$lib/server/auth";
import { collections } from "$lib/server/database";
import { connectAgentRpcRuntime } from "./AgentRpcRuntime";
import { sandboxIdentityFilter, workspaceOperationAvailableFilter } from "./WorkspaceOperation";
import type { Workspace } from "$lib/types/Workspace";

/** Owner-scoped lookup followed by lifecycle-aware runtime connection. */
export async function connectOwnedAgentRpcSession(id: string, locals: App.Locals) {
	if (!ObjectId.isValid(id)) error(404, "Conversation not found");
	const conversationId = new ObjectId(id);
	const owner = authCondition(locals);
	const conversation = await collections.conversations.findOne({ _id: conversationId, ...owner });
	if (!conversation) error(404, "Conversation not found");
	if (conversation.mode !== "agent" || !conversation.workspaceId) {
		error(409, "This conversation has no agent workspace");
	}

	const workspaceFilter = {
		_id: conversation.workspaceId,
		conversationId,
		...owner,
	};
	const workspace = await collections.workspaces.findOne(workspaceFilter);
	if (!workspace) error(404, "Agent workspace not found");

	const validateRuntimeWorkspace = async (runtimeWorkspace: Workspace) => {
		if (
			!runtimeWorkspace._id.equals(workspace._id) ||
			!runtimeWorkspace.conversationId.equals(conversationId)
		) {
			return false;
		}
		const stillOwned = await collections.conversations.findOne({
			_id: conversationId,
			...owner,
			mode: "agent",
			workspaceId: workspace._id,
		});
		if (!stillOwned) return false;
		const sameIncarnation = await collections.workspaces.findOneAndUpdate(
			{
				...workspaceFilter,
				updatedAt: runtimeWorkspace.updatedAt,
				retirement: { $exists: false },
				...workspaceOperationAvailableFilter(),
				...sandboxIdentityFilter(runtimeWorkspace.sandbox),
			},
			{ $unset: { operation: "" } },
			{ returnDocument: "after", includeResultMetadata: false }
		);
		return Boolean(sameIncarnation);
	};
	let connectedWorkspace: Workspace | undefined;
	const manager = await connectAgentRpcRuntime({
		workspace,
		locals,
		reloadWorkspace: () => collections.workspaces.findOne(workspaceFilter),
		validateRuntimeWorkspace: async (runtimeWorkspace) => {
			const valid = await validateRuntimeWorkspace(runtimeWorkspace);
			if (valid) connectedWorkspace = runtimeWorkspace;
			return valid;
		},
	});
	if (!connectedWorkspace) error(409, "The agent workspace changed while reconnecting");
	const exactConnectedWorkspace: Workspace = connectedWorkspace;
	return {
		manager,
		revalidate: () => validateRuntimeWorkspace(exactConnectedWorkspace),
	};
}