import { z } from 'zod';

const MODELS_DEV_URL = 'https://models.dev/api.json';

const MODELS_DEV_PROVIDER_ALIASES: Record<string, string> = {
	'amazon-bedrock': 'aws-bedrock',
	azure: 'azure-openai',
	'azure-cognitive-services': 'azure-openai',
};

const AGENT_PROVIDER_NAMES: Record<string, string> = {
	'aws-bedrock': 'AWS Bedrock',
	'azure-openai': 'Azure OpenAI',
	'google-vertex-anthropic': 'Google Vertex Anthropic',
};

/** Cost per million tokens. */
export interface ModelCost {
	/** Cost per million input tokens (USD). */
	input: number;
	/** Cost per million output tokens (USD). */
	output: number;
	/** Cost per million cached input tokens (USD). */
	cacheRead?: number;
	/** Cost per million cache write tokens (USD). */
	cacheWrite?: number;
}

/** Model context/output limits. */
export interface ModelLimits {
	/** Maximum context window size in tokens. */
	context?: number;
	/** Maximum output tokens. */
	output?: number;
}

/** Input and output types supported by a model. */
export interface ModelModalities {
	/** Supported input types. */
	input?: string[];
	/** Supported output types. */
	output?: string[];
}

/** Information about a single model. */
export interface ModelInfo {
	/** Model ID (e.g. 'claude-sonnet-4-5'). */
	id: string;
	/** Human-readable name (e.g. 'Claude Sonnet 4.5'). */
	name: string;
	/** Release date in ISO date format when available from models.dev. */
	releaseDate?: string;
	/** Whether the model supports reasoning / thinking, when reported by models.dev. */
	reasoning?: boolean;
	/**
	 * Whether the model accepts a temperature sampling parameter, when
	 * reported by models.dev. Reasoning-family models typically report false.
	 */
	temperature?: boolean;
	/** Whether the model supports tool calling. */
	toolCall: boolean;
	/** Input and output types supported by the model. */
	modalities?: ModelModalities;
	/** Cost per million tokens. */
	cost?: ModelCost;
	/** Token limits. */
	limits?: ModelLimits;
}

/** Information about a provider. */
export interface ProviderInfo {
	/** Provider ID (e.g. 'anthropic'). */
	id: string;
	/** Human-readable name (e.g. 'Anthropic'). */
	name: string;
	/** Available (non-deprecated) models keyed by model ID. */
	models: Record<string, ModelInfo>;
	/**
	 * Model IDs models.dev marks as deprecated for this provider. Kept out of
	 * `models` so existing consumers never offer them, but exposed so callers
	 * can distinguish known-retired IDs (hard error) from merely-absent ones.
	 */
	deprecatedModelIds?: string[];
}

/** The full catalog of providers and their models. */
export type ProviderCatalog = Record<string, ProviderInfo>;

const modelsDevModelSchema = z.object({
	id: z.string().min(1),
	name: z.string().min(1),
	release_date: z.string().optional(),
	reasoning: z.boolean().optional(),
	temperature: z.boolean().optional(),
	tool_call: z.boolean().optional(),
	status: z.string().optional(),
	modalities: z
		.object({
			input: z.array(z.string()).optional(),
			output: z.array(z.string()).optional(),
		})
		.optional(),
	cost: z
		.object({
			input: z.number().optional(),
			output: z.number().optional(),
			cache_read: z.number().optional(),
			cache_write: z.number().optional(),
		})
		.optional(),
	limit: z
		.object({
			context: z.number().optional(),
			output: z.number().optional(),
		})
		.optional(),
});

const modelsDevProviderSchema = z.object({
	id: z.string().min(1),
	name: z.string().min(1),
	models: z.record(z.string(), z.unknown()).optional(),
});

const modelsDevCatalogSchema = z.record(z.string(), z.unknown());

function toAgentProviderId(modelsDevProviderId: string): string {
	return MODELS_DEV_PROVIDER_ALIASES[modelsDevProviderId] ?? modelsDevProviderId;
}

const LATEST_NAME_SUFFIX = /\s*\(latest\)$/i;

