/**
 * Provider override layer.
 *
 * For each managed provider, registers a replacement that resolves
 * credentials from the vault backend before each API request.
 *
 * Key insight: pi-ai's built-in providers already read `options.apiKey`
 * from SimpleStreamOptions. We do NOT need to create a wrapper stream,
 * iterate events, and forward them. We just resolve the credential and
 * pass it through `options.apiKey` to the original provider's streamSimple.
 *
 * The only case where a full stream wrapper is justified is when the
 * extension needs to intercept mid-stream events (e.g., multicodex's
 * quota-based retry/rotation). Vault does not need that.
 */
import type { ExtensionAPI, ProviderConfig } from "@mariozechner/pi-coding-agent";
import {
	type Api,
	type AssistantMessageEventStream,
	type Context,
	type Model,
	type SimpleStreamOptions,
	createAssistantMessageEventStream,
	getApiProvider,
} from "@mariozechner/pi-ai";
import type { VaultController } from "../config/controller.js";
import { mirrorProvider, getRegisteredProviderIds } from "./mirror.js";
import { createOAuthWrapper } from "./oauth-wrapper.js";

// TODO: extract to pi-provider-utils (duplicated in multicodex stream-wrapper.ts)
function getErrorMessage(error: unknown): string {
	if (error instanceof Error) return error.message;
	return typeof error === "string" ? error : JSON.stringify(error);
}

/**
 * Resolve the API key for a provider from the vault backend.
 * Handles both api_key and oauth entries, including token refresh.
 */
async function resolveApiKey(
	providerId: string,
	controller: VaultController,
): Promise<string> {
	const backend = controller.getBackend();
	const entry = await backend.get(providerId);

	if (!entry) {
		throw new Error(
			`No credentials in vault for "${providerId}". Use /login or /vault setup.`,
		);
	}

	if (entry.type === "api_key") {
		return entry.key;
	}

	// OAuth entry -- check if token needs refresh
	if (Date.now() < entry.expires) {
		return entry.access;
	}

	// Token expired -- attempt refresh
	const oauthWrapper = createOAuthWrapper(providerId, () => backend);
	if (oauthWrapper) {
		const refreshed = await oauthWrapper.refreshToken({
			access: entry.access,
			refresh: entry.refresh,
			expires: entry.expires,
			accountId: entry.accountId,
		});
		return oauthWrapper.getApiKey(refreshed);
	}

	// No OAuth provider available for refresh -- use expired token as-is
	// (the provider will surface an auth error if it's truly expired)
	return entry.access;
}

/**
 * Build an error AssistantMessage for a model.
 * // TODO: extract to pi-provider-utils (duplicated in multicodex)
 */
function buildErrorMessage(model: Model<Api>, message: string) {
	return {
		role: "assistant" as const,
		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: "error" as const,
		errorMessage: message,
		timestamp: Date.now(),
	};
}

/** Push an error event into a stream. */
function pushError(stream: AssistantMessageEventStream, model: Model<Api>, message: string): void {
	stream.push({ type: "error", reason: "error", error: buildErrorMessage(model, message) });
}

/** Create a stream that immediately emits a single error and closes. */
function createErrorStream(model: Model<Api>, message: string): AssistantMessageEventStream {
	const stream = createAssistantMessageEventStream();
	pushError(stream, model, message);
	return stream;
}

// TODO: extract to pi-provider-utils (generic stream forwarding utility)
/**
 * Pipe all events from source into target, then end the target.
 *
 * push() marks the target as done when it sees a terminal event (done/error),
 * which prevents further pushes. But end() is still needed to flush any
 * consumers waiting on the target's async iterator after the last event.
 */
async function pipeStream(
	source: AssistantMessageEventStream,
	target: AssistantMessageEventStream,
): Promise<void> {
	for await (const event of source) {
		target.push(event);
	}
	target.end();
}

/**
 * Create a streamSimple function that resolves credentials from the vault
 * then delegates to the original provider.
 *
 * streamSimple must return synchronously, but credential resolution is async.
 * We return a forwarding stream and kick off the async resolve + pipe.
 */
function createVaultStreamSimple(
	providerId: string,
	controller: VaultController,
) {
	const originalApi = getApiProvider(providerId);

	return (
		model: Model<Api>,
		context: Context,
		options?: SimpleStreamOptions,
	): AssistantMessageEventStream => {
		if (!originalApi) {
			return createErrorStream(model, `No API provider for "${providerId}".`);
		}

		const stream = createAssistantMessageEventStream();

		void resolveApiKey(providerId, controller)
			.then((apiKey) => {
				const inner = originalApi.streamSimple(model, context, { ...options, apiKey });
				return pipeStream(inner, stream);
			})
			.catch((error: unknown) => {
				pushError(stream, model, `Vault: ${getErrorMessage(error)}`);
			});

		return stream;
	};
}

/**
 * Override all managed providers with vault-backed credential injection.
 * Returns the list of provider IDs that were overridden.
 */
export function overrideManagedProviders(
	pi: ExtensionAPI,
	controller: VaultController,
): string[] {
	const overridden: string[] = [];
	const providerIds = getRegisteredProviderIds();

	for (const providerId of providerIds) {
		if (!controller.isProviderManaged(providerId)) {
			continue;
		}

		const mirror = mirrorProvider(providerId);
		if (!mirror) {
			continue;
		}

		const oauthWrapper = createOAuthWrapper(providerId, () =>
			controller.getBackend(),
		);

		const providerConfig: ProviderConfig = {
			baseUrl: mirror.baseUrl,
			apiKey: "managed-by-vault",
			api: mirror.api,
			streamSimple: createVaultStreamSimple(providerId, controller),
			models: mirror.models.map((m) => ({
				...m,
				input: [...m.input],
				cost: { ...m.cost },
			})),
		};

		if (oauthWrapper) {
			providerConfig.oauth = oauthWrapper;
		}

		pi.registerProvider(providerId, providerConfig);
		overridden.push(providerId);
	}

	return overridden;
}