import type { AuthInteraction, AuthPrompt, OAuthCredential } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import type { CreditStatusRuntime } from "./credits.js";
import { HYPER_API_BASE_URL, HYPER_API_KEY, PROVIDER_DISPLAY_NAME, PROVIDER_NAME } from "./hyper.js";
import { fetchHyperModels } from "./models.js";
import { createNotifier } from "./notify.js";
import { loginHyper, refreshHyperToken } from "./oauth.js";

// OMP (Oh My Pi) serves @earendil-works/pi-ai imports through a compat shim
// that lacks createProvider/envApiKeyAuth/openAICompletionsApi, and its
// registerProvider expects a flat ProviderConfig (api as a string id,
// fetchDynamicModels, oauth.{login,refreshToken,getApiKey}) rather than
// upstream's createProvider() result. OMP's runtime streams built-in API ids
// like "openai-completions" natively, so no streamSimple is needed; passing
// upstream's lazy API object here is what crashed OMP's dispatch path.

interface OmpOAuthLoginCallbacks {
	onAuth(info: { url: string; launchUrl?: string; instructions?: string }): void;
	onPrompt(prompt: { message: string; placeholder?: string; allowEmpty?: boolean }): Promise<string>;
	onProgress?(message: string): void;
	signal?: AbortSignal;
}

interface OmpOAuthCredentials {
	refresh: string;
	access: string;
	expires: number;
	[key: string]: unknown;
}

// A real Charm Hyper credential is either an sk-hyper-* API key or a 3-segment
// JWT (the OAuth access token). Anything else in HYPER_API_KEY — a leftover
// placeholder, an unquoted shell fragment, whitespace — is guaranteed to 401,
// so flag it loudly at startup instead of letting the first request fail with
// Charm's opaque "token is malformed" error.
const HYPER_API_KEY_PREFIX = "sk-hyper-";
const JWT_SEGMENT_COUNT = 3;

function looksValidHyperApiKey(key: string): boolean {
	const trimmed = key.trim();
	if (trimmed.startsWith(HYPER_API_KEY_PREFIX)) return true;
	return trimmed.split(".").length === JWT_SEGMENT_COUNT;
}

// OMP's ProviderConfig. It diverges from upstream's @earendil-works/pi-coding-agent
// ProviderConfig: OMP uses `fetchDynamicModels` where upstream uses `refreshModels`,
// and its oauth credentials type omits the `type` tag. Declared locally because the
// dev-installed upstream types do not expose OMP's shape.
interface OmpProviderConfig {
	baseUrl?: string;
	apiKey?: string;
	api?: string;
	headers?: Record<string, string>;
	authHeader?: boolean;
	fetchDynamicModels?: (apiKey: string | undefined) => Promise<readonly unknown[]>;
	oauth?: {
		name: string;
		login(callbacks: OmpOAuthLoginCallbacks): Promise<OmpOAuthCredentials>;
		refreshToken?(credentials: OmpOAuthCredentials): Promise<OmpOAuthCredentials>;
		getApiKey?(credentials: OmpOAuthCredentials): string;
	};
}

function toAuthInteraction(callbacks: OmpOAuthLoginCallbacks): AuthInteraction {
	return {
		signal: callbacks.signal,
		prompt: (prompt: AuthPrompt) => callbacks.onPrompt(prompt),
		notify: (event) => {
			switch (event.type) {
				case "device_code":
					callbacks.onAuth({
						url: event.verificationUri,
						instructions: `Enter code: ${event.userCode}`,
					});
					break;
				case "auth_url":
					callbacks.onAuth({ url: event.url, instructions: event.instructions });
					break;
				case "progress":
				case "info":
					callbacks.onProgress?.(event.message);
					break;
			}
		},
	};
}

function toOmpCredentials(credential: OAuthCredential): OmpOAuthCredentials {
	const { type: _type, ...credentials } = credential;
	return credentials;
}

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

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

