/**
 * MiMo provider package extracted from pi-excalibur.
 *
 * Registers the MiMo models pi-excalibur uses.
 * The provider name stays `mimo` so existing literal model references keep working
 * while the transport details move out of feature code.
 */

import {
	type Api,
	type AssistantMessageEventStream,
	type Context,
	type Model,
	type SimpleStreamOptions,
	streamSimpleOpenAICompletions,
} from "@mariozechner/pi-ai";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

const BASE_URL = "https://token-plan-cn.xiaomimimo.com/v1";
const MIMO_API = "mimo-openai-completions" as Api;

const MIMO_COMPAT = {
	supportsStore: false,
	supportsDeveloperRole: true,
	supportsUsageInStreaming: true,
	maxTokensField: "max_completion_tokens" as const,
};

function buildMimoModel(model: Model<Api>): Model<"openai-completions"> {
	return {
		...model,
		api: "openai-completions" as Api,
		baseUrl: BASE_URL,
		compat: { ...MIMO_COMPAT, ...(model.compat ?? {}) },
	} as Model<"openai-completions">;
}

function streamMiMo(
	model: Model<Api>,
	context: Context,
	options?: SimpleStreamOptions,
): AssistantMessageEventStream {
	const mimoModel = buildMimoModel(model);
	const originalOnPayload = options?.onPayload;

	const wrappedOptions: SimpleStreamOptions = {
		...options,
		async onPayload(params, payloadModel) {
			delete (params as any).reasoning_effort;
			delete (params as any).enable_thinking;
			delete (params as any).chat_template_kwargs;

			if (model.reasoning) {
				(params as any).thinking = {
					type: options?.reasoning ? "enabled" : "disabled",
				};
			}

			if (originalOnPayload) {
				return await originalOnPayload(params, payloadModel);
			}
			return params;
		},
	};

	return streamSimpleOpenAICompletions(mimoModel, context, wrappedOptions);
}

const MODELS = [
	{
		id: "mimo-v2.5-pro",
		name: "MiMo V2.5 Pro",
		reasoning: true,
		input: ["text"] as ("text" | "image")[],
		cost: { input: 1, output: 3, cacheRead: 0.2, cacheWrite: 0 },
		contextWindow: 1_000_000,
		maxTokens: 131_072,
		compat: MIMO_COMPAT,
	},
	{
		id: "mimo-v2-pro",
		name: "MiMo V2 Pro",
		reasoning: true,
		input: ["text"] as ("text" | "image")[],
		cost: { input: 1, output: 3, cacheRead: 0.2, cacheWrite: 0 },
		contextWindow: 1_000_000,
		maxTokens: 131_072,
		compat: MIMO_COMPAT,
	},
];

export default function (pi: ExtensionAPI) {
	pi.registerProvider("mimo", {
		baseUrl: BASE_URL,
		apiKey: "MIMO_API_KEY",
		api: MIMO_API,
		models: MODELS,
		streamSimple: streamMiMo,
	});
}