/**
 * Antigravity stream provider.
 *
 * Wire protocol, headers, retry/failover behavior, and the per-request
 * spoofing envelope (sessionId/requestId/labels) are copied from oh-my-pi's
 * google-gemini-cli provider for the google-antigravity provider.
 */
import { createHash, randomUUID } from "node:crypto";
import {
	type Api,
	type AssistantMessage,
	type AssistantMessageEventStream,
	type Context,
	calculateCost,
	createAssistantMessageEventStream,
	type Model,
	type TextContent,
	type ThinkingContent,
	type ThinkingLevel,
	type ToolCall,
} from "@earendil-works/pi-ai";
import {
	type AntigravitySimpleStreamOptions,
	buildBaseOptions,
	clampReasoning,
	convertMessages,
	convertTools,
	type FunctionCallingConfigMode,
	type GoogleContent,
	type GoogleThinkingConfig,
	type GoogleThinkingLevel,
	headersToRecord,
	isThinkingPart,
	mapStopReasonString,
	mapToolChoice,
	retainThoughtSignature,
	sanitizeSurrogates,
} from "./helpers.js";
import { type AntigravityModelProfile, getAntigravityProfile } from "./models.js";
import {
	ANTIGRAVITY_DAILY_ENDPOINT,
	ANTIGRAVITY_ENDPOINT_FALLBACKS,
	ANTIGRAVITY_SANDBOX_ENDPOINT,
	CLAUDE_THINKING_BETA_HEADER,
	getAntigravityModelWireProfile,
	getAntigravityUserAgent,
} from "./wire.js";

const GOOGLE_THINKING: Partial<Record<ThinkingLevel, number>> = {
	minimal: 1024,
	low: 4096,
	medium: 8192,
	high: 16384,
	xhigh: 24575,
	max: 32768,
};

const MIN_OUTPUT_TOKENS = 1024;

/** Map a normalized effort to Google's thinkingLevel enum for google-level models. */
function mapEffortToGoogleLevel(effort: ThinkingLevel, minimalIsLow: boolean): GoogleThinkingLevel {
	if (effort === "minimal") return minimalIsLow ? "LOW" : "MINIMAL";
	switch (effort) {
		case "low":
			return "LOW";
		case "medium":
			return "MEDIUM";
		default:
			return "HIGH";
	}
}

const FORCED_TOOL_DIRECTIVE =
	"TOOL-ONLY TURN. This turn accepts a tool call and nothing else; a text reply here is discarded unread and you will be re-prompted. Emit the tool call now.\n";

const MAX_RETRIES = 3;
const BASE_DELAY_MS = 1000;
const MAX_EMPTY_STREAM_RETRIES = 2;
const EMPTY_STREAM_BASE_DELAY_MS = 500;
const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000;

let toolCallCounter = 0;

interface CloudCodeAssistRequest {
	project: string;
	model: string;
	request: {
		contents: GoogleContent[];
		sessionId?: string;
		systemInstruction?: { role?: string; parts: { text: string }[] };
		generationConfig?: {
			maxOutputTokens?: number;
			temperature?: number;
			topP?: number;
			topK?: number;
			minP?: number;
			presencePenalty?: number;
			repetitionPenalty?: number;
			thinkingConfig?: GoogleThinkingConfig;
		};
		tools?: { functionDeclarations: Record<string, unknown>[] }[] | undefined;
		toolConfig?: {
			functionCallingConfig: {
				mode: FunctionCallingConfigMode;
				allowedFunctionNames?: string[];
			};
		};
		labels?: Record<string, string>;
	};
	requestType?: string;
	userAgent?: string;
	requestId?: string;
}

interface CloudCodeAssistResponseChunk {
	response?: {
		candidates?: Array<{
			content?: {
				role: string;
				parts?: Array<{
					text?: string;
					thought?: boolean;
					thoughtSignature?: string;
					functionCall?: { name: string; args: Record<string, unknown>; id?: string };
				}>;
			};
			finishReason?: string;
		}>;
		usageMetadata?: {
			promptTokenCount?: number;
			candidatesTokenCount?: number;
			thoughtsTokenCount?: number;
			totalTokenCount?: number;
			cachedContentTokenCount?: number;
		};
		modelVersion?: string;
		responseId?: string;
		promptFeedback?: { blockReason?: string; blockReasonMessage?: string };
	};
	error?: { code?: number; message?: string; status?: string };
	traceId?: string;
}