/**
 * models.dev names versionless alias models with a " (latest)" suffix
 * (e.g. `claude-opus-4-5` → "Claude Opus 4.5 (latest)"), which quickly goes
 * stale as newer models ship. Strip the suffix from every model name, and when
 * a pinned snapshot shares the resulting name with an alias (e.g.
 * `claude-opus-4-5-20251101` "Claude Opus 4.5"), drop the snapshot so the
 * alias is the single entry for that model.
 */
function normalizeLatestModelNames(models: Record<string, ModelInfo>): void {
	const aliasNames = new Set<string>();
	for (const model of Object.values(models)) {
		if (LATEST_NAME_SUFFIX.test(model.name)) {
			aliasNames.add(model.name.replace(LATEST_NAME_SUFFIX, ''));
		}
	}

	for (const [modelId, model] of Object.entries(models)) {
		if (LATEST_NAME_SUFFIX.test(model.name)) {
			model.name = model.name.replace(LATEST_NAME_SUFFIX, '');
		} else if (aliasNames.has(model.name)) {
			delete models[modelId];
		}
	}
}

/**
 * Fetch the provider/model catalog from models.dev.
 *
 * Returns a map of provider ID → ProviderInfo with all available models.
 * The catalog is fetched once and can be cached by the caller.
 *
 * @example
 * ```typescript
 * import { fetchProviderCatalog } from '@n8n/agents';
 *
 * const catalog = await fetchProviderCatalog();
 * console.log(Object.keys(catalog)); // ['anthropic', 'openai', ...]
 * console.log(catalog.anthropic.models['claude-sonnet-4-5'].reasoning); // true
 * ```
 */
export async function fetchProviderCatalog(options?: {
	signal?: AbortSignal;
}): Promise<ProviderCatalog> {
	const response = await fetch(MODELS_DEV_URL, { signal: options?.signal });
	if (!response.ok) {
		throw new Error(`Failed to fetch provider catalog: ${response.statusText}`);
	}

	const data = modelsDevCatalogSchema.parse(await response.json());
	const catalog: ProviderCatalog = {};

	for (const [key, rawProvider] of Object.entries(data)) {
		const providerResult = modelsDevProviderSchema.safeParse(rawProvider);
		if (!providerResult.success) continue;
		const provider = providerResult.data;
		if (!provider.models || Object.keys(provider.models).length === 0) continue;

		const providerId = toAgentProviderId(key);
		const models: Record<string, ModelInfo> = {};
		const deprecatedModelIds = new Set<string>(catalog[providerId]?.deprecatedModelIds);
		for (const [modelId, rawModel] of Object.entries(provider.models)) {
			const modelResult = modelsDevModelSchema.safeParse(rawModel);
			if (!modelResult.success) continue;
			const model = modelResult.data;
			// Deprecated models still 404 at call time when the provider retires
			// them, so never offer them in `models` — but record the id so
			// validation can flag known-retired picks as hard errors.
			if (model.status === 'deprecated') {
				deprecatedModelIds.add(modelId);
				deprecatedModelIds.add(model.id);
				continue;
			}
			const info: ModelInfo = {
				id: model.id,
				name: model.name,
				...(model.release_date !== undefined && { releaseDate: model.release_date }),
				...(model.reasoning !== undefined && { reasoning: model.reasoning }),
				...(model.temperature !== undefined && { temperature: model.temperature }),
				toolCall: model.tool_call ?? false,
				...(model.modalities !== undefined && { modalities: model.modalities }),
			};
			if (model.cost?.input !== undefined && model.cost?.output !== undefined) {
				info.cost = {
					input: model.cost.input,
					output: model.cost.output,
					...(model.cost.cache_read !== undefined && { cacheRead: model.cost.cache_read }),
					...(model.cost.cache_write !== undefined && { cacheWrite: model.cost.cache_write }),
				};
			}
			if (model.limit) {
				info.limits = {
					...(model.limit.context !== undefined && { context: model.limit.context }),
					...(model.limit.output !== undefined && { output: model.limit.output }),
				};
			}
			models[modelId] = info;
		}

		if (Object.keys(models).length === 0 && deprecatedModelIds.size === 0) continue;

		catalog[providerId] = {
			id: providerId,
			name: catalog[providerId]?.name ?? AGENT_PROVIDER_NAMES[providerId] ?? provider.name,
			models: {
				...(catalog[providerId]?.models ?? {}),
				...models,
			},
			deprecatedModelIds: Array.from(deprecatedModelIds),
		};
	}

	for (const provider of Object.values(catalog)) {
		normalizeLatestModelNames(provider.models);
	}

	return catalog;
}

