import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { readStoredCredential } from "@earendil-works/pi-coding-agent";
import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
import type {
	AssistantMessage,
	AssistantMessageEvent,
	AssistantMessageEventStream,
	Context,
	Model,
	OAuthAuth,
	Provider,
	ProviderAuth,
} from "@earendil-works/pi-ai";
import { builtinProviders } from "@earendil-works/pi-ai/providers/all";

const BUILTIN_CODEX_PROVIDER = "openai-codex";
const CODEX_API = "openai-codex-responses" as const;

const ALIASES = [
	{ id: "openai-codex-personal", name: "ChatGPT Plus/Pro (Codex Subscription) - Personal" },
	{ id: "openai-codex-work", name: "ChatGPT Plus/Pro (Codex Subscription) - Work" },
] as const;

const CODEX_ENV_KEYS = [
	"OPENAI_CODEX_OAUTH_TOKEN",
	"OPENAI_CODEX_ACCESS_TOKEN",
	"OPENAI_CODEX_ACCOUNT_ID",
	"CHATGPT_ACCOUNT_ID",
] as const;
const [OAUTH_TOKEN_ENV, ACCESS_TOKEN_ENV, ACCOUNT_ID_ENV, CHATGPT_ACCOUNT_ID_ENV] = CODEX_ENV_KEYS;

type CodexEnvKey = (typeof CODEX_ENV_KEYS)[number];

type StoredCredential = {
	type?: string;
	accountId?: unknown;
};

export default function codexAliases(pi: ExtensionAPI) {
	const codexProvider = getBuiltinCodexProvider();

	for (const alias of ALIASES) {
		pi.registerProvider(createCodexAliasProvider(codexProvider, alias.id, alias.name));
	}

	const originalSubBarEnv = Object.fromEntries(
		CODEX_ENV_KEYS.map((key) => [key, process.env[key]]),
	) as Record<CodexEnvKey, string | undefined>;
	const refreshTimers = new Set<ReturnType<typeof setTimeout>>();
	let ownsSubBarEnv = false;
	let syncGeneration = 0;

	function restoreSubBarEnv(): boolean {
		if (!ownsSubBarEnv) return false;

		for (const key of CODEX_ENV_KEYS) {
			const value = originalSubBarEnv[key];
			if (value === undefined) delete process.env[key];
			else process.env[key] = value;
		}
		ownsSubBarEnv = false;
		return true;
	}

	function isAliasProvider(provider: string | undefined): provider is (typeof ALIASES)[number]["id"] {
		return provider !== undefined && ALIASES.some((alias) => alias.id === provider);
	}

	async function syncSubBarCodexEnv(ctx: Pick<ExtensionContext, "model" | "modelRegistry">): Promise<boolean> {
		const generation = ++syncGeneration;
		const provider = ctx.model?.provider;
		if (!isAliasProvider(provider)) return restoreSubBarEnv();

		// Resolve through Pi first so an expired OAuth credential is refreshed before
		// exposing it to sub-bar. Then read the stored credential for its account ID.
		const access = await ctx.modelRegistry.getApiKeyForProvider(provider);
		if (generation !== syncGeneration) return false;
		if (!access) return restoreSubBarEnv();

		const stored = readStoredCredential(provider) as StoredCredential | undefined;
		process.env[OAUTH_TOKEN_ENV] = access;
		process.env[ACCESS_TOKEN_ENV] = access;
		if (stored?.type === "oauth" && typeof stored.accountId === "string") {
			process.env[ACCOUNT_ID_ENV] = stored.accountId;
			process.env[CHATGPT_ACCOUNT_ID_ENV] = stored.accountId;
		} else {
			delete process.env[ACCOUNT_ID_ENV];
			delete process.env[CHATGPT_ACCOUNT_ID_ENV];
		}
		ownsSubBarEnv = true;
		return true;
	}

	function scheduleSubBarRefresh(delay: number) {
		const timer = setTimeout(() => {
			refreshTimers.delete(timer);
			pi.events.emit("sub-core:action", { type: "refresh", force: true });
		}, delay);
		refreshTimers.add(timer);
	}

	function refreshSubBarSoon() {
		scheduleSubBarRefresh(0);
		scheduleSubBarRefresh(250);
	}

	async function syncAndRefresh(_event: unknown, ctx: Pick<ExtensionContext, "model" | "modelRegistry">) {
		if (await syncSubBarCodexEnv(ctx)) refreshSubBarSoon();
	}

	async function syncOnly(_event: unknown, ctx: Pick<ExtensionContext, "model" | "modelRegistry">) {
		await syncSubBarCodexEnv(ctx);
	}

	pi.on("session_start", syncAndRefresh);
	pi.on("model_select", syncAndRefresh);
	pi.on("before_agent_start", syncOnly);
	pi.on("turn_start", syncOnly);

	pi.on("session_shutdown", async () => {
		++syncGeneration;
		for (const timer of refreshTimers) clearTimeout(timer);
		refreshTimers.clear();
		restoreSubBarEnv();
	});
}