interface ParsedCredentials {
	accessToken: string;
	projectId: string;
	email?: string;
}

function parseCredentials(apiKeyRaw: string): ParsedCredentials {
	let raw: unknown;
	try {
		raw = JSON.parse(apiKeyRaw);
	} catch {
		throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate.");
	}
	if (typeof raw !== "object" || raw === null) {
		throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate.");
	}
	const parsed = raw as { token?: unknown; projectId?: unknown; project_id?: unknown; email?: unknown };
	const token = typeof parsed.token === "string" ? parsed.token : undefined;
	const projectId =
		typeof parsed.projectId === "string"
			? parsed.projectId
			: typeof parsed.project_id === "string"
				? parsed.project_id
				: undefined;
	const email = typeof parsed.email === "string" && parsed.email.length > 0 ? parsed.email : undefined;
	if (!token || !projectId) {
		throw new Error("Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate.");
	}
	return { accessToken: token, projectId, email };
}

function sleep(ms: number, signal?: AbortSignal): Promise<void> {
	return new Promise((resolve, reject) => {
		if (signal?.aborted) {
			reject(new Error("Request was aborted"));
			return;
		}
		const timeout = setTimeout(resolve, ms);
		signal?.addEventListener("abort", () => {
			clearTimeout(timeout);
			reject(new Error("Request was aborted"));
		});
	});
}

function extractRetryDelay(errorText: string, response?: Response | Headers): number | undefined {
	const normalizeDelay = (ms: number): number | undefined => (ms > 0 ? Math.ceil(ms + 1000) : undefined);
	const headers = response instanceof Headers ? response : response?.headers;
	if (headers) {
		const retryAfter = headers.get("retry-after");
		if (retryAfter) {
			const retryAfterSeconds = Number(retryAfter);
			if (Number.isFinite(retryAfterSeconds)) {
				const delay = normalizeDelay(retryAfterSeconds * 1000);
				if (delay !== undefined) return delay;
			}
			const retryAfterDate = new Date(retryAfter);
			const retryAfterMs = retryAfterDate.getTime();
			if (!Number.isNaN(retryAfterMs)) {
				const delay = normalizeDelay(retryAfterMs - Date.now());
				if (delay !== undefined) return delay;
			}
		}
		const rateLimitReset = headers.get("x-ratelimit-reset");
		if (rateLimitReset) {
			const resetSeconds = Number.parseInt(rateLimitReset, 10);
			if (!Number.isNaN(resetSeconds)) {
				const delay = normalizeDelay(resetSeconds * 1000 - Date.now());
				if (delay !== undefined) return delay;
			}
		}
		const rateLimitResetAfter = headers.get("x-ratelimit-reset-after");
		if (rateLimitResetAfter) {
			const resetAfterSeconds = Number(rateLimitResetAfter);
			if (Number.isFinite(resetAfterSeconds)) {
				const delay = normalizeDelay(resetAfterSeconds * 1000);
				if (delay !== undefined) return delay;
			}
		}
	}

	const durationMatch = errorText.match(/reset after (?:(\d+)h)?(?:(\d+)m)?(\d+(?:\.\d+)?)s/i);
	if (durationMatch) {
		const hours = durationMatch[1] ? parseInt(durationMatch[1], 10) : 0;
		const minutes = durationMatch[2] ? parseInt(durationMatch[2], 10) : 0;
		const seconds = parseFloat(durationMatch[3]);
		if (!Number.isNaN(seconds)) {
			const totalMs = ((hours * 60 + minutes) * 60 + seconds) * 1000;
			const delay = normalizeDelay(totalMs);
			if (delay !== undefined) return delay;
		}
	}

	const retryInMatch = errorText.match(/Please retry in ([0-9.]+)(ms|s)/i);
	if (retryInMatch?.[1]) {
		const value = parseFloat(retryInMatch[1]);
		if (!Number.isNaN(value) && value > 0) {
			const ms = retryInMatch[2].toLowerCase() === "ms" ? value : value * 1000;
			const delay = normalizeDelay(ms);
			if (delay !== undefined) return delay;
		}
	}

	const retryDelayMatch = errorText.match(/"retryDelay":\s*"([0-9.]+)(ms|s)"/i);
	if (retryDelayMatch?.[1]) {
		const value = parseFloat(retryDelayMatch[1]);
		if (!Number.isNaN(value) && value > 0) {
			const ms = retryDelayMatch[2].toLowerCase() === "ms" ? value : value * 1000;
			const delay = normalizeDelay(ms);
			if (delay !== undefined) return delay;
		}
	}

	return undefined;
}

