/**
 * pi-fastmode-web: automatic Fast Mode for pi CLI and pi-web.
 */
import { existsSync, readFileSync } from "node:fs";
import { dirname, join, resolve, sep } from "node:path";
import { fileURLToPath } from "node:url";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	ExtensionFactory,
} from "@earendil-works/pi-coding-agent";

const DEFAULT_EXTENSION_DIR = dirname(fileURLToPath(import.meta.url));
const STATE_ENTRY_TYPE = "pi-fastmode-web-state";

type FastPath = "off" | "tier" | "model" | "unsupported";

type ModelRef = {
	provider: string;
	id: string;
	thinkingLevel?: string;
};

type FastModeState = {
	enabled: boolean;
	path: FastPath;
	original?: ModelRef;
	switchingModel: boolean;
	started: boolean;
};

export interface FastModeConfig {
	enabled: boolean;
	tierProviders: string[];
	fallbackModels: string[];
	serviceTier: string;
}

const DEFAULT_CONFIG: FastModeConfig = {
	enabled: true,
	tierProviders: ["openai", "openai-codex"],
	fallbackModels: ["gpt-5.3-codex-spark", "gpt-5.4-mini"],
	serviceTier: "fast",
};

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function normalizeStringArray(value: unknown, fallback: string[]): string[] {
	if (
		!Array.isArray(value) ||
		!value.every((item) => typeof item === "string")
	) {
		return fallback;
	}
	const normalized = value.map((item) => item.trim()).filter(Boolean);
	return normalized.length ? normalized : fallback;
}

function normalizeServiceTier(value: unknown): string {
	if (typeof value !== "string" || !value.trim())
		return DEFAULT_CONFIG.serviceTier;
	const normalized = value.trim();
	return normalized === "priority" ? "fast" : normalized;
}

function normalizeConfig(raw: unknown): FastModeConfig {
	const cfg = { ...DEFAULT_CONFIG };
	if (!isRecord(raw)) return cfg;
	if (typeof raw.enabled === "boolean") cfg.enabled = raw.enabled;
	cfg.tierProviders = normalizeStringArray(
		raw.tierProviders,
		cfg.tierProviders,
	);
	cfg.fallbackModels = normalizeStringArray(
		raw.fallbackModels,
		cfg.fallbackModels,
	);
	cfg.serviceTier = normalizeServiceTier(raw.serviceTier);
	return cfg;
}

function parseConfig(json: string): FastModeConfig {
	try {
		return normalizeConfig(JSON.parse(json));
	} catch {
		return { ...DEFAULT_CONFIG };
	}
}

/** Load config: project-level first, then user-level, then defaults. */
function loadConfig(cwd: string, extensionDir: string): FastModeConfig {
	const projectPath = join(
		resolve(cwd),
		".pi",
		"pi-fastmode-web",
		"config.json",
	);
	if (existsSync(projectPath)) {
		return parseConfig(readFileSync(projectPath, "utf8"));
	}
	const isProjectLocal =
		resolve(extensionDir) === resolve(cwd, ".pi") ||
		resolve(extensionDir).startsWith(resolve(cwd, ".pi") + sep);
	if (!isProjectLocal) {
		const userPath = join(
			process.env.PI_AGENT_DIR ?? join(process.env.HOME ?? "", ".pi", "agent"),
			"extensions",
			"pi-fastmode-web",
			"config.json",
		);
		if (existsSync(userPath)) {
			return parseConfig(readFileSync(userPath, "utf8"));
		}
	}
	return { ...DEFAULT_CONFIG };
}

function toModelRef(
	model: unknown,
): { provider: string; id: string } | undefined {
	if (!isRecord(model)) return undefined;
	const { provider, id } = model;
	if (typeof provider !== "string" || typeof id !== "string") return undefined;
	if (!provider || !id) return undefined;
	return { provider, id };
}

function isGpt5(modelId: string): boolean {
	return /^gpt-5(\.|$)/.test(modelId);
}

