// Provider registration: turn scan results into pi models and register
// the provider with pi. Lazy-loaded on first discovery completion.

import {
	DEFAULT_API_KEY,
	PROVIDER_NAME,
	baseName,
	idSafeHost,
	makeCompat,
	modelOptions,
	modelOptionsFor,
	saveConfig,
	saveModelsCache,
	shared,
} from "./core.ts";
import type { ExtensionAPI } from "./types.ts";
import type {
	EndpointResult,
	InfraConfig,
	LlamaCppModel,
	ModelMetadata,
	PiModel,
	ScanResult,
} from "./types.ts";
import { cleanModelName, lmStudioContextLength } from "./scan.ts";
import { createLongTimeoutOpenAICompletionsStream } from "./runtime.ts";

/** Compact badge string appended to display names when enabled. */
function badgeSuffix(
	meta: { vision?: boolean; drafter?: string; routerStatus?: string },
	showBadges: boolean,
): string {
	if (!showBadges) return "";
	let badges = "";
	if (meta.vision) badges += "👁️";
	if (meta.drafter) badges += "🚀";
	if (meta.routerStatus && meta.routerStatus !== "loaded") badges += "💤";
	return badges ? ` [${badges}]` : "";
}

function toPiModel(
	model: LlamaCppModel,
	ep: EndpointResult,
	srv: import("./types.ts").ServerConfig,
	settings: import("./types.ts").SettingsConfig,
	modelMeta: ModelMetadata,
): PiModel {
	const rawId = String(model.id ?? "");
	const hostPort = `${idSafeHost(srv.host)}:${ep.port}`;
	const kind = ep.server;

	// Per-model thinking-budget field override: wins over the per-kind default
	// (e.g. force vLLM's `thinking_token_budget` on a server whose kind was not
	// detected). pi reads it from `compat.thinkingTokenBudgetField` directly.
	const compat = makeCompat(kind, modelOptionsFor(rawId)?.thinkingBudgetField);

	// A server may publish a model id that already embeds this machine's tag —
	// e.g. a vLLM started with `--served-model-name "Example-27B (gpu-host:8000)"`.
	// The tag is only needed to disambiguate across machines, so add it only
	// when the raw id doesn't already carry it. Compare case-insensitively:
	// hostPort is lowercased by idSafeHost, but a served-model-name may use any
	// casing. An empty raw id gets no tag (nothing to disambiguate). Crucially,
	// `serverModelId` must ALWAYS stay the raw server id: it is what gets sent
	// in the request payload, and a display-shaped id (with the tag) makes
	// OpenAI-compatible servers (vLLM/SGLang/TGI) answer 404 "model does not
	// exist".
	const hasTag = rawId.toLowerCase().includes(`(${hostPort})`);
	const machineTag = settings.prefixModelIds && rawId !== "" && !hasTag ? ` (${hostPort})` : "";

	const common = {
		baseUrl: ep.baseUrl,
		input: (modelMeta.vision ? ["text", "image"] : ["text"]) as ("text" | "image")[],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		endpoint: { serverId: srv.id, host: srv.host, port: ep.port, kind, mode: ep.mode },
		quant: modelMeta.quant,
		cacheK: modelMeta.cacheK,
		cacheV: modelMeta.cacheV,
		drafter: modelMeta.drafter,
		routerStatus: modelMeta.routerStatus,
		...(srv.apiKey ? { headers: { Authorization: `Bearer ${srv.apiKey}` } } : {}),
	};

	// LM Studio: OpenAI-compatible runtime with rich REST model metadata.
	if (kind === "lmstudio") {
		const info = model.lmStudio;
		const contextWindow =
			lmStudioContextLength(info) ?? model.context_window ?? model.context_length ?? 32768;
		const displayName = info?.display_name ? cleanModelName(info.display_name) : cleanModelName(rawId);
		return {
			...common,
			id: `${displayName}${machineTag}`,
			serverModelId: rawId,
			name: `${displayName}${badgeSuffix(modelMeta, settings.showBadgesInNames)}`,
			reasoning: false,
			contextWindow,
			maxTokens: model.max_tokens ?? Math.min(contextWindow, settings.maxOutputTokens),
			compat,
		};
	}

	// lucebox: alias as source of truth, rich metadata from /props.
	if (kind === "lucebox") {
		const modelPath = ep.props?.model_path || rawId;
		const alias = ep.props?.model_alias || rawId;
		const nCtx = ep.props?.default_generation_settings?.n_ctx;
		const contextWindow = model.context_length ?? nCtx ?? 8192;
		const displayName = cleanModelName(baseName(alias !== rawId ? alias : modelPath));
		return {
			...common,
			id: `${displayName}${machineTag}`,
			serverModelId: alias,
			name: `${displayName}${badgeSuffix(modelMeta, settings.showBadgesInNames)}`,
			reasoning: ep.props?.capabilities?.reasoning_supported ?? true,
			contextWindow,
			maxTokens: model.max_tokens ?? Math.min(contextWindow, settings.maxOutputTokens),
			compat,
		};
	}

	// vLLM / SGLang / TGI: OpenAI-compatible, context in max_model_len, no
	// meta.n_ctx, no GGUF KV-cache, thinking driven by the chat template.
	if (kind === "vllm") {
		const displayName = cleanModelName(rawId);
		const contextWindow =
			model.max_model_len ?? model.context_window ?? model.context_length ?? 32768;
		// vLLM serves one model per process (no router/status/slots), and
		// cacheK/cacheV are llama.cpp/GGUF-specific.
		const { routerStatus: _routerStatus, cacheK: _cacheK, cacheV: _cacheV, ...commonVllm } = common;
		return {
			...commonVllm,
			id: `${displayName}${machineTag}`,
			serverModelId: rawId,
			name: `${displayName}${badgeSuffix(modelMeta, settings.showBadgesInNames)}`,
			reasoning: true,
			contextWindow,
			maxTokens: model.max_tokens ?? Math.min(contextWindow, settings.maxOutputTokens),
			compat,
		};
	}

	// ZINC / llama.cpp / dwarfstar.
	const cleanName = cleanModelName(rawId);
	const nameSource = model.path ?? rawId;
	const cleanFromSource = cleanModelName(baseName(nameSource));
	const baseName2 =
		model.name && String(model.name) !== rawId ? cleanModelName(String(model.name)) : cleanFromSource || cleanName;
	const nameMapName = ep.nameMap?.get(rawId);
	const displayName = nameMapName ? cleanModelName(nameMapName) : baseName2;

	const contextWindow =
		model.meta?.n_ctx ?? model.meta?.n_ctx_train ?? model.max_model_len ??
		model.context_window ?? model.context_length ?? 32768;

	const isLlamaFamily = kind === "llamacpp";

	return {
		...common,
		id: `${displayName}${machineTag}`,
		serverModelId: rawId,
		name: `${displayName}${badgeSuffix(modelMeta, settings.showBadgesInNames)}`,
		reasoning: isLlamaFamily,
		contextWindow,
		maxTokens: model.max_tokens ?? Math.min(contextWindow, settings.maxOutputTokens),
		compat,
	};
}