function isRetryableError(status: number, errorText: string): boolean {
	if (status === 408 || status === 429 || status >= 500) {
		return true;
	}
	return /resource.?exhausted|rate.?limit|overloaded|service.?unavailable|other.?side.?closed/i.test(errorText);
}

function extractErrorMessage(errorText: string): string {
	try {
		const parsed = JSON.parse(errorText) as { error?: { message?: string } };
		if (parsed.error?.message) {
			return parsed.error.message;
		}
	} catch {
		// Not JSON, return as-is
	}
	return errorText;
}

function extractGoogleValidationUrl(errorBody: string): string | undefined {
	if (!errorBody.includes("VALIDATION_REQUIRED")) return undefined;
	const start = errorBody.indexOf("{");
	if (start === -1) return undefined;
	try {
		const parsed = JSON.parse(errorBody.slice(start)) as {
			error?: { details?: Array<{ reason?: string; metadata?: { validation_url?: string } }> };
		};
		const detail = parsed.error?.details?.find(
			(d) => d.reason === "VALIDATION_REQUIRED" && typeof d.metadata?.validation_url === "string",
		);
		return detail?.metadata?.validation_url;
	} catch {
		return undefined;
	}
}

function formatAntigravityApiError(status: number, errorText: string, email?: string): string {
	const validationUrl = extractGoogleValidationUrl(errorText);
	if (validationUrl) {
		const account = email ? ` for ${email}` : "";
		return `Cloud Code Assist API error (${status}): Account verification required${account}. Visit ${validationUrl} to continue, then retry your request.`;
	}
	return `Cloud Code Assist API error (${status}): ${extractErrorMessage(errorText)}`;
}

const INT63_MASK = (1n << 63n) - 1n;
function formatSignedDecimalSessionId(value: bigint): string {
	return `-${value.toString()}`;
}

function deriveSignedDecimalFromHash(text: string): string {
	const digest = createHash("sha256").update(text).digest();
	let value = 0n;
	for (let index = 0; index < 8; index += 1) {
		value = (value << 8n) | BigInt(digest[index] ?? 0);
	}
	return formatSignedDecimalSessionId(value & INT63_MASK);
}

function getFirstUserText(context: Context): string | undefined {
	for (const message of context.messages) {
		if (message.role !== "user") continue;
		if (typeof message.content === "string") return message.content;
		if (Array.isArray(message.content)) {
			const firstTextPart = message.content.find((item): item is TextContent => item.type === "text");
			return firstTextPart?.text;
		}
		return undefined;
	}
	return undefined;
}

function deriveAntigravitySessionId(context: Context): string {
	const text = getFirstUserText(context);
	if (text && text.trim().length > 0) {
		return deriveSignedDecimalFromHash(text);
	}
	const digest = createHash("sha256").update(randomUUID()).digest();
	let value = 0n;
	for (let index = 0; index < 8; index += 1) {
		value = (value << 8n) | BigInt(digest[index] ?? 0);
	}
	return formatSignedDecimalSessionId(value & INT63_MASK);
}

interface AntigravitySessionState {
	agentId: string;
	trajectoryId: string;
	sessionId: string;
	stepIndex: number;
	lastExecutionId?: string;
}

const sessionStates = new Map<string, AntigravitySessionState>();

