/**
 * Alibaba Model Studio "token plan" provider for Prime Agent.
 *
 * `/login alibaba` stores a pasted API key in ~/.prime/agent/auth.json, and that is the
 * only place the key lives. ALIBABA_API_KEY is the sole fallback; when neither is present
 * the provider reports an unresolved key instead of sending a placeholder, so a missing
 * credential cannot be mistaken for an expired one.
 * ALIBABA_BASE_URL overrides the endpoint (e.g. another region).
 *
 * The model list is baked into models.json and regenerated by regen.ts, which
 * probes the live endpoint. Nothing here is fetched at startup except the
 * account's entitled ids, so capabilities never drift between runs.
 */

import type { Api } from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";

export const PROVIDER = "alibaba";
export const DEFAULT_BASE_URL = "https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1";
export const CATALOG = join(fileURLToPath(new URL(".", import.meta.url)), "models.json");
const FETCH_TIMEOUT_MS = 10000;
const AUTH_FILE = join(homedir(), ".prime", "agent", "auth.json");

/**
 * What Prime Agent resolves when no credential is stored. `!`-prefixed values run through
 * the shell (bash -c), and an empty result is an unresolved key rather than an empty
 * string, which is what turns "no key" into a readable error instead of a 401.
 */
const API_KEY_FALLBACK = '!printf %s "$ALIBABA_API_KEY"';
const ENTITLEMENT_CACHE = join(homedir(), ".prime", "agent", "alibaba-entitlement.json");

export interface BakedModel {
	id: string;
	name: string;
	reasoning: boolean;
	input: ("text" | "image")[];
	contextWindow: number;
	maxTokens: number;
}

export interface BakedCatalog {
	generatedAt: string;
	baseUrl: string;
	models: BakedModel[];
	/** entitled ids the bake considered and refused */
	skipped: string[];
}

type Family = "qwen" | "deepseek" | "zai" | "other";

export const familyOf = (id: string): Family =>
	id.startsWith("qwen") ? "qwen" : id.startsWith("deepseek") ? "deepseek" : id.startsWith("glm") ? "zai" : "other";

/**
 * Per-family request quirks, mirroring what Prime's own catalog applies to the same
 * model families. Without the DeepSeek entry, replayed thinking blocks from another
 * provider land in the assistant text instead of `reasoning_content`, and the model
 * answers by writing its reasoning into the visible reply.
 */
const FAMILY_COMPAT: Record<Family, Record<string, unknown>> = {
	qwen: { thinkingFormat: "qwen", supportsReasoningEffort: false, cacheControlFormat: "anthropic" },
	deepseek: { thinkingFormat: "deepseek", requiresReasoningContentOnAssistantMessages: true },
	// glm depth is reasoning_effort, no enable_thinking: thinking is on by default here and
	// glm-5.3 rejects enable_thinking false outright
	zai: { supportsReasoningEffort: true },
	other: {},
};

/**
 * Model Studio's own levels, not the vendors' defaults.
 *
 * DeepSeek takes the full `reasoning_effort` ladder with high as its default. Qwen takes a
 * boolean `enable_thinking` only. GLM takes `reasoning_effort` too, with a value set that
 * varies by model, so the family entry is a superset and `MODEL_THINKING_LEVELS` narrows it.
 */
const THINKING_LEVELS: Record<Family, Record<string, string | null> | undefined> = {
	deepseek: { minimal: null, low: "low", medium: "medium", high: "high", xhigh: "xhigh", max: "max" },
	qwen: { minimal: null, low: null, medium: null, high: "high" },
	zai: { off: "none", minimal: "minimal", low: "low", medium: "medium", high: "high", xhigh: "xhigh", max: "max" },
	other: undefined,
};

/** glm-5.3 takes only low/high/max, and 400s on every attempt to stop it thinking */
const MODEL_THINKING_LEVELS: Record<string, Record<string, string | null>> = {
	"glm-5.3": { off: null, minimal: null, low: "low", medium: null, high: "high", xhigh: null, max: "max" },
};