export default function (pi: ExtensionAPI) {
	const notifier = createNotifier();
	let creditStatusState: CreditStatusState = { kind: "idle" };
	let pendingCreditStatusRefresh: PendingCreditStatusRefresh | undefined;
	let creditStatusRefreshWork: NodeJS.Immediate | 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("Hyper status support was disposed"));
		}

		const operation = import("./credits.js").then(({ createCreditStatusRuntime }) => {
			const runtime = createCreditStatusRuntime(notifier.warn);
			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 schedulePendingCreditStatusRefresh(): void {
		if (!pendingCreditStatusRefresh || creditStatusRefreshWork !== undefined) return;

		// The module transform evaluates credits.ts synchronously once the
		// import starts. A macrotask lets OMP finish its awaited lifecycle
		// dispatch before that work.
		const scheduled = setImmediate(() => {
			void loadCreditStatus()
				.then((runtime) => {
					if (creditStatusRefreshWork !== scheduled) return;
					const refresh = pendingCreditStatusRefresh;
					pendingCreditStatusRefresh = undefined;
					creditStatusRefreshWork = undefined;
					if (refresh) {
						void runtime.refresh(refresh.ctx, refresh.model).catch((error: unknown) => {
							if (creditStatusState.kind !== "disposed") {
								notifier.warn(`Unable to refresh Hyper status: ${String(error)}`);
							}
						});
					}
				})
				.catch((error: unknown) => {
					if (creditStatusRefreshWork === scheduled && creditStatusState.kind !== "disposed") {
						pendingCreditStatusRefresh = undefined;
						notifier.warn(`Unable to load Hyper status support: ${String(error)}`);
					}
				})
				.finally(() => {
					if (creditStatusRefreshWork !== scheduled) return;
					creditStatusRefreshWork = undefined;
					schedulePendingCreditStatusRefresh();
				});
		});
		creditStatusRefreshWork = scheduled;
	}

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

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

	pi.on("session_start", (_event, ctx) => {
		notifier.activate(ctx);
		if (!ctx.hasUI) return;
		if (ctx.model?.provider !== PROVIDER_NAME) {
			deactivateCreditStatus(ctx, ctx.model);
			return;
		}
		scheduleCreditStatusRefresh(ctx, ctx.model);
	});

	if (HYPER_API_KEY !== undefined && !looksValidHyperApiKey(HYPER_API_KEY)) {
		notifier.warn(
			`HYPER_API_KEY is set but is not a recognizable Charm credential (expected an ${HYPER_API_KEY_PREFIX}* key or a JWT). ` +
				`Every hyper request will 401. Fix or unset HYPER_API_KEY, or use /login ${PROVIDER_NAME} for OAuth.`,
		);
	}

	pi.registerProvider(PROVIDER_NAME, {
		baseUrl: HYPER_API_BASE_URL,
		apiKey: HYPER_API_KEY,
		api: "openai-completions",
		// OMP hands the resolved credential to the callback when it has one;
		// without one the request stays headerless so OMP's cached catalog
		// keeps restoring without authentication.
		fetchDynamicModels: (apiKey) => fetchHyperModels(apiKey ? { token: apiKey } : {}),
		oauth: {
			name: PROVIDER_DISPLAY_NAME,
			login: (callbacks: OmpOAuthLoginCallbacks) => loginHyper(toAuthInteraction(callbacks)).then(toOmpCredentials),
			refreshToken: (credentials: OmpOAuthCredentials) =>
				refreshHyperToken({ type: "oauth", ...credentials } as OAuthCredential).then(toOmpCredentials),
			getApiKey: (credentials: OmpOAuthCredentials) => credentials.access,
		},
	} satisfies OmpProviderConfig as unknown as Parameters<ExtensionAPI["registerProvider"]>[1]);

	pi.registerCommand("hyper-status", {
		description: "Configure the Charm Hyper footer status",
		handler: async (args, ctx) => {
			try {
				const runtime = await loadCreditStatus();
				if (creditStatusState.kind === "disposed") return;
				await runtime.handleCommand(args, ctx);
			} catch (error) {
				if (creditStatusState.kind === "disposed") return;
				ctx.ui.notify(`Unable to load Hyper status support: ${String(error)}`, "warning");
			}
		},
	});

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

	pi.on("message_end", (event, ctx) => {
		if (creditStatusState.kind === "disposed" || !ctx.hasUI) return;
		if (event.message.role === "assistant" && ctx.model?.provider === PROVIDER_NAME) {
			scheduleCreditStatusRefresh(ctx, ctx.model);
		}
	});

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