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

const EFFORT_LEVELS = [
	"off",
	"minimal",
	"low",
	"medium",
	"high",
	"xhigh",
	"max",
] as const;
type EffortLevel = (typeof EFFORT_LEVELS)[number];

const FAST_STATE_ENTRY = "effort-controls.fast-state";

function supportedEfforts(ctx: ExtensionContext): EffortLevel[] {
	const model = ctx.model;
	if (!model?.reasoning) return ["off"];

	return EFFORT_LEVELS.filter((level) => {
		const mapped = model.thinkingLevelMap?.[level];
		if (mapped === null) return false;
		return level !== "xhigh" && level !== "max"
			? true
			: mapped !== undefined && mapped !== null;
	});
}

function supportsFastMode(ctx: ExtensionContext): boolean {
	const model = ctx.model;
	if (!model) return false;

	const isOpenAIResponses =
		(model.provider === "openai" && model.api === "openai-responses") ||
		(model.provider === "openai-codex" && model.api === "openai-codex-responses");

	return isOpenAIResponses && /^gpt-5\.(?:[4-9]|\d{2,})(?:-|$)/.test(model.id);
}

function updateFastStatus(ctx: ExtensionContext, fastEnabled: boolean): void {
	ctx.ui.setStatus(
		"effort-controls.fast",
		fastEnabled && supportsFastMode(ctx)
			? ctx.ui.theme.fg("accent", "fast:on")
			: undefined,
	);
}

export default function effortControlsExtension(pi: ExtensionAPI) {
	let fastEnabled = false;

	pi.registerCommand("effort", {
		description: "Choose the model reasoning effort",
		getArgumentCompletions: (prefix) => {
			const levels = EFFORT_LEVELS.filter((level) =>
				level.startsWith(prefix.toLowerCase()),
			);
			return levels.length > 0
				? levels.map((level) => ({ value: level, label: level }))
				: null;
		},
		handler: async (args, ctx) => {
			const requested = args.trim().toLowerCase();
			const supported = supportedEfforts(ctx);

			if (requested) {
				if (!EFFORT_LEVELS.includes(requested as EffortLevel)) {
					ctx.ui.notify(
						`Unknown effort "${requested}". Choose: ${supported.join(", ")}`,
						"warning",
					);
					return;
				}
				if (!supported.includes(requested as EffortLevel)) {
					ctx.ui.notify(
						`Effort "${requested}" is unavailable for ${ctx.model?.id ?? "this model"}.`,
						"warning",
					);
					return;
				}
				pi.setThinkingLevel(requested as EffortLevel);
				ctx.ui.notify(`Reasoning effort: ${pi.getThinkingLevel()}`, "info");
				return;
			}

			if (!ctx.hasUI) {
				ctx.ui.notify(
					`Available efforts: ${supported.join(", ")}. Use /effort <level>.`,
					"info",
				);
				return;
			}

			const choice = await ctx.ui.select("Reasoning effort", supported);
			if (!choice) return;

			pi.setThinkingLevel(choice as EffortLevel);
			ctx.ui.notify(`Reasoning effort: ${pi.getThinkingLevel()}`, "info");
		},
	});

	pi.registerCommand("fast", {
		description:
			"Toggle OpenAI priority (fast) mode for eligible GPT-5.4+ models",
		getArgumentCompletions: (prefix) => {
			const values = ["on", "off"];
			const matches = values.filter((value) =>
				value.startsWith(prefix.toLowerCase()),
			);
			return matches.length > 0
				? matches.map((value) => ({ value, label: value }))
				: null;
		},
		handler: async (args, ctx) => {
			if (!supportsFastMode(ctx)) {
				ctx.ui.notify(
					"Fast mode requires an OpenAI Responses or Codex Responses GPT-5.4+ model.",
					"warning",
				);
				return;
			}

			const requested = args.trim().toLowerCase();
			if (requested && requested !== "on" && requested !== "off") {
				ctx.ui.notify("Use /fast, /fast on, or /fast off.", "warning");
				return;
			}

			fastEnabled = requested ? requested === "on" : !fastEnabled;
			pi.appendEntry(FAST_STATE_ENTRY, { enabled: fastEnabled });
			updateFastStatus(ctx, fastEnabled);
			ctx.ui.notify(`Fast mode ${fastEnabled ? "on" : "off"}.`, "info");
		},
	});

	pi.on("before_provider_request", (event, ctx) => {
		if (
			!fastEnabled ||
			!supportsFastMode(ctx) ||
			typeof event.payload !== "object" ||
			event.payload === null ||
			Array.isArray(event.payload)
		) {
			return;
		}
		return { ...event.payload, service_tier: "priority" };
	});

	pi.on("session_start", (_event, ctx) => {
		const latestState = ctx.sessionManager
			.getBranch()
			.filter(
				(entry) => entry.type === "custom" && entry.customType === FAST_STATE_ENTRY,
			)
			.pop() as { data?: { enabled?: unknown } } | undefined;
		fastEnabled = latestState?.data?.enabled === true;
		updateFastStatus(ctx, fastEnabled);
	});

	pi.on("model_select", (_event, ctx) => {
		updateFastStatus(ctx, fastEnabled);
	});
}