const thinkingLevelsFor = (id: string): Record<string, string | null> | undefined =>
	MODEL_THINKING_LEVELS[id] ?? THINKING_LEVELS[familyOf(id)];

export const baseUrlFromEnv = (): string => (process.env.ALIBABA_BASE_URL || DEFAULT_BASE_URL).replace(/\/+$/, "");

const readFile = (path: string): string | undefined => {
	try {
		return readFileSync(path, "utf-8");
	} catch {
		return undefined;
	}
};

/** Key saved by `/login alibaba`; the only place this provider reads a key from. */
const readStoredKey = (): string | undefined => {
	const raw = readFile(AUTH_FILE);
	if (!raw) return undefined;
	try {
		return (JSON.parse(raw) as Record<string, { key?: string }>)[PROVIDER]?.key;
	} catch {
		return undefined;
	}
};

export const resolveKey = (): string | undefined => readStoredKey() ?? process.env.ALIBABA_API_KEY;

export const loadCatalog = (): BakedCatalog => JSON.parse(readFileSync(CATALOG, "utf-8")) as BakedCatalog;

/** Model ids the account can actually call, straight from the endpoint. */
export async function accountModelIds(baseUrl: string, apiKey: string): Promise<string[]> {
	const response = await fetch(`${baseUrl}/models`, {
		headers: { Authorization: `Bearer ${apiKey}` },
		signal: AbortSignal.timeout(FETCH_TIMEOUT_MS),
	});
	if (!response.ok) throw new Error(`${baseUrl}/models ${response.status}`);
	const payload = (await response.json()) as { data?: Array<{ id: string }> };
	const ids = (payload.data ?? []).map((m) => m.id);
	if (ids.length === 0) throw new Error("no entitled models");
	return ids.sort();
}

/** Entitled ids, cached on disk: a slow or failed lookup reuses the last known list. */
async function entitledIds(baseUrl: string, apiKey: string | undefined): Promise<Set<string> | undefined> {
	const cached = (): Set<string> | undefined => {
		const raw = readFile(ENTITLEMENT_CACHE);
		if (!raw) return undefined;
		try {
			const ids = JSON.parse(raw) as string[];
			return ids.length > 0 ? new Set(ids) : undefined;
		} catch {
			return undefined;
		}
	};
	if (!apiKey) return cached();
	try {
		const ids = await accountModelIds(baseUrl, apiKey);
		writeFileSync(ENTITLEMENT_CACHE, `${JSON.stringify(ids, null, "\t")}\n`);
		return new Set(ids);
	} catch {
		return cached();
	}
}

export default async function (pi: ExtensionAPI) {
	const baseUrl = baseUrlFromEnv();
	const catalog = loadCatalog();
	const entitled = await entitledIds(baseUrl, resolveKey());

	const models = catalog.models
		.filter((m) => !entitled || entitled.has(m.id))
		.map((m) => ({
			...m,
			// The token plan bills a flat subscription, so per-token cost is zero.
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
			compat: {
				supportsDeveloperRole: false,
				maxTokensField: "max_tokens" as const,
				...FAMILY_COMPAT[familyOf(m.id)],
			},
			...(m.reasoning && thinkingLevelsFor(m.id) ? { thinkingLevelMap: thinkingLevelsFor(m.id) } : {}),
		}));

	// entitlement can only hide, never add, so an id gained after the last bake stays
	// invisible until regen.ts runs
	const known = new Set([...catalog.models.map((m) => m.id), ...catalog.skipped]);
	const unbaked = entitled === undefined ? [] : [...entitled].filter((id) => !known.has(id)).sort();
	pi.on("session_start", (_event, ctx) => {
		if (unbaked.length > 0) {
			ctx.ui.notify(`alibaba: entitled but not in models.json: ${unbaked.join(", ")} - run npx tsx regen.ts`, "warning");
		}
	});

	pi.registerProvider(PROVIDER, {
		name: "Alibaba Model Studio (token plan)",
		baseUrl,
		apiKey: API_KEY_FALLBACK,
		api: "openai-completions" as Api,
		models,
	});
}