function getSessionState(key: string, context: Context): AntigravitySessionState {
	const existing = sessionStates.get(key);
	if (existing) return existing;
	const state: AntigravitySessionState = {
		agentId: randomUUID(),
		trajectoryId: randomUUID(),
		sessionId: deriveAntigravitySessionId(context),
		stepIndex: 1,
	};
	sessionStates.set(key, state);
	return state;
}

function buildAntigravityRequestEnvelope(
	profile: AntigravityModelProfile,
	state: AntigravitySessionState,
): { sessionId: string; requestId: string; labels: Record<string, string> } {
	state.stepIndex = state.stepIndex + 1;
	const step = state.stepIndex;
	const requestId = `agent/${state.agentId}/${Date.now()}/${state.trajectoryId}/${step}`;
	const labels: Record<string, string> = {};
	if (state.lastExecutionId) labels.last_execution_id = state.lastExecutionId;
	labels.last_step_index = String(step - 1);
	if (profile.modelEnum !== undefined) labels.model_enum = profile.modelEnum;
	labels.trajectory_id = state.trajectoryId;
	const usageLabel = String(profile.identityClass === "anthropic");
	labels.used_claude = usageLabel;
	labels.used_claude_conservative = usageLabel;
	return { sessionId: state.sessionId, requestId, labels };
}

export interface AntigravityStreamOptions extends AntigravitySimpleStreamOptions {
	toolChoice?: "auto" | "none" | "any";
	topP?: number;
	topK?: number;
	minP?: number;
	presencePenalty?: number;
	repetitionPenalty?: number;
}

export function buildRequest(
	model: Model<Api>,
	context: Context,
	projectId: string,
	options: AntigravityStreamOptions = {},
	profile: AntigravityModelProfile = getAntigravityProfile(model.id),
	state: AntigravitySessionState = getSessionState(options?.sessionId ?? "default", context),
): CloudCodeAssistRequest {
	const contents = convertMessages(model, context);

	const generationConfig: CloudCodeAssistRequest["request"]["generationConfig"] = {};
	if (options.temperature !== undefined) {
		generationConfig.temperature = options.temperature;
	}
	if (options.maxTokens !== undefined) {
		generationConfig.maxOutputTokens = options.maxTokens;
	}
	if (options.topP !== undefined) {
		generationConfig.topP = options.topP;
	}
	if (options.topK !== undefined) {
		generationConfig.topK = options.topK;
	}
	if (options.minP !== undefined) {
		generationConfig.minP = options.minP;
	}
	if (options.presencePenalty !== undefined) {
		generationConfig.presencePenalty = options.presencePenalty;
	}
	if (options.repetitionPenalty !== undefined) {
		generationConfig.repetitionPenalty = options.repetitionPenalty;
	}

	const thinking = profile.thinking;
	if (model.reasoning && options.reasoning && options.reasoning !== "off") {
		const effort = clampReasoning(options.reasoning)!;
		if (thinking.mode === "google-level") {
			generationConfig.thinkingConfig = {
				includeThoughts: true,
				thinkingLevel: mapEffortToGoogleLevel(effort, thinking.minimalLevelIsLow ?? false),
			};
		} else {
			const customBudgets = options.thinkingBudgets;
			let thinkingBudget =
				customBudgets?.[effort] ?? thinking.effortBudgets?.[effort] ?? GOOGLE_THINKING[effort] ?? 0;
			const maxTokens = Math.min((options.maxTokens ?? model.maxTokens) + thinkingBudget, model.maxTokens);
			if (maxTokens <= thinkingBudget) {
				thinkingBudget = Math.max(0, maxTokens - MIN_OUTPUT_TOKENS);
			}
			if (thinkingBudget > 0) {
				generationConfig.thinkingConfig = { includeThoughts: true, thinkingBudget };
				generationConfig.maxOutputTokens = maxTokens;
			}
		}
	} else if (model.reasoning && thinking.suppressWhenOff) {
		// Omit thinkingConfig when off is not special-cased; otherwise CCA
		// re-applies the per-id baked server default and bills hidden reasoning.
		generationConfig.thinkingConfig = {
			includeThoughts: false,
			...(thinking.mode === "google-level" ? { thinkingLevel: "MINIMAL" } : { thinkingBudget: 0 }),
		};
	}

	const request: CloudCodeAssistRequest["request"] = { contents };

	if (context.systemPrompt) {
		request.systemInstruction = {
			role: "user",
			parts: [{ text: sanitizeSurrogates(context.systemPrompt) }],
		};
	}

	if (context.tools && context.tools.length > 0) {
		// Antigravity / Cloud Code Assist requires the legacy `parameters` field
		// regardless of model family.
		request.tools = convertTools(context.tools, true);
		if (options.toolChoice) {
			request.toolConfig = {
				functionCallingConfig: {
					mode: mapToolChoice(options.toolChoice),
				},
			};
			// Cloud Code Assist drops toolConfig on Antigravity's Gemini routes:
			// the backend answers in text under mode ANY and still emits calls
			// under NONE. Claude routes implement it, so only Gemini needs the
			// forced choice restated in the transcript.
			if (profile.identityClass !== "anthropic" && request.toolConfig?.functionCallingConfig.mode === "ANY") {
				contents.push({ role: "user", parts: [{ text: FORCED_TOOL_DIRECTIVE }] });
			}
		}
		// Antigravity's default tool mode is VALIDATED.
		if (!request.toolConfig) {
			request.toolConfig = {
				functionCallingConfig: { mode: "VALIDATED" },
			};
		}
	}

	// Claude on Antigravity always forces VALIDATED, even with no tools declared.
	if (profile.identityClass === "anthropic") {
		request.toolConfig = {
			functionCallingConfig: { mode: "VALIDATED" },
		};
	}

	const wireModelId = profile.wireModelId;
	const wireProfile = getAntigravityModelWireProfile(wireModelId);
	if (wireProfile) {
		generationConfig.maxOutputTokens = wireProfile.maxOutputTokens;
	}

	const envelope = buildAntigravityRequestEnvelope(profile, state);
	request.labels = envelope.labels;
	if (Object.keys(generationConfig).length > 0) {
		request.generationConfig = generationConfig;
	}
	request.sessionId = envelope.sessionId;

	return {
		project: projectId,
		requestId: envelope.requestId,
		request,
		model: wireModelId,
		userAgent: "antigravity",
		requestType: "agent",
	};
}