function validOriginal(value: unknown): ModelRef | undefined {
	const ref = toModelRef(value);
	if (!ref) return undefined;
	if (isRecord(value) && typeof value.thinkingLevel === "string") {
		return { ...ref, thinkingLevel: value.thinkingLevel };
	}
	return ref;
}

function validStoredState(
	value: unknown,
): Pick<FastModeState, "enabled" | "original"> | undefined {
	if (!isRecord(value) || typeof value.enabled !== "boolean") return undefined;
	return { enabled: value.enabled, original: validOriginal(value.original) };
}

function latestStoredState(
	ctx: ExtensionContext,
): Pick<FastModeState, "enabled" | "original"> | undefined {
	const entries =
		ctx.sessionManager?.getBranch?.() ??
		ctx.sessionManager?.getEntries?.() ??
		[];
	for (let index = entries.length - 1; index >= 0; index -= 1) {
		const entry = entries[index];
		if (
			isRecord(entry) &&
			entry.type === "custom" &&
			entry.customType === STATE_ENTRY_TYPE
		) {
			const state = validStoredState(entry.data);
			if (state) return state;
		}
	}
	return undefined;
}

function persistState(pi: ExtensionAPI, state: FastModeState): void {
	pi.appendEntry(STATE_ENTRY_TYPE, {
		enabled: state.enabled,
		original: state.original,
	});
}

function statusText(path: FastPath, modelId?: string): string {
	return path === "off"
		? "FAST · off"
		: path === "tier"
			? "FAST · tier requested"
			: path === "model" && modelId
				? `FAST · model ${modelId}`
				: "FAST · unsupported";
}

function setStatus(
	ctx: ExtensionContext,
	path: FastPath,
	modelId?: string,
): void {
	if (!ctx.hasUI) return;
	ctx.ui.setStatus("fast-mode", statusText(path, modelId));
}

function notifyStatus(
	ctx: ExtensionContext,
	path: FastPath,
	modelId?: string,
): void {
	if (!ctx.hasUI) return;
	ctx.ui.notify(statusText(path, modelId));
}

function findFallback(
	ctx: ExtensionContext,
	config: FastModeConfig,
	provider: string,
): NonNullable<ExtensionContext["model"]> | undefined {
	for (const id of config.fallbackModels) {
		const fallback = ctx.modelRegistry.find(provider, id);
		if (fallback) return fallback;
	}
	return undefined;
}

