import type { Api, Model } from "@earendil-works/pi-ai";

/**
 * The registry hands back models for whichever API their catalog entry uses.
 */
type AnyModel = Model<Api>;

/**
 * A provider/model pair recorded on the session branch.
 */
export interface SessionModel {
	provider: string;
	modelId: string;
}

/**
 * Minimal view of the model registry needed to reconcile a session model.
 */
export interface SessionModelRegistry {
	find(provider: string, modelId: string): AnyModel | undefined;
	hasConfiguredAuth(model: AnyModel): boolean;
	refresh(provider: string): Promise<unknown>;
}

export interface RestoreSessionModelOptions {
	providerId: string;
	entries: readonly unknown[];
	activeModel: { provider: string; id: string } | undefined;
	registry: SessionModelRegistry;
	setModel(model: AnyModel): void;
	attempts?: number;
	retryDelayMs?: number;
	sleep?: (ms: number) => Promise<void>;
}

export type RestoreSessionModelResult = "skipped" | "restored" | "unavailable";

export const DEFAULT_RESTORE_ATTEMPTS = 10;
export const DEFAULT_RESTORE_RETRY_DELAY_MS = 50;

function asRecord(value: unknown): Record<string, unknown> | undefined {
	return value !== null && typeof value === "object" && !Array.isArray(value)
		? (value as Record<string, unknown>)
		: undefined;
}

function asString(value: unknown): string | undefined {
	return typeof value === "string" && value.length > 0 ? value : undefined;
}

/**
 * Recover the provider/model the session was last using.
 *
 * Pi folds `model_change` entries and the `provider`/`model` recorded on assistant
 * messages into the same "current model" notion, so both are considered here.
 */
export function recordedSessionModel(entries: readonly unknown[]): SessionModel | undefined {
	let recorded: SessionModel | undefined;
	for (const entry of entries) {
		const record = asRecord(entry);
		if (!record) {
			continue;
		}
		if (record.type === "model_change") {
			const provider = asString(record.provider);
			const modelId = asString(record.modelId);
			if (provider && modelId) {
				recorded = { provider, modelId };
			}
			continue;
		}
		if (record.type === "message") {
			const message = asRecord(record.message);
			if (message?.role !== "assistant") {
				continue;
			}
			const provider = asString(message.provider);
			const modelId = asString(message.model);
			if (provider && modelId) {
				recorded = { provider, modelId };
			}
		}
	}
	return recorded;
}

/**
 * Re-select the session's provider model when pi could not restore it on its own.
 *
 * Pi restores the session model during `createAgentSession` by requiring both the
 * catalog entry *and* `hasConfiguredAuth()`. The catalog is restored synchronously
 * from the model store, but the configured-provider snapshot is filled in by an
 * asynchronous availability refresh, so a session recorded on this provider can
 * fall back to the default model with
 * `Could not restore model <provider>/<modelId>`.
 *
 * This runs from `session_start`, after registry composition, and repairs that
 * fallback by waiting (bounded) for the provider to become available again.
 */
export async function restoreSessionModel(
	options: RestoreSessionModelOptions,
): Promise<RestoreSessionModelResult> {
	const recorded = recordedSessionModel(options.entries);
	if (!recorded || recorded.provider !== options.providerId) {
		return "skipped";
	}
	if (options.activeModel?.provider === recorded.provider && options.activeModel.id === recorded.modelId) {
		return "skipped";
	}

	const model = options.registry.find(recorded.provider, recorded.modelId);
	if (!model) {
		// The slug is gone from the catalog; let pi's fallback model stand.
		return "skipped";
	}

	if (!options.registry.hasConfiguredAuth(model)) {
		// Provider-scoped, offline refresh: it rebuilds the configured-provider
		// snapshot without triggering a catalog fetch.
		await options.registry.refresh(options.providerId).catch(() => undefined);
	}

	const attempts = options.attempts ?? DEFAULT_RESTORE_ATTEMPTS;
	const retryDelayMs = options.retryDelayMs ?? DEFAULT_RESTORE_RETRY_DELAY_MS;
	const sleep = options.sleep ?? defaultSleep;

	for (let attempt = 0; attempt < attempts; attempt++) {
		if (options.registry.hasConfiguredAuth(model)) {
			options.setModel(model);
			return "restored";
		}
		await sleep(retryDelayMs);
	}
	return "unavailable";
}

function defaultSleep(ms: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, ms));
}