import { AgentSessionManager } from "./AgentSessionManager";
import { getAgentJobsToken } from "./credentials";
import {
	workspaceProvisioner,
	WorkspaceProvisioningError,
	type WorkspaceRuntimeConnection,
} from "./WorkspaceProvisioner";
import type { Workspace } from "$lib/types/Workspace";

export class AgentRpcRuntimeError extends Error {
	constructor(
		message: string,
		readonly status: number
	) {
		super(message);
		this.name = "AgentRpcRuntimeError";
	}
}

export interface ConnectAgentRpcRuntimeOptions {
	workspace: Workspace;
	locals: App.Locals;
	reloadWorkspace: () => Promise<Workspace | null>;
	/** Revalidate the owning conversation and exact sandbox after reconnect awaits. */
	validateRuntimeWorkspace: (workspace: Workspace) => Promise<boolean>;
}

/**
 * Reconnect an already-running runtime before opening pi RPC.
 *
 * This helper is used by read-only context/export requests. Those requests must
 * never create a billable sandbox as a side effect; only explicit run/resume
 * actions are allowed to provision one.
 */
export async function connectAgentRpcRuntime({
	workspace,
	locals,
	reloadWorkspace,
	validateRuntimeWorkspace,
}: ConnectAgentRpcRuntimeOptions): Promise<AgentSessionManager> {
	let runtime: WorkspaceRuntimeConnection;

	try {
		if (workspace.sandbox.status === "ready" || workspace.sandbox.status === "idle") {
			try {
				runtime = await workspaceProvisioner.reconnect({
					workspace,
					requestHfToken: locals.token,
					observational: true,
				});
			} catch (cause) {
				const refreshed = await reloadWorkspace();
				if (!refreshed) throw new AgentRpcRuntimeError("Agent workspace not found", 404);
				if (refreshed.sandbox.status === "provisioning") {
					throw new AgentRpcRuntimeError("The agent workspace is still provisioning", 409);
				}
				throw cause;
			}
		} else if (
			workspace.sandbox.status === "unprovisioned" ||
			workspace.sandbox.status === "dead" ||
			workspace.sandbox.status === "error"
		) {
			throw new AgentRpcRuntimeError(
				"The agent workspace is offline; resume it before using this action",
				409
			);
		} else {
			throw new AgentRpcRuntimeError("The agent workspace is still provisioning", 409);
		}
	} catch (cause) {
		if (cause instanceof AgentRpcRuntimeError) throw cause;
		if (cause instanceof WorkspaceProvisioningError && cause.code === "invalid_configuration") {
			throw new AgentRpcRuntimeError("Agent runtime is not configured", 503);
		}
		throw new AgentRpcRuntimeError("Unable to connect to the agent workspace", 409);
	}
	if (!(await validateRuntimeWorkspace(runtime.workspace))) {
		throw new AgentRpcRuntimeError("The agent workspace changed while reconnecting", 409);
	}

	let jobsToken: string;
	try {
		jobsToken = getAgentJobsToken(locals);
	} catch {
		throw new AgentRpcRuntimeError("Agent runtime is not configured", 503);
	}
	try {
		return AgentSessionManager.connect({
			workspace: runtime.workspace,
			hfToken: jobsToken,
			bridgeUrl: runtime.bridgeUrl,
		});
	} catch {
		throw new AgentRpcRuntimeError("Unable to connect to the coding agent", 409);
	}
}