export function createFastModeDefaultExtension(
	options: { extensionDir?: string } = {},
): ExtensionFactory {
	const extensionDir = options.extensionDir ?? DEFAULT_EXTENSION_DIR;

	return function fastModeDefaultExtension(pi: ExtensionAPI): void {
		let config: FastModeConfig = { ...DEFAULT_CONFIG };
		let loadedCwd: string | undefined;
		const state: FastModeState = {
			enabled: DEFAULT_CONFIG.enabled,
			path: "off",
			switchingModel: false,
			started: false,
		};

		function reloadConfig(ctx: { cwd?: string }, force = false): void {
			if (force || !loadedCwd || loadedCwd !== ctx.cwd) {
				config = loadConfig(ctx.cwd ?? "", extensionDir);
				loadedCwd = ctx.cwd;
			}
		}

		function ensureConfig(ctx: { cwd?: string }): void {
			reloadConfig(ctx);
		}

		async function applyFast(ctx: ExtensionContext): Promise<void> {
			const model = toModelRef(ctx.model);
			if (!state.enabled || !model) {
				state.path = "off";
				setStatus(ctx, state.path);
				persistState(pi, state);
				return;
			}

			if (config.tierProviders.includes(model.provider) && isGpt5(model.id)) {
				state.path = "tier";
				setStatus(ctx, state.path);
				persistState(pi, state);
				return;
			}

			if (config.fallbackModels.includes(model.id)) {
				state.path = "model";
				setStatus(ctx, state.path, model.id);
				persistState(pi, state);
				return;
			}

			const fallback = findFallback(ctx, config, model.provider);
			if (!fallback) {
				state.path = "unsupported";
				setStatus(ctx, state.path);
				persistState(pi, state);
				return;
			}

			if (!state.original) {
				const thinkingLevel = pi.getThinkingLevel();
				state.original = {
					...model,
					...(thinkingLevel ? { thinkingLevel } : {}),
				};
			}
			state.switchingModel = true;
			try {
				if (await pi.setModel(fallback)) {
					state.path = "model";
					setStatus(ctx, state.path, fallback.id);
				} else {
					state.path = "unsupported";
					setStatus(ctx, state.path);
				}
			} finally {
				state.switchingModel = false;
			}
			persistState(pi, state);
		}

		async function restoreOriginal(
			ctx: ExtensionCommandContext,
		): Promise<void> {
			state.enabled = false;
			state.path = "off";
			const original = state.original;
			if (original) {
				const model = ctx.modelRegistry.find(original.provider, original.id);
				if (model) {
					state.switchingModel = true;
					try {
						if (await pi.setModel(model)) {
							if (original.thinkingLevel) {
								pi.setThinkingLevel(
									original.thinkingLevel as Parameters<
										ExtensionAPI["setThinkingLevel"]
									>[0],
								);
							}
							state.original = undefined;
						} else if (ctx.hasUI) {
							ctx.ui.notify(
								"FAST · unable to restore original model",
								"warning",
							);
						}
					} finally {
						state.switchingModel = false;
					}
				} else if (ctx.hasUI) {
					ctx.ui.notify("FAST · unable to restore original model", "warning");
				}
			}
			setStatus(ctx, state.path);
			persistState(pi, state);
		}

		pi.registerFlag("fast", {
			description: "Start with Fast Mode enabled",
			type: "boolean",
		});

		pi.registerCommand("fast", {
			description:
				"Fast Mode status and controls. Usage: /fast [status|on|off]",
			getArgumentCompletions: (prefix: string) => {
				const p = prefix.trim().toLowerCase();
				return ["status", "on", "off"]
					.filter((o) => o.startsWith(p))
					.map((value) => ({ value, label: value }));
			},
			handler: async (
				args: string,
				ctx: ExtensionCommandContext,
			): Promise<void> => {
				const arg = args.trim().toLowerCase() || "status";
				if (arg === "status") {
					notifyStatus(ctx, state.path, ctx.model?.id);
					return;
				}
				if (arg !== "on" && arg !== "off") {
					if (ctx.hasUI) ctx.ui.notify("Usage: /fast [status|on|off]", "error");
					return;
				}
				reloadConfig(ctx, true);
				if (arg === "on") {
					state.enabled = true;
					await applyFast(ctx);
				} else {
					await restoreOriginal(ctx);
				}
			},
		});

		pi.on("session_start", async (_event, ctx) => {
			reloadConfig(ctx, true);
			const stored = latestStoredState(ctx);
			if (stored) {
				state.enabled = stored.enabled;
				state.original = stored.original;
			} else {
				state.enabled = config.enabled;
			}
			state.started = true;
			await applyFast(ctx);
		});

		pi.on("model_select", async (_event, ctx) => {
			if (!state.started || state.switchingModel) return;
			ensureConfig(ctx);
			state.original = undefined;
			await applyFast(ctx);
		});

		pi.on("before_provider_request", (event, ctx) => {
			ensureConfig(ctx);
			const model = toModelRef(ctx.model);
			if (!state.enabled || state.path !== "tier" || !model) return undefined;
			if (!config.tierProviders.includes(model.provider) || !isGpt5(model.id))
				return undefined;
			if (!isRecord(event.payload)) return undefined;
			return {
				...event.payload,
				service_tier: config.serviceTier,
			};
		});
	};
}

const fastModeDefaultExtension = createFastModeDefaultExtension();

export default fastModeDefaultExtension;