import { createProvider, envApiKeyAuth } from "@earendil-works/pi-ai";
import type { ModelsStoreEntry } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import type { CreditStatusRuntime } from "./credits.ts";
import { explabsResponsesApi } from "./responses.ts";
import {
	API_KEY_ENV,
	BASE_URL_ENV,
	DEFAULT_API_BASE_URL,
	PROVIDER_ID,
	PROVIDER_NAME,
	USER_AGENT,
} from "./config.ts";
import { restoreSessionModel } from "./session-model.ts";
import { normalizeApiBaseUrl } from "./urls.ts";

function resolveBaseUrl(): string {
	const fromEnv = process.env[BASE_URL_ENV];
	return normalizeApiBaseUrl(fromEnv && fromEnv.length > 0 ? fromEnv : DEFAULT_API_BASE_URL);
}

export function rebaseCachedExplabsModels(
	stored: Readonly<ModelsStoreEntry> | undefined,
	baseUrl: string,
): ModelsStoreEntry | undefined {
	if (!stored) {
		return undefined;
	}
	return {
		...stored,
		models: stored.models.map((model) =>
			model.provider === PROVIDER_ID ? { ...model, baseUrl } : model,
		),
	};
}

export default function experientialLabs(pi: ExtensionAPI): void {
	const baseUrl = resolveBaseUrl();
	const provider = createProvider({
		id: PROVIDER_ID,
		name: PROVIDER_NAME,
		baseUrl,
		headers: { "User-Agent": USER_AGENT },
		auth: {
			apiKey: envApiKeyAuth(`${PROVIDER_NAME} API key`, [API_KEY_ENV]),
		},
		models: [],
		fetchModels: async ({ credential, signal }) => {
			const { apiKeyFromCredential, loadModels } = await import("./catalog.ts");
			return loadModels({
				baseUrl,
				apiKey: apiKeyFromCredential(credential),
				signal,
			});
		},
		api: explabsResponsesApi(),
	});
	const refreshModels = provider.refreshModels;
	if (refreshModels) {
		provider.refreshModels = (context) =>
			refreshModels.call(provider, {
				...context,
				stored: rebaseCachedExplabsModels(context.stored, baseUrl),
			});
	}

	pi.registerProvider(provider);

	type CreditStatusState =
		| { kind: "idle" }
		| { kind: "loading"; operation: Promise<CreditStatusRuntime> }
		| { kind: "ready"; runtime: CreditStatusRuntime }
		| { kind: "disposed" };

	type PendingRefresh = { ctx: ExtensionContext; model: ExtensionContext["model"] };

	let creditStatusState: CreditStatusState = { kind: "idle" };
	let pendingRefresh: PendingRefresh | undefined;
	let creditStatusWork: ReturnType<typeof setImmediate> | undefined;

	function loadCreditStatus(): Promise<CreditStatusRuntime> {
		if (creditStatusState.kind === "ready") {
			return Promise.resolve(creditStatusState.runtime);
		}
		if (creditStatusState.kind === "loading") {
			return creditStatusState.operation;
		}
		if (creditStatusState.kind === "disposed") {
			return Promise.reject(new Error("Experiential Labs credits status was disposed"));
		}

		const operation = import("./credits.ts").then(({ createCreditStatusRuntime }) => {
			const runtime = createCreditStatusRuntime(baseUrl);
			if (creditStatusState.kind === "disposed") {
				runtime.dispose();
				return runtime;
			}
			creditStatusState = { kind: "ready", runtime };
			return runtime;
		});
		creditStatusState = { kind: "loading", operation };
		void operation.catch(() => {
			if (creditStatusState.kind === "loading" && creditStatusState.operation === operation) {
				creditStatusState = { kind: "idle" };
			}
		});
		return operation;
	}

	function schedulePendingRefresh(): void {
		if (!pendingRefresh || creditStatusWork !== undefined) {
			return;
		}

		const scheduled = setImmediate(() => {
			void loadCreditStatus()
				.then((runtime) => {
					if (creditStatusWork !== scheduled) {
						return;
					}
					const refresh = pendingRefresh;
					pendingRefresh = undefined;
					creditStatusWork = undefined;
					if (refresh) {
						void runtime.refresh(refresh.ctx, refresh.model);
					}
				})
				.catch(() => {
					if (creditStatusWork === scheduled && creditStatusState.kind !== "disposed") {
						pendingRefresh = undefined;
					}
				})
				.finally(() => {
					if (creditStatusWork !== scheduled) {
						return;
					}
					creditStatusWork = undefined;
					schedulePendingRefresh();
				});
		});
		creditStatusWork = scheduled;
	}

	function scheduleCreditStatusRefresh(ctx: ExtensionContext, model: ExtensionContext["model"]): void {
		pendingRefresh = { ctx, model };
		schedulePendingRefresh();
	}

	function deactivateCreditStatus(ctx: ExtensionContext, model: ExtensionContext["model"]): void {
		pendingRefresh = undefined;
		if (creditStatusWork !== undefined) {
			clearImmediate(creditStatusWork);
			creditStatusWork = undefined;
		}
		if (creditStatusState.kind === "ready") {
			void creditStatusState.runtime.refresh(ctx, model);
		}
		if (ctx.hasUI) {
			ctx.ui.setStatus(PROVIDER_ID, undefined);
		}
	}

	pi.on("session_start", async (_event, ctx) => {
		await restoreSessionModel({
			providerId: PROVIDER_ID,
			entries: ctx.sessionManager.getBranch(),
			activeModel: ctx.model,
			registry: {
				find: (provider, modelId) => ctx.modelRegistry.find(provider, modelId),
				hasConfiguredAuth: (model) => ctx.modelRegistry.hasConfiguredAuth(model),
				refresh: (provider) => ctx.modelRegistry.refresh({ providers: [provider], allowNetwork: false }),
			},
			setModel: (model) => pi.setModel(model),
		});
		if (!ctx.hasUI) {
			return;
		}
		if (ctx.model?.provider !== PROVIDER_ID) {
			deactivateCreditStatus(ctx, ctx.model);
			return;
		}
		scheduleCreditStatusRefresh(ctx, ctx.model);
	});

	pi.on("model_select", (event, ctx) => {
		if (!ctx.hasUI) {
			return;
		}
		if (event.model.provider !== PROVIDER_ID) {
			deactivateCreditStatus(ctx, event.model);
			return;
		}
		scheduleCreditStatusRefresh(ctx, event.model);
	});

	pi.on("agent_settled", (_event, ctx) => {
		if (creditStatusState.kind === "disposed" || !ctx.hasUI || ctx.model?.provider !== PROVIDER_ID) {
			return;
		}
		scheduleCreditStatusRefresh(ctx, ctx.model);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		pendingRefresh = undefined;
		if (creditStatusWork !== undefined) {
			clearImmediate(creditStatusWork);
		}
		creditStatusWork = undefined;
		if (creditStatusState.kind === "ready") {
			creditStatusState.runtime.dispose();
		}
		creditStatusState = { kind: "disposed" };
		if (ctx.hasUI) {
			ctx.ui.setStatus(PROVIDER_ID, undefined);
		}
	});
}