import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

/**
 * ai& Provider Extension for pi
 *
 * Registers ai& (https://aiand.com) as an OpenAI-compatible inference provider.
 *
 * Authentication:
 *   Run `pi --login aiand` and enter your API key when prompted.
 *   Or pass `--api-key sk-...` on the CLI.
 *
 * Usage:
 *   pi --model deepseek-ai/deepseek-v4-flash "your prompt"
 *
 * This extension uses a static list for fast startup and offline availability.
 * The list is synced from the models available to the maintainer's ai&
 * organization; availability and pricing can vary by organization.
 */

async function loginAiand(callbacks: any) {
	const key = await callbacks.onPrompt({ message: "Enter your ai& API key:" });
	return {
		refresh: key,
		access: key,
		expires: Date.now() + 100 * 365 * 24 * 60 * 60 * 1000, // 100 years
	};
}

async function refreshAiand(credentials: any) {
	return credentials; // API keys don't expire
}

function getApiKey(credentials: any) {
	return credentials.access;
}

export default function (pi: ExtensionAPI) {
	pi.registerProvider("aiand", {
		name: "aiand",
		baseUrl: "https://api.aiand.com/v1",
		api: "openai-completions",
		apiKey: "$AIAND_API_KEY",
		authHeader: true,
		oauth: {
			name: "aiand",
			login: loginAiand,
			refreshToken: refreshAiand,
			getApiKey,
		},
		// pi's model metadata has no billing-currency field. These cost values are
		// USD-equivalents of the synced JPY rates at ¥160/USD; ai& bills in JPY.
		models: [
			{
				id: "qwen/qwen3.6-27b",
				name: "Qwen 3.6 27B",
				reasoning: true,
				input: ["text"],
				contextWindow: 262_144,
				maxTokens: 16_384,
				cost: { input: 0.05, output: 0.51, cacheRead: 0.03, cacheWrite: 0 },
			},
			{
				id: "deepseek-ai/deepseek-v4-flash",
				name: "DeepSeek V4 Flash",
				reasoning: true,
				input: ["text"],
				contextWindow: 1_000_000,
				maxTokens: 16_384,
				cost: { input: 0.025, output: 0.04, cacheRead: 0.01, cacheWrite: 0 },
			},
			{
				id: "google/gemma-4-31b-it",
				name: "Gemma 4 31B IT",
				reasoning: true,
				input: ["text", "image"],
				contextWindow: 262_144,
				maxTokens: 16_384,
				cost: { input: 0.03, output: 0.08, cacheRead: 0.008, cacheWrite: 0 },
			},
			{
				id: "openai/gpt-oss-120b",
				name: "GPT-OSS 120B",
				reasoning: true,
				input: ["text"],
				contextWindow: 131_072,
				maxTokens: 16_384,
				cost: { input: 0.025, output: 0.095, cacheRead: 0.01, cacheWrite: 0 },
			},
			{
				id: "deepseek-ai/deepseek-v4-pro",
				name: "DeepSeek V4 Pro",
				reasoning: true,
				input: ["text"],
				contextWindow: 1_000_000,
				maxTokens: 16_384,
				cost: { input: 0.16, output: 0.40, cacheRead: 0.03, cacheWrite: 0 },
			},
			{
				id: "moonshotai/kimi-k2.7-code",
				name: "Kimi K2.7 Code",
				reasoning: true,
				input: ["text", "image"],
				contextWindow: 262_144,
				maxTokens: 16_384,
				cost: { input: 0.125, output: 0.56, cacheRead: 0.03, cacheWrite: 0 },
			},
			{
				id: "moonshotai/kimi-k3",
				name: "Kimi K3",
				reasoning: true,
				input: ["text", "image"],
				contextWindow: 1_048_576,
				maxTokens: 16_384,
				cost: { input: 0.48, output: 2.00, cacheRead: 0.08, cacheWrite: 0 },
			},
			{
				id: "zai-org/glm-5.2",
				name: "GLM 5.2",
				reasoning: true,
				input: ["text"],
				contextWindow: 1_000_000,
				maxTokens: 16_384,
				cost: { input: 0.16, output: 0.65, cacheRead: 0.04, cacheWrite: 0 },
			},
		],
	});
}