/**
 * Build pi models from a scan and (re)register the provider.
 * Returns the registered model list.
 */
export function buildAndRegisterProvider(
	pi: ExtensionAPI,
	scan: ScanResult,
	config: InfraConfig,
	options?: { persistCache?: boolean },
): PiModel[] {
	shared.zincModelIds.clear();
	shared.serverModelIds.clear();
	shared.compactModelIds.clear();
	const settings = config.settings;
	let configDirty = false;

	const piModels: PiModel[] = [];
	const seenIds = new Set<string>();
	for (const ep of scan.endpoints) {
		if (!ep.ok) continue;
		const srv = config.servers.find((s) => s.id === ep.serverId && s.host === ep.host);
		if (!srv) continue;

		const visibleModels = ep.models.filter((m) => {
			const lmType = m.lmStudio?.type;
			if (ep.server === "lmstudio") return lmType !== "embedding" && lmType !== "embeddings";
			const st = ep.meta.get(String(m.id ?? ""))?.routerStatus;
			return !st || st === "loaded" || settings.includeUnloadedRouterModels;
		});

		const nameCount = new Map<string, number>();
		for (const model of visibleModels) {
			const n = cleanModelName(String(model.name ?? model.id));
			nameCount.set(n, (nameCount.get(n) ?? 0) + 1);
		}

		for (const model of visibleModels) {
			const rawId = String(model.id ?? "");
			const modelMeta = ep.meta.get(rawId) ?? {};
			const pm = toPiModel(model, ep, srv, settings, modelMeta);
			const hostPort = `${idSafeHost(srv.host)}:${ep.port}`;

			// Per-model output cap override (modelOptions[id].maxTokens): wins over
			// both the server-reported max_tokens and settings.maxOutputTokens, and
			// is clamped to the model contextWindow. Invalid values are ignored.
			const overrideMax = modelOptionsFor(rawId)?.maxTokens;
			if (typeof overrideMax === "number" && Number.isFinite(overrideMax) && overrideMax > 0) {
				pm.maxTokens = Math.min(Math.floor(overrideMax), pm.contextWindow);
			}

			// ID collision guard. `toPiModel` already appends the (host:port) tag
			// when the raw id doesn't carry it; here we only add it for the raw id
			// itself (which the tag is meant to disambiguate) and then a numeric
			// suffix for any remaining clash. Never touch serverModelId: it must
			// stay the raw server id or the request payload goes out with a
			// display-shaped id and OpenAI-compatible servers answer 404.
			if (seenIds.has(pm.id)) {
				if (!pm.id.includes(`(${hostPort})`)) pm.id = `${pm.id} (${hostPort})`;
				let n = 2;
				while (seenIds.has(pm.id)) pm.id = `${pm.id}-${n++}`;
			}
			seenIds.add(pm.id);

			// Thinking budgets: resolve by registered (compact) id or legacy
			// "host:port/model" keys, then migrate legacy keys forward.
			const opts = modelOptions();
			let budgets = opts[pm.id]?.thinkingBudgets;
			if (!budgets) {
				for (const legacyKey of [
					`${hostPort}/${pm.serverModelId.replace(/^\/+/, "")}`,
					pm.serverModelId,
				]) {
					const entry = opts[legacyKey];
					if (entry?.thinkingBudgets) {
						if (!opts[pm.id]) opts[pm.id] = entry;
						delete opts[legacyKey];
						configDirty = true;
						budgets = entry.thinkingBudgets;
						break;
					}
				}
			}
			if (budgets && ep.server !== "lmstudio") {
				pm.thinkingBudgets = budgets;
				pm.reasoning = true;
			}

			// Duplicate display names get the raw id appended.
			const candidateName = cleanModelName(String(model.name ?? model.id));
			if (nameCount.get(candidateName)! > 1) {
				pm.name = `${candidateName}${badgeSuffix(modelMeta, settings.showBadgesInNames)} (${pm.serverModelId})`;
			}
			piModels.push(pm);
			if (ep.server === "zinc") {
				shared.zincModelIds.add(pm.id);
				shared.zincModelIds.add(pm.serverModelId);
			}
			shared.serverModelIds.set(pm.id, pm.serverModelId);
			if (!shared.compactModelIds.has(pm.serverModelId)) shared.compactModelIds.set(pm.serverModelId, pm.id);
		}
	}

	if (configDirty) saveConfig(config);

	shared.endpointKinds.clear();
	for (const ep of scan.endpoints) {
		if (ep.ok) shared.endpointKinds.set(ep.baseUrl, ep.server);
	}
	shared.modelBaseUrls.clear();
	for (const pm of piModels) shared.modelBaseUrls.set(pm.id, pm.baseUrl);

	const first = config.servers.find((s) => s.enabled && s.ports.length > 0);
	const defaultBaseUrl =
		scan.endpoints.find((e) => e.ok)?.baseUrl ?? `http://${first?.host ?? "127.0.0.1"}:${first?.ports[0] ?? 8080}/v1`;
	const defaultApiKey = first?.apiKey || DEFAULT_API_KEY;

	try {
		pi.unregisterProvider(PROVIDER_NAME);
	} catch {
		// not registered yet
	}

	pi.registerProvider(PROVIDER_NAME, {
		name: `🦙 llama-infra (${scan.totalModels} on ${scan.serversUp}/${scan.serversTotal} servers)`,
		baseUrl: defaultBaseUrl,
		apiKey: defaultApiKey,
		api: "openai-completions",
		streamSimple: createLongTimeoutOpenAICompletionsStream,
		// `compat.thinkingTokenBudgetField: false` is pi's documented runtime
		// opt-out ("Default: false"), but its TS type only allows the field-name
		// union → cast the models array here.
		models: piModels as any,
	});

	// Persist the last known good models so a fresh pi process can register them
	// synchronously at boot (before async discovery completes) — required for CLI
	// consumers that resolve --model right after startup (e.g. trimegisto spawns).
	// Tests opt out so they never touch the real per-user cache.
	if (piModels.length > 0 && options?.persistCache !== false) saveModelsCache(piModels);

	return piModels;
}