/**
 * Antigravity model catalog: static fallback list plus per-model wire profiles
 * and thinking metadata, copied from oh-my-pi's google-antigravity catalog.
 *
 * The extension registers wire ids directly (there is no effort-tier routing
 * layer in Prime's extension API), so each entry below is a concrete Cloud
 * Code Assist wire model id. Live discovery (discovery.ts) replaces the static
 * list when it succeeds.
 */

import type { ThinkingLevel } from "@earendil-works/pi-ai";
import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
import { getAntigravityModelWireProfile } from "./wire.js";

export type AntigravityThinkingMode = "budget" | "google-level";

export interface AntigravityModelThinking {
	/** `thinkingLevel` (google-level) vs `thinkingBudget` (budget) wire field. */
	mode: AntigravityThinkingMode;
	/** Per-effort token budgets for budget-mode models. Falls back to Google defaults. */
	effortBudgets?: Partial<Record<ThinkingLevel, number>>;
	/** Emit an explicit thinking-off config; otherwise omit thinkingConfig entirely. */
	suppressWhenOff: boolean;
	/** Reasoning is mandatory; "off" is hidden in the UI unless suppressWhenOff. */
	requiresEffort: boolean;
	/** `-low` SKUs that also serve "minimal" must emit LOW (CCA rejects MINIMAL there). */
	minimalLevelIsLow?: boolean;
}

export interface AntigravityModelProfile {
	/** Claude vs Gemini/GPT-OSS wire paths. */
	identityClass: "anthropic" | "gemini";
	/** Upstream wire model id (equals the registered id for this extension). */
	wireModelId: string;
	/** labels.model_enum token captured from the real antigravity/hub client. */
	modelEnum?: string;
	/** Fixed generationConfig.maxOutputTokens enforced by the backend. */
	maxOutputTokens?: number;
	thinking: AntigravityModelThinking;
}

export function identityClassFor(modelId: string = ""): AntigravityModelProfile["identityClass"] {
	return (modelId || "").startsWith("claude-") ? "anthropic" : "gemini";
}

export interface AntigravityModelCost {
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
}

