import { randomUUID } from "node:crypto";
import { authCondition } from "$lib/server/auth";
import { collections } from "$lib/server/database";
import { getAgentJobsToken } from "$lib/server/agent/credentials";
import { logger } from "$lib/server/logger";
import { refreshWorkspaceSandboxCost } from "$lib/server/agent/sandboxCost";
import { workspaceProvisioner } from "$lib/server/agent/WorkspaceProvisioner";
import { workspaceOperationAvailableFilter } from "$lib/server/agent/WorkspaceOperation";
import type { Workspace } from "$lib/types/Workspace";

const RETIREMENT_LEASE_MS = 15 * 60 * 1_000;

export class WorkspaceRetirementConflict extends Error {
	constructor() {
		super("Agent workspace changed while it was being retired");
		this.name = "WorkspaceRetirementConflict";
	}
}

export interface RetiredAgentWorkspace {
	workspace: Workspace;
	leaseId: string;
}

/**
 * Cancels Jobs before their owning records are deleted. This is deliberately
 * sequential so a bulk conversation delete cannot burst Jobs cancellation calls.
 */
export async function retireAgentWorkspaces(
	workspaces: Workspace[],
	locals: App.Locals
): Promise<RetiredAgentWorkspace[]> {
	const owner = authCondition(locals);
	const claimed: RetiredAgentWorkspace[] = [];
	// Claim every target before stopping the first Job. A multi-workspace delete
	// must either own the whole set or leave the records intact for a retry.
	try {
		for (const candidate of workspaces) {
			const leaseId = randomUUID();
			const claimNow = new Date();
			const retirementStaleBefore = new Date(claimNow.getTime() - RETIREMENT_LEASE_MS);
			const canReclaimLease =
				candidate.retirement &&
				(candidate.sandbox.status === "error" ||
					candidate.retirement.claimedAt < retirementStaleBefore);
			// A provisioner owns the interval between its database claim and publishing
			// the newly-created Job identity. Retiring that fresh claim could mark the
			// workspace dead (and later delete it) before the provisioner has a Job id
			// to stop. Let the in-flight provision finish or fail; its exact writes are
			// then fenced by the retirement claim on the next attempt.
			const provisioningClaimedAt = candidate.sandbox.claimedAt ?? candidate.updatedAt;
			if (
				candidate.sandbox.status === "provisioning" &&
				provisioningClaimedAt >= retirementStaleBefore &&
				!canReclaimLease
			) {
				throw new WorkspaceRetirementConflict();
			}
			const reclaimFailedLease =
				canReclaimLease && candidate.retirement
					? {
							"retirement.id": candidate.retirement.id,
							"retirement.claimedAt": candidate.retirement.claimedAt,
						}
					: { retirement: { $exists: false } };
			const workspace = await collections.workspaces.findOneAndUpdate(
				{
					_id: candidate._id,
					...owner,
					updatedAt: candidate.updatedAt,
					"sandbox.status": candidate.sandbox.status,
					...(candidate.sandbox.id ? { "sandbox.id": candidate.sandbox.id } : {}),
					...(candidate.sandbox.nonce ? { "sandbox.nonce": candidate.sandbox.nonce } : {}),
					...(candidate.sandbox.generation
						? { "sandbox.generation": candidate.sandbox.generation }
						: {}),
					$and: [
						{
							$or: [
								{ execution: { $exists: false } },
								{ "execution.expiresAt": { $lte: claimNow } },
							],
						},
						workspaceOperationAvailableFilter(claimNow),
					],
					...reclaimFailedLease,
				},
				{
					$set: {
						retirement: { id: leaseId, claimedAt: claimNow },
					},
					$unset: { execution: "", operation: "" },
				},
				{ returnDocument: "after", includeResultMetadata: false }
			);
			if (!workspace) throw new WorkspaceRetirementConflict();
			claimed.push({ workspace, leaseId });
		}
	} catch (cause) {
		await releaseWorkspaceRetirements(claimed, locals);
		throw cause;
	}

	for (let index = 0; index < claimed.length; index += 1) {
		const retirement = claimed[index];
		let workspace = retirement.workspace;
		try {
			try {
				workspace = await refreshWorkspaceSandboxCost(workspace, getAgentJobsToken(locals));
			} catch (cause) {
				logger.warn(
					{ cause, workspaceId: workspace._id.toString() },
					"Unable to finalize sandbox cost before deletion"
				);
			}
			const stopped = await workspaceProvisioner.kill({
				workspace,
				requestHfToken: locals.token,
			});
			claimed[index] = { workspace: stopped, leaseId: retirement.leaseId };
		} catch (cause) {
			// Keep the failed target leased: its Job could not be confirmed stopped.
			// Mark only our exact lease stale so the next explicit teardown can retry
			// immediately, including failures that happened before kill() persisted an
			// error state. A process crash still observes the normal lease timeout.
			await collections.workspaces.updateOne(
				{ _id: workspace._id, ...owner, "retirement.id": retirement.leaseId },
				{ $set: { "retirement.claimedAt": new Date(0) } }
			);
			await releaseWorkspaceRetirements(
				claimed.filter((_, claimedIndex) => claimedIndex !== index),
				locals
			);
			throw cause;
		}
	}
	return claimed;
}

/** Release leases only when the caller still owns them. */
export async function releaseWorkspaceRetirements(
	retired: RetiredAgentWorkspace[],
	locals: App.Locals
): Promise<void> {
	const owner = authCondition(locals);
	for (const { workspace, leaseId } of retired) {
		await collections.workspaces.updateOne(
			{ _id: workspace._id, ...owner, "retirement.id": leaseId },
			{ $unset: { retirement: "" } }
		);
	}
}

/** Exact filter for the stopped incarnation protected by a retirement lease. */
export function retiredWorkspaceFilter(
	retired: RetiredAgentWorkspace,
	locals: App.Locals
): Record<string, unknown> {
	const { workspace, leaseId } = retired;
	return {
		_id: workspace._id,
		...authCondition(locals),
		"retirement.id": leaseId,
		updatedAt: workspace.updatedAt,
		"sandbox.status": workspace.sandbox.status,
		...(workspace.sandbox.id ? { "sandbox.id": workspace.sandbox.id } : {}),
		...(workspace.sandbox.nonce ? { "sandbox.nonce": workspace.sandbox.nonce } : {}),
		...(workspace.sandbox.generation ? { "sandbox.generation": workspace.sandbox.generation } : {}),
	};
}