import { getModels, type Model } from "@mariozechner/pi-ai";
import { openaiCodexOAuthProvider } from "@mariozechner/pi-ai/oauth";
import type { ExtensionAPI, ProviderModelConfig } from "@mariozechner/pi-coding-agent";

const PROVIDER = "openai-codex" as const;
const BASE_MODEL_ID = "gpt-5.4" as const;
const TARGET_MODEL_ID = "gpt-5.5" as const;
const TARGET_CONTEXT_WINDOW = 1_000_000;

function cloneCompat<T extends Model<any>["compat"]>(compat: T): T {
	if (!compat) return compat;
	return JSON.parse(JSON.stringify(compat)) as T;
}

function toProviderModelConfig(model: Model<any>): ProviderModelConfig {
	return {
		id: model.id,
		name: model.name,
		api: model.api,
		reasoning: model.reasoning,
		input: [...model.input],
		cost: { ...model.cost },
		contextWindow: model.contextWindow,
		maxTokens: model.maxTokens,
		headers: model.headers ? { ...model.headers } : undefined,
		compat: cloneCompat(model.compat),
	};
}

export default function (pi: ExtensionAPI) {
	const builtInModels = getModels(PROVIDER);
	const baseModel = builtInModels.find((model) => model.id === BASE_MODEL_ID);

	if (!baseModel) {
		throw new Error(`Could not find ${PROVIDER}/${BASE_MODEL_ID} to derive ${TARGET_MODEL_ID}.`);
	}

	const targetModel: ProviderModelConfig = {
		...toProviderModelConfig(baseModel),
		id: TARGET_MODEL_ID,
		name: "GPT-5.5",
		cost: {
			input: baseModel.cost.input * 2,
			output: baseModel.cost.output * 2,
			cacheRead: baseModel.cost.cacheRead * 2,
			cacheWrite: baseModel.cost.cacheWrite * 2,
		},
		contextWindow: TARGET_CONTEXT_WINDOW,
	};

	const models: ProviderModelConfig[] = [];
	let insertedTarget = false;

	for (const model of builtInModels) {
		if (model.id === TARGET_MODEL_ID) {
			models.push(targetModel);
			insertedTarget = true;
			continue;
		}

		models.push(toProviderModelConfig(model));
		if (model.id === BASE_MODEL_ID && !insertedTarget) {
			models.push(targetModel);
			insertedTarget = true;
		}
	}

	if (!insertedTarget) {
		models.push(targetModel);
	}

	pi.registerProvider(PROVIDER, {
		baseUrl: baseModel.baseUrl,
		api: baseModel.api,
		models,
		oauth: {
			name: openaiCodexOAuthProvider.name,
			login: openaiCodexOAuthProvider.login,
			refreshToken: openaiCodexOAuthProvider.refreshToken,
			getApiKey: openaiCodexOAuthProvider.getApiKey,
			modifyModels: openaiCodexOAuthProvider.modifyModels,
		},
	});
}