export function streamAntigravity(
	model: Model<Api>,
	context: Context,
	options?: AntigravityStreamOptions,
): AssistantMessageEventStream {
	const stream = createAssistantMessageEventStream();

	(async () => {
		const output: AssistantMessage = {
			role: "assistant",
			content: [],
			api: model.api,
			provider: model.provider,
			model: model.id,
			usage: {
				input: 0,
				output: 0,
				cacheRead: 0,
				cacheWrite: 0,
				totalTokens: 0,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
			},
			stopReason: "stop",
			timestamp: Date.now(),
		};

		try {
			const apiKeyRaw = options?.apiKey;
			if (!apiKeyRaw) {
				throw new Error(
					"Antigravity requires OAuth authentication. Use /login google-antigravity to authenticate.",
				);
			}
			const credentials = parseCredentials(apiKeyRaw);
			const profile = getAntigravityProfile(model.id);
			const state = getSessionState(options?.sessionId ?? "default", context);

			const baseUrl = model.baseUrl?.trim();
			const cleanUrl = baseUrl ? baseUrl.replace(/\/+$/, "") : "";
			let endpoints: string[];
			if (cleanUrl && cleanUrl !== ANTIGRAVITY_DAILY_ENDPOINT && cleanUrl !== ANTIGRAVITY_SANDBOX_ENDPOINT) {
				endpoints = [cleanUrl];
			} else {
				endpoints = [...ANTIGRAVITY_ENDPOINT_FALLBACKS];
			}

			let requestBody = buildRequest(model, context, credentials.projectId, options ?? {}, profile, state);
			const nextRequestBody = await options?.onPayload?.(requestBody, model);
			if (nextRequestBody !== undefined) {
				requestBody = nextRequestBody as CloudCodeAssistRequest;
			}

			const requestHeaders: Record<string, string> = {
				Authorization: `Bearer ${credentials.accessToken}`,
				"Content-Type": "application/json",
				Accept: "text/event-stream",
				"User-Agent": getAntigravityUserAgent(),
				...(profile.identityClass === "anthropic" && model.reasoning
					? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER }
					: {}),
				...options?.headers,
			};
			const requestBodyJson = JSON.stringify(requestBody);

			let response: Response | undefined;
			let lastError: Error | undefined;
			let requestUrl: string | undefined;
			let endpointIndex = 0;

			for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
				if (options?.signal?.aborted) {
					throw new Error("Request was aborted");
				}

				try {
					const endpoint = endpoints[endpointIndex];
					requestUrl = `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
					response = await fetch(requestUrl, {
						method: "POST",
						headers: requestHeaders,
						body: requestBodyJson,
						signal: options?.signal,
					});
					await options?.onResponse?.(
						{ status: response.status, headers: headersToRecord(response.headers) },
						model,
					);

					if (response.ok) {
						break;
					}

					const errorText = await response.text();

					if ((response.status === 403 || response.status === 404) && endpointIndex < endpoints.length - 1) {
						endpointIndex++;
						continue;
					}

					if (attempt < MAX_RETRIES && isRetryableError(response.status, errorText)) {
						if (endpointIndex < endpoints.length - 1) {
							endpointIndex++;
						}
						const serverDelay = extractRetryDelay(errorText, response);
						const delayMs = serverDelay ?? BASE_DELAY_MS * 2 ** attempt;
						const maxDelayMs = options?.maxRetryDelayMs ?? RATE_LIMIT_BUDGET_MS;
						if (maxDelayMs > 0 && serverDelay && serverDelay > maxDelayMs) {
							const delaySeconds = Math.ceil(serverDelay / 1000);
							throw new Error(
								`Server requested ${delaySeconds}s retry delay (max: ${Math.ceil(maxDelayMs / 1000)}s). ${extractErrorMessage(errorText)}`,
							);
						}
						await sleep(delayMs, options?.signal);
						continue;
					}

					throw new Error(formatAntigravityApiError(response.status, errorText, credentials.email));
				} catch (error) {
					if (error instanceof Error) {
						if (error.name === "AbortError" || error.message === "Request was aborted") {
							throw new Error("Request was aborted");
						}
					}
					lastError = error instanceof Error ? error : new Error(String(error));
					if (lastError.message === "fetch failed" && lastError.cause instanceof Error) {
						lastError = new Error(`Network error: ${lastError.cause.message}`);
					}
					if (attempt < MAX_RETRIES) {
						const delayMs = BASE_DELAY_MS * 2 ** attempt;
						await sleep(delayMs, options?.signal);
						continue;
					}
					throw lastError;
				}
			}

			if (!response || !response.ok) {
				throw lastError ?? new Error("Failed to get response after retries");
			}

			let started = false;
			const ensureStarted = () => {
				if (!started) {
					stream.push({ type: "start", partial: output });
					started = true;
				}
			};

			const resetOutput = () => {
				output.content = [];
				output.usage = {
					input: 0,
					output: 0,
					cacheRead: 0,
					cacheWrite: 0,
					totalTokens: 0,
					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
				};
				output.stopReason = "stop";
				output.errorMessage = undefined;
				output.timestamp = Date.now();
				started = false;
			};

			const streamResponse = async (activeResponse: Response): Promise<boolean> => {
				if (!activeResponse.body) {
					throw new Error("No response body");
				}

				let hasContent = false;
				let currentBlock: TextContent | ThinkingContent | null = null;
				const blocks = output.content;
				const blockIndex = () => blocks.length - 1;

				const reader = activeResponse.body.getReader();
				const decoder = new TextDecoder();
				let buffer = "";

				const abortHandler = () => {
					void reader.cancel().catch(() => {});
				};
				options?.signal?.addEventListener("abort", abortHandler);

				try {
					while (true) {
						if (options?.signal?.aborted) {
							throw new Error("Request was aborted");
						}

						const { done, value } = await reader.read();
						if (done) break;

						buffer += decoder.decode(value, { stream: true });
						const lines = buffer.split("\n");
						buffer = lines.pop() || "";

						for (const line of lines) {
							if (!line.startsWith("data:")) continue;

							const jsonStr = line.slice(5).trim();
							if (!jsonStr) continue;

							let chunk: CloudCodeAssistResponseChunk;
							try {
								chunk = JSON.parse(jsonStr);
							} catch {
								continue;
							}

							if (chunk.error) {
								const detail = chunk.error.message || chunk.error.status || "unknown error";
								throw new Error(`Cloud Code Assist stream error: ${detail}`);
							}

							const responseData = chunk.response;
							if (!responseData) continue;
							output.responseId ||= responseData.responseId;

							if (!responseData.candidates?.length && responseData.promptFeedback?.blockReason) {
								const detail = responseData.promptFeedback.blockReasonMessage;
								throw new Error(
									`Request blocked by Google (${responseData.promptFeedback.blockReason})${detail ? `: ${detail}` : ""}`,
								);
							}

							const candidate = responseData.candidates?.[0];
							if (candidate?.content?.parts) {
								for (const part of candidate.content.parts) {
									if (part.text !== undefined && part.text !== "") {
										hasContent = true;
										const isThinking = isThinkingPart(part);
										if (
											!currentBlock ||
											(isThinking && currentBlock.type !== "thinking") ||
											(!isThinking && currentBlock.type !== "text")
										) {
											if (currentBlock) {
												if (currentBlock.type === "text") {
													stream.push({
														type: "text_end",
														contentIndex: blockIndex(),
														content: currentBlock.text,
														partial: output,
													});
												} else {
													stream.push({
														type: "thinking_end",
														contentIndex: blockIndex(),
														content: currentBlock.thinking,
														partial: output,
													});
												}
											}
											if (isThinking) {
												currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
												output.content.push(currentBlock);
												ensureStarted();
												stream.push({
													type: "thinking_start",
													contentIndex: blockIndex(),
													partial: output,
												});
											} else {
												currentBlock = { type: "text", text: "" };
												output.content.push(currentBlock);
												ensureStarted();
												stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
											}
										}
										if (currentBlock.type === "thinking") {
											currentBlock.thinking += part.text;
											currentBlock.thinkingSignature = retainThoughtSignature(
												currentBlock.thinkingSignature,
												part.thoughtSignature,
											);
											stream.push({
												type: "thinking_delta",
												contentIndex: blockIndex(),
												delta: part.text,
												partial: output,
											});
										} else {
											currentBlock.text += part.text;
											currentBlock.textSignature = retainThoughtSignature(
												currentBlock.textSignature,
												part.thoughtSignature,
											);
											stream.push({
												type: "text_delta",
												contentIndex: blockIndex(),
												delta: part.text,
												partial: output,
											});
										}
									} else if (part.text === "" && part.thoughtSignature && !part.functionCall) {
										// Terminal thought signatures on empty text parts (e.g. at stream end).
										if (currentBlock) {
											if (currentBlock.type === "thinking") {
												currentBlock.thinkingSignature = retainThoughtSignature(
													currentBlock.thinkingSignature,
													part.thoughtSignature,
												);
											} else {
												currentBlock.textSignature = retainThoughtSignature(
													currentBlock.textSignature,
													part.thoughtSignature,
												);
											}
										}
									}

									if (part.functionCall) {
										hasContent = true;
										if (currentBlock) {
											if (currentBlock.type === "text") {
												stream.push({
													type: "text_end",
													contentIndex: blockIndex(),
													content: currentBlock.text,
													partial: output,
												});
											} else {
												stream.push({
													type: "thinking_end",
													contentIndex: blockIndex(),
													content: currentBlock.thinking,
													partial: output,
												});
											}
											currentBlock = null;
										}

										const providedId = part.functionCall.id;
										const needsNewId =
											!providedId ||
											output.content.some((b) => b.type === "toolCall" && b.id === providedId);
										const toolCallId = needsNewId
											? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
											: providedId;

										const toolCall: ToolCall = {
											type: "toolCall",
											id: toolCallId,
											name: part.functionCall.name || "",
											arguments: (part.functionCall.args as Record<string, unknown>) ?? {},
											...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
										};

										output.content.push(toolCall);
										ensureStarted();
										stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
										stream.push({
											type: "toolcall_delta",
											contentIndex: blockIndex(),
											delta: JSON.stringify(toolCall.arguments),
											partial: output,
										});
										stream.push({
											type: "toolcall_end",
											contentIndex: blockIndex(),
											toolCall,
											partial: output,
										});
									}
								}
							}

							if (candidate?.finishReason) {
								output.stopReason = mapStopReasonString(candidate.finishReason);
								if (output.content.some((b) => b.type === "toolCall")) {
									output.stopReason = "toolUse";
								}
							}

							if (responseData.usageMetadata) {
								const promptTokens = responseData.usageMetadata.promptTokenCount || 0;
								const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0;
								const thinkingTokens = responseData.usageMetadata.thoughtsTokenCount || 0;
								output.usage = {
									input: promptTokens - cacheReadTokens,
									output: (responseData.usageMetadata.candidatesTokenCount || 0) + thinkingTokens,
									cacheRead: cacheReadTokens,
									cacheWrite: 0,
									totalTokens: responseData.usageMetadata.totalTokenCount || 0,
									cost: {
										input: 0,
										output: 0,
										cacheRead: 0,
										cacheWrite: 0,
										total: 0,
									},
								};
								calculateCost(model, output.usage);
							}
						}
					}
				} finally {
					options?.signal?.removeEventListener("abort", abortHandler);
				}

				if (currentBlock) {
					if (currentBlock.type === "text") {
						stream.push({
							type: "text_end",
							contentIndex: blockIndex(),
							content: currentBlock.text,
							partial: output,
						});
					} else {
						stream.push({
							type: "thinking_end",
							contentIndex: blockIndex(),
							content: currentBlock.thinking,
							partial: output,
						});
					}
				}

				return hasContent;
			};

			let receivedContent = false;
			let currentResponse = response;

			for (let emptyAttempt = 0; emptyAttempt <= MAX_EMPTY_STREAM_RETRIES; emptyAttempt++) {
				if (options?.signal?.aborted) {
					throw new Error("Request was aborted");
				}

				if (emptyAttempt > 0) {
					const backoffMs = EMPTY_STREAM_BASE_DELAY_MS * 2 ** (emptyAttempt - 1);
					await sleep(backoffMs, options?.signal);

					if (!requestUrl) {
						throw new Error("Missing request URL");
					}

					currentResponse = await fetch(requestUrl, {
						method: "POST",
						headers: requestHeaders,
						body: requestBodyJson,
						signal: options?.signal,
					});
					await options?.onResponse?.(
						{ status: currentResponse.status, headers: headersToRecord(currentResponse.headers) },
						model,
					);

					if (!currentResponse.ok) {
						const retryErrorText = await currentResponse.text();
						throw new Error(`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`);
					}
				}

				const streamed = await streamResponse(currentResponse);
				if (streamed) {
					receivedContent = true;
					break;
				}

				if (emptyAttempt < MAX_EMPTY_STREAM_RETRIES) {
					resetOutput();
				}
			}

			if (!receivedContent) {
				throw new Error("Cloud Code Assist API returned an empty response");
			}

			if (options?.signal?.aborted) {
				throw new Error("Request was aborted");
			}

			if (output.stopReason === "aborted" || output.stopReason === "error") {
				throw new Error(output.errorMessage ?? "An unknown error occurred");
			}

			state.lastExecutionId = output.responseId;
			stream.push({ type: "done", reason: output.stopReason, message: output });
			stream.end();
		} catch (error) {
			output.stopReason = options?.signal?.aborted ? "aborted" : "error";
			output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
			stream.push({ type: "error", reason: output.stopReason, error: output });
			stream.end();
		}
	})();

	return stream;
}

export function streamSimpleAntigravity(
	model: Model<Api>,
	context: Context,
	options?: AntigravitySimpleStreamOptions,
): AssistantMessageEventStream {
	const apiKey = options?.apiKey;
	if (!apiKey) {
		throw new Error("Antigravity requires OAuth authentication. Use /login google-antigravity to authenticate.");
	}
	const base = buildBaseOptions(model, options, apiKey);
	return streamAntigravity(model, context, {
		...base,
		reasoning: options?.reasoning,
		thinkingBudgets: options?.thinkingBudgets,
	});
}