// --- Global cached catalog for internal use ---

let cachedCatalog: ProviderCatalog | undefined;
let catalogFetchPromise: Promise<ProviderCatalog | undefined> | undefined;

/**
 * Get the cached catalog, fetching once if needed.
 * Returns undefined if the fetch fails (offline, timeout, etc.).
 * On failure, clears the in-flight promise so the next call retries.
 * @internal
 */
export async function getCachedCatalog(): Promise<ProviderCatalog | undefined> {
	if (cachedCatalog) return cachedCatalog;

	catalogFetchPromise ??= fetchProviderCatalog()
		.then((c) => {
			cachedCatalog = c;
			return c;
		})
		.catch((error: unknown) => {
			// Clear so subsequent calls retry
			catalogFetchPromise = undefined;
			console.warn(
				'[agents] Failed to fetch model catalog from models.dev — cost data will be unavailable:',
				error instanceof Error ? error.message : error,
			);
			return undefined;
		});

	return await catalogFetchPromise;
}

/**
 * Look up cost info for a model by its full ID (e.g. 'anthropic/claude-sonnet-4-5').
 * Returns undefined if catalog is unavailable or model not found.
 * @internal
 */
export async function getModelCost(modelId: string): Promise<ModelCost | undefined> {
	const catalog = await getCachedCatalog();
	if (!catalog) return undefined;

	const [provider, ...rest] = modelId.split('/');
	const modelName = rest.join('/');

	return catalog[provider]?.models[modelName]?.cost;
}

/**
 * models.dev's `cacheWrite` rate encodes Anthropic's 5-minute-TTL write premium
 * (1.25x base input). A 1h breakpoint costs ~1.6x more to write (2x vs 1.25x base
 * input), so scale the catalog rate when the caller reports a 1h TTL. Only
 * Anthropic populates `inputTokenDetails.cacheWrite`, so this is safe to apply
 * unconditionally — it's a no-op whenever there are no cache-write tokens.
 */
function resolveCacheWriteRate(
	cost: ModelCost,
	anthropicCacheTtl: '5m' | '1h' | undefined,
): number {
	const isOneHour = anthropicCacheTtl === '1h';
	return cost.cacheWrite !== undefined
		? cost.cacheWrite * (isOneHour ? 1.6 : 1)
		: cost.input * (isOneHour ? 2 : 1.25);
}

/**
 * Compute the cost in USD from token usage and per-million-token pricing.
 * When `usage.inputTokenDetails` is present, prompt tokens are billed per
 * cache tier (no-cache / cache read / cache write) using the catalog's cache
 * rates. When a tier's catalog rate is unavailable, cache reads fall back to
 * the flat input rate; cache writes fall back to the input rate scaled by
 * Anthropic's write premium (see {@link resolveCacheWriteRate}).
 * `anthropicCacheTtl` scales the cache-write rate for Anthropic's 1h tier;
 * ignored when there are no cache-write tokens.
 */
export function computeCost(
	usage: {
		promptTokens: number;
		completionTokens: number;
		inputTokenDetails?: { noCache?: number; cacheRead?: number; cacheWrite?: number };
	},
	cost: ModelCost,
	options?: { anthropicCacheTtl?: '5m' | '1h' },
): number {
	const details = usage.inputTokenDetails;
	let inputCost: number;
	if (details) {
		const noCache = details.noCache ?? 0;
		const cacheRead = details.cacheRead ?? 0;
		const cacheWrite = details.cacheWrite ?? 0;
		// Any prompt tokens not covered by the breakdown are billed at the full input rate.
		const remaining = Math.max(usage.promptTokens - (noCache + cacheRead + cacheWrite), 0);
		inputCost =
			((noCache + remaining) / 1_000_000) * cost.input +
			(cacheRead / 1_000_000) * (cost.cacheRead ?? cost.input) +
			(cacheWrite / 1_000_000) * resolveCacheWriteRate(cost, options?.anthropicCacheTtl);
	} else {
		inputCost = (usage.promptTokens / 1_000_000) * cost.input;
	}

	const outputCost = (usage.completionTokens / 1_000_000) * cost.output;
	return inputCost + outputCost;
}