function getBuiltinCodexProvider(): Provider<typeof CODEX_API> {
	const provider = builtinProviders().find((candidate) => candidate.id === BUILTIN_CODEX_PROVIDER);
	if (!provider) throw new Error("Pi's built-in OpenAI Codex provider is unavailable.");
	if (!provider.auth.oauth) throw new Error("Pi's built-in OpenAI Codex provider does not expose OAuth.");
	return provider as Provider<typeof CODEX_API>;
}

export function createCodexAliasProvider(
	codexProvider: Provider<typeof CODEX_API>,
	providerId: string,
	displayName: string,
): Provider<typeof CODEX_API> {
	const models = codexProvider.getModels().map((model) => ({
		...model,
		provider: providerId,
		name: `${model.name ?? model.id} (${displayName})`,
	}));

	return {
		id: providerId,
		name: displayName,
		baseUrl: codexProvider.baseUrl,
		headers: codexProvider.headers,
		auth: createAliasAuth(codexProvider.auth, displayName),
		getModels: () => models,
		stream(model, context, options) {
			const inner = codexProvider.stream(
				toBuiltinCodexModel(model),
				normalizeCodexAliasContext(context, providerId),
				options,
			);
			return rewriteStreamProvider(inner, providerId);
		},
		streamSimple(model, context, options) {
			const inner = codexProvider.streamSimple(
				toBuiltinCodexModel(model),
				normalizeCodexAliasContext(context, providerId),
				options,
			);
			return rewriteStreamProvider(inner, providerId);
		},
	};
}

function createAliasAuth(auth: ProviderAuth, displayName: string): ProviderAuth {
	const oauth = auth.oauth;
	if (!oauth) throw new Error("Pi's built-in OpenAI Codex provider does not expose OAuth.");

	const aliasOAuth: OAuthAuth = {
		...oauth,
		name: displayName,
	};

	return {
		...auth,
		oauth: aliasOAuth,
	};
}

function toBuiltinCodexModel(model: Model<typeof CODEX_API>): Model<typeof CODEX_API> {
	return {
		...model,
		provider: BUILTIN_CODEX_PROVIDER,
		api: CODEX_API,
	};
}

function normalizeCodexAliasContext(context: Context, providerId: string): Context {
	return {
		...context,
		messages: context.messages.map((message) => {
			if (message.role !== "assistant" || message.provider !== providerId) return message;
			return {
				...message,
				provider: BUILTIN_CODEX_PROVIDER,
			};
		}),
	};
}

function rewriteStreamProvider(
	inner: AssistantMessageEventStream,
	providerId: string,
): AssistantMessageEventStream {
	const outer = createAssistantMessageEventStream();

	void (async () => {
		for await (const event of inner) {
			outer.push(rewriteEventProvider(event, providerId));
		}
	})();

	return outer;
}

function rewriteEventProvider(event: AssistantMessageEvent, providerId: string): AssistantMessageEvent {
	const eventWithMessages = event as AssistantMessageEvent & {
		partial?: AssistantMessage;
		message?: AssistantMessage;
		error?: AssistantMessage;
	};

	if (eventWithMessages.partial) eventWithMessages.partial.provider = providerId;
	if (eventWithMessages.message) eventWithMessages.message.provider = providerId;
	if (eventWithMessages.error) eventWithMessages.error.provider = providerId;

	return event;
}