export const ANTIGRAVITY_MODEL_COSTS = {
	"claude-opus": { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 },
	"claude-sonnet": { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
	"claude-haiku": { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
	"gemini-flash": { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 },
	"gemini-flash-lite": { input: 0.25, output: 1.5, cacheRead: 0.025, cacheWrite: 0 },
	"gemini-pro": { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 },
	"gpt-oss": { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
	"gemini-2.5-flash": { input: 0.3, output: 2.5, cacheRead: 0.03, cacheWrite: 0 },
	"gemini-2.5-flash-lite": { input: 0.1, output: 0.4, cacheRead: 0.01, cacheWrite: 0 },
	"gemini-2.5-pro": { input: 1.25, output: 10, cacheRead: 0.125, cacheWrite: 0 },
	"gemini-3.5-flash-extra-low": { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 },
	"gemini-3.5-flash-low": { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 },
	"gemini-3-flash-agent": { input: 0.5, output: 3, cacheRead: 0.05, cacheWrite: 0 },
} as const satisfies Record<string, AntigravityModelCost>;

const ZERO_COST: AntigravityModelCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };

export function resolveAntigravityModelCost(modelId: string = ""): AntigravityModelCost {
	if (modelId in ANTIGRAVITY_MODEL_COSTS) {
		return ANTIGRAVITY_MODEL_COSTS[modelId as keyof typeof ANTIGRAVITY_MODEL_COSTS];
	}
	const lower = (modelId || "").toLowerCase();
	if (lower.startsWith("claude-") && lower.includes("opus")) return ANTIGRAVITY_MODEL_COSTS["claude-opus"];
	if (lower.startsWith("claude-") && lower.includes("sonnet")) return ANTIGRAVITY_MODEL_COSTS["claude-sonnet"];
	if (lower.startsWith("claude-") && lower.includes("haiku")) return ANTIGRAVITY_MODEL_COSTS["claude-haiku"];
	if (lower.startsWith("gpt-oss")) return ANTIGRAVITY_MODEL_COSTS["gpt-oss"];
	if (lower.startsWith("gemini-") && lower.includes("flash-lite")) return ANTIGRAVITY_MODEL_COSTS["gemini-flash-lite"];
	if (lower.startsWith("gemini-") && (lower.includes("flash") || lower.includes("flash-agent"))) {
		return ANTIGRAVITY_MODEL_COSTS["gemini-flash"];
	}
	if (lower.startsWith("gemini-") && lower.includes("pro")) return ANTIGRAVITY_MODEL_COSTS["gemini-pro"];
	return ZERO_COST;
}

interface AntigravityModelDef {
	id: string;
	name: string;
	reasoning: boolean;
	input: ("text" | "image")[];
	contextWindow: number;
	maxTokens: number;
	cost: AntigravityModelCost;
	thinking: AntigravityModelThinking;
}

const THINK_BUDGET_DEFAULT: AntigravityModelThinking = {
	mode: "budget",
	suppressWhenOff: false,
	requiresEffort: false,
};

const GEMINI_FLASH_BUDGETS: Partial<Record<ThinkingLevel, number>> = {
	minimal: 1000,
	low: 1000,
	medium: 4000,
	high: 10000,
};

const GEMINI_PRO_BUDGETS: Partial<Record<ThinkingLevel, number>> = {
	low: 1001,
	high: 10001,
};

const MODEL_DEFS: AntigravityModelDef[] = [
	{
		id: "claude-opus-4-5-thinking",
		name: "Claude Opus 4.5",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 200000,
		maxTokens: 64000,
		cost: resolveAntigravityModelCost("claude-opus-4-5-thinking"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "claude-opus-4-6-thinking",
		name: "Claude Opus 4.6",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 250000,
		maxTokens: 64000,
		cost: resolveAntigravityModelCost("claude-opus-4-6-thinking"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "claude-sonnet-4-5",
		name: "Claude Sonnet 4.5",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1000000,
		maxTokens: 64000,
		cost: resolveAntigravityModelCost("claude-sonnet-4-5"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "claude-sonnet-4-6",
		name: "Claude Sonnet 4.6",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 250000,
		maxTokens: 64000,
		cost: resolveAntigravityModelCost("claude-sonnet-4-6"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "gemini-2.5-flash",
		name: "Gemini 2.5 Flash",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-2.5-flash"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "gemini-2.5-flash-lite",
		name: "Gemini 2.5 Flash Lite",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-2.5-flash-lite"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
	{
		id: "gemini-2.5-pro",
		name: "Gemini 2.5 Pro",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-2.5-pro"),
		thinking: { ...THINK_BUDGET_DEFAULT, requiresEffort: true },
	},
	{
		id: "gemini-3-pro-low",
		name: "Gemini 3 Pro (Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-3-pro-low"),
		thinking: { mode: "google-level", suppressWhenOff: true, requiresEffort: true },
	},
	{
		id: "gemini-3.1-flash-lite",
		name: "Gemini 3.1 Flash Lite",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-3.1-flash-lite"),
		thinking: { mode: "google-level", suppressWhenOff: false, requiresEffort: true },
	},
	{
		id: "gemini-3.1-pro-low",
		name: "Gemini 3.1 Pro (Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-3.1-pro-low"),
		thinking: {
			mode: "budget",
			effortBudgets: GEMINI_PRO_BUDGETS,
			suppressWhenOff: true,
			requiresEffort: true,
		},
	},
	{
		id: "gemini-pro-agent",
		name: "Gemini Pro Agent",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65535,
		cost: resolveAntigravityModelCost("gemini-pro-agent"),
		thinking: {
			mode: "budget",
			effortBudgets: GEMINI_PRO_BUDGETS,
			suppressWhenOff: true,
			requiresEffort: true,
		},
	},
	{
		id: "gemini-3.5-flash-extra-low",
		name: "Gemini 3.5 Flash (Extra Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-3.5-flash-extra-low"),
		thinking: {
			mode: "budget",
			effortBudgets: GEMINI_FLASH_BUDGETS,
			suppressWhenOff: true,
			requiresEffort: true,
		},
	},
	{
		id: "gemini-3.5-flash-low",
		name: "Gemini 3.5 Flash (Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-3.5-flash-low"),
		thinking: {
			mode: "budget",
			effortBudgets: GEMINI_FLASH_BUDGETS,
			suppressWhenOff: true,
			requiresEffort: true,
		},
	},
	{
		id: "gemini-3-flash-agent",
		name: "Gemini 3 Flash Agent",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-3-flash-agent"),
		thinking: {
			mode: "budget",
			effortBudgets: GEMINI_FLASH_BUDGETS,
			suppressWhenOff: true,
			requiresEffort: true,
		},
	},
	{
		id: "gemini-3.6-flash-low",
		name: "Gemini 3.6 Flash (Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-3.6-flash-low"),
		thinking: {
			mode: "google-level",
			suppressWhenOff: false,
			requiresEffort: true,
			minimalLevelIsLow: true,
		},
	},
	{
		id: "gemini-3.7-flash-low",
		name: "Gemini 3.7 Flash (Low)",
		reasoning: true,
		input: ["text", "image"],
		contextWindow: 1048576,
		maxTokens: 65536,
		cost: resolveAntigravityModelCost("gemini-3.7-flash-low"),
		thinking: {
			mode: "google-level",
			suppressWhenOff: false,
			requiresEffort: true,
			minimalLevelIsLow: true,
		},
	},
	{
		id: "gpt-oss-120b-medium",
		name: "GPT OSS 120B",
		reasoning: true,
		input: ["text"],
		contextWindow: 131072,
		maxTokens: 32768,
		cost: resolveAntigravityModelCost("gpt-oss-120b-medium"),
		thinking: { ...THINK_BUDGET_DEFAULT },
	},
];

const MODEL_DEF_BY_ID = new Map(MODEL_DEFS.map((model) => [model.id, model]));

/**
 * Thinking-mode inference for live-discovered ids not in the static table,
 * mirroring oh-my-pi's per-family google-antigravity compat rules.
 */
function inferAntigravityThinking(modelId: string = ""): AntigravityModelThinking {
	const lower = (modelId || "").toLowerCase();
	if (lower.startsWith("claude-") || lower.startsWith("gpt-oss")) return { ...THINK_BUDGET_DEFAULT };
	if (lower.startsWith("gemini-3-pro")) {
		return { mode: "google-level", suppressWhenOff: true, requiresEffort: true };
	}
	if (/^gemini-3\.(6|7)-flash/.test(lower)) {
		return {
			mode: "google-level",
			suppressWhenOff: false,
			requiresEffort: true,
			minimalLevelIsLow: lower.endsWith("-low"),
		};
	}
	if (lower === "gemini-3.1-flash-lite") {
		return { mode: "google-level", suppressWhenOff: false, requiresEffort: true };
	}
	if (/^gemini-3(?:\.\d+)?-flash/.test(lower) || lower.includes("flash-agent")) {
		return { mode: "budget", effortBudgets: GEMINI_FLASH_BUDGETS, suppressWhenOff: true, requiresEffort: true };
	}
	if (/^gemini-3(?:\.\d+)?-pro/.test(lower) || lower === "gemini-pro-agent") {
		return { mode: "budget", effortBudgets: GEMINI_PRO_BUDGETS, suppressWhenOff: true, requiresEffort: true };
	}
	return { ...THINK_BUDGET_DEFAULT };
}

export function getAntigravityProfile(modelId: string): AntigravityModelProfile {
	const def = MODEL_DEF_BY_ID.get(modelId);
	const wire = getAntigravityModelWireProfile(modelId);
	return {
		identityClass: def ? identityClassFor(def.id) : identityClassFor(modelId),
		wireModelId: modelId,
		modelEnum: wire?.modelEnum,
		maxOutputTokens: wire?.maxOutputTokens,
		thinking: def?.thinking ?? inferAntigravityThinking(modelId),
	};
}

/**
 * Prime UI availability map. Levels omitted from the map default to supported
 * (except xhigh/max, which require an explicit mapping), and `null` hides a
 * level. Mirrors oh-my-pi's per-model supported efforts.
 */
function thinkingLevelMap(def: AntigravityModelDef): ProviderModelConfig["thinkingLevelMap"] {
	const { thinking } = def;
	const levelMap: NonNullable<ProviderModelConfig["thinkingLevelMap"]> = {
		xhigh: null,
		max: null,
	};

	if (thinking.requiresEffort && !thinking.suppressWhenOff) {
		levelMap.off = null;
	}

	if (thinking.mode === "google-level") {
		levelMap.minimal = thinking.minimalLevelIsLow ? "LOW" : "MINIMAL";
		levelMap.low = "LOW";
		levelMap.medium = "MEDIUM";
		levelMap.high = "HIGH";
	}

	return levelMap;
}

export function toProviderModelConfig(def: AntigravityModelDef): ProviderModelConfig {
	return {
		id: def.id,
		name: def.name,
		reasoning: def.reasoning,
		input: def.input,
		cost: def.cost,
		contextWindow: def.contextWindow,
		maxTokens: def.maxTokens,
		thinkingLevelMap: thinkingLevelMap(def),
	};
}

export const STATIC_PROVIDER_MODELS: ProviderModelConfig[] = MODEL_DEFS.map(toProviderModelConfig);