import { randomUUID } from "node:crypto";
import type { Filter, ObjectId } from "mongodb";
import { collections } from "$lib/server/database";
import type { Workspace, WorkspaceOperationLease, WorkspaceSandbox } from "$lib/types/Workspace";

export const WORKSPACE_OPERATION_LEASE_MS = 20 * 60 * 1_000;

export class WorkspaceOperationConflict extends Error {
	constructor(message = "Agent workspace changed while the operation was starting") {
		super(message);
		this.name = "WorkspaceOperationConflict";
	}
}

export interface ClaimedWorkspaceOperation {
	workspace: Workspace;
	lease: WorkspaceOperationLease;
}

/**
 * Admission predicate shared by every workspace producer/relay/RPC path.
 * Callers which mutate the workspace must pair it with `$unset: { operation:
 * "" }`; the predicate then makes reclaiming an orphaned lease atomic while a
 * lease younger than the timeout remains a hard fence.
 */
export function workspaceOperationAvailableFilter(now: Date = new Date()): Filter<Workspace> {
	const staleBefore = new Date(now.getTime() - WORKSPACE_OPERATION_LEASE_MS);
	return {
		$or: [{ operation: { $exists: false } }, { "operation.claimedAt": { $lte: staleBefore } }],
	};
}

/**
 * Claims the exact workspace revision only while its conversation still owns it
 * in Agent mode. Retirement uses the reciprocal `operation` guard, so either
 * this claim or a mode change wins; neither can cross the other's async work.
 */
export async function claimWorkspaceOperation(options: {
	conversationId: ObjectId;
	workspace: Workspace;
	owner: Record<string, unknown>;
	kind: WorkspaceOperationLease["kind"];
	action: string;
	now?: Date;
}): Promise<ClaimedWorkspaceOperation> {
	const now = options.now ?? new Date();
	const lease: WorkspaceOperationLease = {
		id: randomUUID(),
		kind: options.kind,
		action: options.action.slice(0, 64),
		claimedAt: now,
		workspaceUpdatedAt: options.workspace.updatedAt,
	};
	const workspace = await collections.workspaces.findOneAndUpdate(
		{
			_id: options.workspace._id,
			conversationId: options.conversationId,
			...options.owner,
			updatedAt: options.workspace.updatedAt,
			...sandboxIdentityFilter(options.workspace.sandbox),
			retirement: { $exists: false },
			$and: [
				workspaceOperationAvailableFilter(now),
				{
					$or: [{ execution: { $exists: false } }, { "execution.expiresAt": { $lte: now } }],
				},
			],
		},
		{
			$set: { operation: lease },
			// The predicate admits only an absent/expired execution. Removing an
			// expired claim fences a delayed producer/relay from renewing into this
			// workspace mutation after the lease has been granted.
			$unset: { execution: "" },
		},
		{ returnDocument: "after", includeResultMetadata: false }
	);
	if (!workspace) throw new WorkspaceOperationConflict();

	const stillOwned = await collections.conversations.findOne({
		_id: options.conversationId,
		...options.owner,
		mode: "agent",
		workspaceId: workspace._id,
	});
	if (!stillOwned) {
		await releaseWorkspaceOperation(workspace._id, options.owner, lease.id);
		throw new WorkspaceOperationConflict(
			"The conversation changed before the workspace operation could start"
		);
	}
	return { workspace, lease };
}

/** Revalidates both collections and the exact workspace revision after an await. */
export async function requireWorkspaceOperation(options: {
	conversationId: ObjectId;
	workspace: Workspace;
	owner: Record<string, unknown>;
	leaseId: string;
}): Promise<Workspace> {
	const conversation = await collections.conversations.findOne({
		_id: options.conversationId,
		...options.owner,
		mode: "agent",
		workspaceId: options.workspace._id,
	});
	if (!conversation) {
		throw new WorkspaceOperationConflict("The conversation changed during this operation");
	}
	const workspace = await collections.workspaces.findOne({
		_id: options.workspace._id,
		conversationId: options.conversationId,
		...options.owner,
		"operation.id": options.leaseId,
		retirement: { $exists: false },
		updatedAt: options.workspace.updatedAt,
		...sandboxIdentityFilter(options.workspace.sandbox),
	});
	if (!workspace)
		throw new WorkspaceOperationConflict("The agent sandbox changed during this operation");
	return workspace;
}

/**
 * Reload an operation after a provisioner deliberately changed the sandbox
 * state before reporting an error. The reciprocal lease still proves no
 * supported mode/reset path crossed the operation.
 */
export async function reloadWorkspaceOperation(options: {
	conversationId: ObjectId;
	workspaceId: ObjectId;
	owner: Record<string, unknown>;
	leaseId: string;
}): Promise<Workspace> {
	const conversation = await collections.conversations.findOne({
		_id: options.conversationId,
		...options.owner,
		mode: "agent",
		workspaceId: options.workspaceId,
	});
	if (!conversation) {
		throw new WorkspaceOperationConflict("The conversation changed during this operation");
	}
	const workspace = await collections.workspaces.findOne({
		_id: options.workspaceId,
		conversationId: options.conversationId,
		...options.owner,
		"operation.id": options.leaseId,
		retirement: { $exists: false },
	});
	if (!workspace) throw new WorkspaceOperationConflict("The agent workspace changed");
	return workspace;
}

export async function releaseWorkspaceOperation(
	workspaceId: ObjectId,
	owner: Record<string, unknown>,
	leaseId: string
): Promise<void> {
	await collections.workspaces.updateOne(
		{ _id: workspaceId, ...owner, "operation.id": leaseId },
		{ $unset: { operation: "" } }
	);
}

export function sandboxIdentityFilter(sandbox: WorkspaceSandbox): Record<string, unknown> {
	return {
		"sandbox.status": sandbox.status,
		...(sandbox.id ? { "sandbox.id": sandbox.id } : { "sandbox.id": { $exists: false } }),
		...(sandbox.nonce
			? { "sandbox.nonce": sandbox.nonce }
			: { "sandbox.nonce": { $exists: false } }),
		...(sandbox.generation
			? { "sandbox.generation": sandbox.generation }
			: { "sandbox.generation": { $exists: false } }),
	};
}