import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import {
	getAgentDir,
	type ExtensionAPI,
	type ExtensionContext,
} from "@earendil-works/pi-coding-agent";

/**
 * pi-fast-all — one fast-mode toggle for every provider that has one.
 *
 * Providers expose "pay more, go faster" through different mechanisms:
 *
 *   Anthropic       body { speed: "fast" }        + anthropic-beta: fast-mode-2026-02-01
 *   OpenAI          body { service_tier: "priority" }
 *   OpenAI Codex    body { service_tier: "priority" }
 *
 * Rather than hardcoding one of those, this extension is a table of rules. A
 * rule matches on provider/api/model and declares a `body` patch and/or header
 * work. Adding a provider is a config edit, not a release.
 */

const EXTENSION_ID = "pi-fast-all";
const STATUS_KEY = "fast";
const WIDGET_KEY = "pi-fast-all";
const DEFAULT_INDICATOR_TEXT = "⚡fast";
/** Soft orange. Reads as "this costs extra" without shouting like an error. */
const DEFAULT_INDICATOR_COLOR = "#ffa657";
const HEX_COLOR = /^#([0-9a-f]{6})$/i;
const ANTHROPIC_OAUTH_BETAS = ["claude-code-20250219", "oauth-2025-04-20"];

/** A single "this is how this provider does fast" declaration. */
type Rule = {
	/** Provider id, e.g. `anthropic`. Required. */
	provider: string;
	/** Optional api narrowing, e.g. `anthropic-messages`. */
	api?: string;
	/** Model ids the rule applies to. `["*"]` matches every model of the provider. */
	models: string[];
	/** Keys merged into the request body. Existing keys are never overwritten. */
	body?: Record<string, unknown>;
	/** Headers set on the model before requests. */
	headers?: Record<string, string>;
	/** Value appended (comma-separated, deduped) to `anthropic-beta`. */
	betaHeader?: string;
	/** What fast multiplies the token bill by. Used for spend estimates. */
	costMultiplier?: number;
	/**
	 * True when the host already prices this tier (pi-ai reads OpenAI's
	 * `service_tier` back off the response). False means pi's reported cost is
	 * the *standard* price and the real bill is `reported * costMultiplier`.
	 */
	costCountedByHost?: boolean;
	/** Human-readable cost warning surfaced by `/fast status`. */
	note?: string;
};

const DEFAULT_COST_MULTIPLIER = 2;

/**
 * Where to draw the "fast is on" marker.
 *
 * `status` is pi's native footer status line, but hosts are free to replace the
 * footer — arc does, and its replacement never reads extension statuses — so a
 * widget above the editor is the portable default.
 */
type Indicator = "widget" | "status" | "both" | "off";

const INDICATORS: Indicator[] = ["widget", "status", "both", "off"];

type Config = {
	/** Fast state at session start. `/fast` overrides it for the session. */
	enabled: boolean;
	/** Where to show the marker while fast mode is active. */
	indicator: Indicator;
	/** Marker text. */
	indicatorText: string;
	/** `#rrggbb`, or a theme colour name such as `warning`. */
	indicatorColor: string;
	/** Recolour the marker once the fast premium passes this many USD. 0 disables. */
	warnAboveUsd: number;
	/** Replaces the built-in rules when present. */
	rules?: Rule[];
	/** Always appended to whichever rule set is in effect. */
	extraRules?: Rule[];
};

const DEFAULT_RULES: Rule[] = [
	{
		provider: "anthropic",
		api: "anthropic-messages",
		// Opus only, and not every Opus: the API rejects `speed` outright on
		// 4.7, on every Sonnet, on Haiku and on Fable (400, not a silent no-op).
		// Verified against api.anthropic.com, 2026-08.
		models: ["claude-opus-5", "claude-opus-4-8", "claude-opus-4-6"],
		body: { speed: "fast" },
		betaHeader: "fast-mode-2026-02-01",
		note: "~2x token cost; separate rate limits; no prompt-cache sharing with standard speed",
	},
	{
		provider: "openai",
		api: "openai-responses",
		models: ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5", "gpt-5.4"],
		body: { service_tier: "priority" },
		costCountedByHost: true,
		note: "priority tier: 2x token cost (2.5x on gpt-5.5)",
	},
	{
		provider: "openai-codex",
		api: "openai-codex-responses",
		models: ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5", "gpt-5.4"],
		body: { service_tier: "priority" },
		costCountedByHost: true,
		note: "priority tier: 2x token cost (2.5x on gpt-5.5)",
	},
];

const DEFAULT_CONFIG: Config = {
	enabled: false,
	indicator: "widget",
	indicatorText: DEFAULT_INDICATOR_TEXT,
	indicatorColor: DEFAULT_INDICATOR_COLOR,
	warnAboveUsd: 5,
};

const WARN_INDICATOR_COLOR = "#ff6b4a";

type Override = "auto" | "on" | "off";

type Spend = {
	/** Requests this extension actually patched. */
	requests: number;
	/** What those turns really cost, in USD. */
	billedUsd: number;
	/** How much of that was the fast premium. */
	extraUsd: number;
	/** True when pi's own cost display is missing some of `billedUsd`. */
	underReported: boolean;
};

type SessionState = {
	config: Config;
	rules: Rule[];
	override: Override;
	/** Armed by `/fast once`: fast for the current turn, then back to normal. */
	once: boolean;
	/** Rule used by the in-flight request, awaiting its assistant message. */
	pending?: Rule;
	spend: Spend;
	lastAppliedAt?: number;
};

type Match = { rule?: Rule; modelKey: string; reason?: string };

// --------------------------------------------------------------------------- config

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

function normalizeRule(raw: unknown): Rule | undefined {
	if (!isRecord(raw)) return undefined;
	if (typeof raw.provider !== "string" || !raw.provider) return undefined;
	if (!Array.isArray(raw.models) || !raw.models.every((m) => typeof m === "string")) return undefined;

	const rule: Rule = { provider: raw.provider, models: raw.models as string[] };
	if (typeof raw.api === "string") rule.api = raw.api;
	if (isRecord(raw.body)) rule.body = raw.body;
	if (isRecord(raw.headers)) {
		const headers: Record<string, string> = {};
		for (const [k, v] of Object.entries(raw.headers)) if (typeof v === "string") headers[k] = v;
		rule.headers = headers;
	}
	if (typeof raw.betaHeader === "string") rule.betaHeader = raw.betaHeader;
	if (typeof raw.costMultiplier === "number" && raw.costMultiplier > 0) {
		rule.costMultiplier = raw.costMultiplier;
	}
	if (typeof raw.costCountedByHost === "boolean") rule.costCountedByHost = raw.costCountedByHost;
	if (typeof raw.note === "string") rule.note = raw.note;

	// A rule that changes nothing is a config mistake, not a no-op we should hide.
	if (!rule.body && !rule.headers && !rule.betaHeader) return undefined;
	return rule;
}

function normalizeRules(raw: unknown): Rule[] | undefined {
	if (!Array.isArray(raw)) return undefined;
	const rules = raw.map(normalizeRule).filter((r): r is Rule => r !== undefined);
	return rules.length > 0 ? rules : undefined;
}

function loadConfig(): Config {
	const path = join(getAgentDir(), "extensions", `${EXTENSION_ID}.json`);
	if (!existsSync(path)) return { ...DEFAULT_CONFIG };

	let raw: unknown;
	try {
		raw = JSON.parse(readFileSync(path, "utf-8"));
	} catch (error) {
		console.error(`[${EXTENSION_ID}] could not parse ${path}: ${error}`);
		return { ...DEFAULT_CONFIG };
	}
	if (!isRecord(raw)) return { ...DEFAULT_CONFIG };

	return {
		enabled: typeof raw.enabled === "boolean" ? raw.enabled : DEFAULT_CONFIG.enabled,
		indicator: INDICATORS.includes(raw.indicator as Indicator)
			? (raw.indicator as Indicator)
			: DEFAULT_CONFIG.indicator,
		indicatorText:
			typeof raw.indicatorText === "string" && raw.indicatorText.trim().length > 0
				? raw.indicatorText
				: DEFAULT_CONFIG.indicatorText,
		indicatorColor:
			typeof raw.indicatorColor === "string" && raw.indicatorColor.length > 0
				? raw.indicatorColor
				: DEFAULT_CONFIG.indicatorColor,
		warnAboveUsd:
			typeof raw.warnAboveUsd === "number" && raw.warnAboveUsd >= 0
				? raw.warnAboveUsd
				: DEFAULT_CONFIG.warnAboveUsd,
		rules: normalizeRules(raw.rules),
		extraRules: normalizeRules(raw.extraRules),
	};
}

function resolveRules(config: Config): Rule[] {
	return [...(config.rules ?? DEFAULT_RULES), ...(config.extraRules ?? [])];
}

// --------------------------------------------------------------------------- matching

function isFastEnabled(state: SessionState): boolean {
	if (state.once) return true;
	if (state.override === "on") return true;
	if (state.override === "off") return false;
	return state.config.enabled;
}

function describeMode(state: SessionState): string {
	if (state.once) return "armed for this turn";
	if (state.override === "on") return "on (session override)";
	if (state.override === "off") return "off (session override)";
	return state.config.enabled ? "on (config default)" : "off (config default)";
}

function findRule(ctx: ExtensionContext, state: SessionState): Match {
	const model = ctx.model;
	if (!model) return { modelKey: "no-model", reason: "no model is selected" };

	const modelKey = `${model.provider}/${model.id}`;
	const providerRules = state.rules.filter((rule) => rule.provider === model.provider);
	if (providerRules.length === 0) {
		return { modelKey, reason: `no fast-mode rule for provider ${model.provider}` };
	}

	const rule = providerRules.find(
		(r) =>
			(r.api === undefined || r.api === model.api) &&
			(r.models.includes(model.id) || r.models.includes("*") || r.models.includes(modelKey)),
	);
	if (!rule) {
		return { modelKey, reason: `${model.provider} has fast mode, but not for ${model.id}` };
	}
	return { rule, modelKey };
}

// --------------------------------------------------------------------------- spend

function newSpend(): Spend {
	return { requests: 0, billedUsd: 0, extraUsd: 0, underReported: false };
}

/**
 * Fold one fast turn into the running total.
 *
 * `reportedUsd` is what pi costed the turn at. For OpenAI, pi-ai reads
 * `service_tier` back off the response and has already applied the premium, so
 * that number is the real bill. For Anthropic it has never heard of `speed`,
 * so the real bill is `reported * multiplier` and pi's own display is short by
 * the difference.
 */
function recordSpend(spend: Spend, rule: Rule, reportedUsd: number): void {
	const multiplier = rule.costMultiplier ?? DEFAULT_COST_MULTIPLIER;
	const billed = rule.costCountedByHost ? reportedUsd : reportedUsd * multiplier;
	spend.requests += 1;
	spend.billedUsd += billed;
	spend.extraUsd += billed - billed / multiplier;
	if (!rule.costCountedByHost && billed > reportedUsd) spend.underReported = true;
}

function formatUsd(value: number): string {
	if (value <= 0) return "$0";
	if (value < 0.01) return "<$0.01";
	return `$${value.toFixed(2)}`;
}

function describeSpend(spend: Spend): string {
	if (spend.requests === 0) return "";
	const turns = `${spend.requests} fast ${spend.requests === 1 ? "request" : "requests"}`;
	const money = `${formatUsd(spend.billedUsd)} billed, ~${formatUsd(spend.extraUsd)} of that the fast premium`;
	const caveat = spend.underReported
		? ` pi's own cost display is short by ~${formatUsd(spend.extraUsd)} — it does not price Anthropic's speed tier.`
		: "";
	return ` ${turns}: ${money}.${caveat}`;
}

// --------------------------------------------------------------------------- rule validation

/**
 * Model ids a rule names that the registry has never heard of. A typo here is
 * the worst failure mode — fast silently does nothing while you believe it is on.
 */
function unknownModels(ctx: ExtensionContext, rule: Rule): string[] {
	return rule.models.filter((id) => {
		if (id === "*") return false;
		const [provider, modelId] = id.includes("/") ? id.split("/", 2) : [rule.provider, id];
		return ctx.modelRegistry.find(provider, modelId) === undefined;
	});
}

// --------------------------------------------------------------------------- headers

function splitBeta(value: string | undefined): string[] {
	return (value ?? "")
		.split(",")
		.map((part) => part.trim())
		.filter(Boolean);
}

/**
 * pi-ai builds `anthropic-beta` inside the provider client and merges
 * `model.headers` *over* it, so writing here wins outright — which also means
 * we have to re-add the Claude Code betas on OAuth accounts or we break
 * subscription auth.
 */
function syncHeaders(ctx: ExtensionContext, state: SessionState): void {
	const model = ctx.model as (typeof ctx.model & { headers?: Record<string, string> }) | undefined;
	if (!model) return;

	const { rule } = findRule(ctx, state);
	const active = isFastEnabled(state) && rule !== undefined;
	const headers = { ...(model.headers ?? {}) };

	// Undo whatever a previous rule/model left behind before re-applying.
	for (const known of state.rules) {
		for (const name of Object.keys(known.headers ?? {})) delete headers[name];
	}
	const betas = splitBeta(headers["anthropic-beta"] ?? headers["Anthropic-Beta"]);
	delete headers["Anthropic-Beta"];
	const knownBetas = new Set(state.rules.map((r) => r.betaHeader).filter(Boolean) as string[]);
	let nextBetas = betas.filter((b) => !knownBetas.has(b));

	if (active && rule) {
		Object.assign(headers, rule.headers ?? {});
		if (rule.betaHeader) {
			const oauth = ctx.modelRegistry.isUsingOAuth(model) ? ANTHROPIC_OAUTH_BETAS : [];
			nextBetas = [...new Set([...nextBetas, ...oauth, rule.betaHeader])];
		}
	}

	if (nextBetas.length > 0) headers["anthropic-beta"] = nextBetas.join(",");
	else delete headers["anthropic-beta"];
	model.headers = headers;
}

/**
 * Colour the marker. A `#rrggbb` value becomes a truecolor SGR sequence — host
 * footers pass status text through verbatim — and anything else is handed to
 * the theme, so `warning` and friends still track the active theme.
 */
function colorize(ctx: ExtensionContext, text: string, color: string): string {
	const hex = HEX_COLOR.exec(color);
	if (!hex) return ctx.ui.theme.fg(color as Parameters<typeof ctx.ui.theme.fg>[0], text);

	const value = Number.parseInt(hex[1], 16);
	const r = (value >> 16) & 0xff;
	const g = (value >> 8) & 0xff;
	const b = value & 0xff;
	return `\u001b[38;2;${r};${g};${b}m${text}\u001b[39m`;
}

function updateStatus(ctx: ExtensionContext, state: SessionState): void {
	syncHeaders(ctx, state);
	if (!ctx.hasUI) return;

	const indicator = state.config.indicator;
	const active = isFastEnabled(state) && findRule(ctx, state).rule !== undefined;
	const showStatus = active && (indicator === "status" || indicator === "both");
	const showWidget = active && (indicator === "widget" || indicator === "both");

	const hot =
		state.config.warnAboveUsd > 0 && state.spend.extraUsd >= state.config.warnAboveUsd;
	const marker =
		showStatus || showWidget
			? colorize(
					ctx,
					state.config.indicatorText,
					hot ? WARN_INDICATOR_COLOR : state.config.indicatorColor,
				)
			: undefined;

	ctx.ui.setStatus(STATUS_KEY, showStatus ? marker : undefined);
	ctx.ui.setWidget(WIDGET_KEY, showWidget && marker ? [marker] : undefined);
}

function statusMessage(ctx: ExtensionContext, state: SessionState): string {
	const enabled = isFastEnabled(state);
	const { rule, modelKey, reason } = findRule(ctx, state);
	const spend = describeSpend(state.spend);

	if (enabled && rule) {
		const patch = Object.entries(rule.body ?? {})
			.map(([k, v]) => `${k}=${JSON.stringify(v)}`)
			.join(" ");
		const note = rule.note ? ` \u2014 ${rule.note}` : "";
		const missing = unknownModels(ctx, rule);
		const warn = missing.length > 0 ? ` Unknown models in this rule: ${missing.join(", ")}.` : "";
		return `Fast is ${describeMode(state)} and active for ${modelKey}: ${patch || "headers only"}${note}.${warn}${spend}`;
	}
	if (enabled) {
		return `Fast is ${describeMode(state)} but inactive for ${modelKey}: ${reason}.${spend}`;
	}
	return `Fast is ${describeMode(state)}. Current model: ${modelKey}.${spend}`;
}

// --------------------------------------------------------------------------- extension

export default function piFastAll(pi: ExtensionAPI) {
	const states = new WeakMap<object, SessionState>();

	pi.registerFlag("fast", {
		description: "Start the session with provider fast mode enabled",
		type: "boolean",
		default: false,
	});

	function newState(): SessionState {
		const config = loadConfig();
		return {
			config,
			rules: resolveRules(config),
			override: pi.getFlag("fast") ? "on" : "auto",
			once: false,
			spend: newSpend(),
		};
	}

	function getState(ctx: ExtensionContext): SessionState {
		let state = states.get(ctx.sessionManager);
		if (!state) {
			state = newState();
			states.set(ctx.sessionManager, state);
		}
		return state;
	}

	pi.on("session_start", (_event, ctx) => {
		const state = newState();
		states.set(ctx.sessionManager, state);
		updateStatus(ctx, state);
	});

	pi.on("model_select", (_event, ctx) => {
		updateStatus(ctx, getState(ctx));
	});

	// Cost lands on the assistant message, so settle the pending request there.
	pi.on("message_end", (event, ctx) => {
		const state = getState(ctx);
		const rule = state.pending;
		if (!rule || event.message.role !== "assistant") return;
		state.pending = undefined;
		recordSpend(state.spend, rule, event.message.usage?.cost?.total ?? 0);
		updateStatus(ctx, state);
	});

	// `/fast once` covers a whole turn including its tool loop, so disarm only
	// when pi has settled and will not continue on its own.
	pi.on("agent_settled", (_event, ctx) => {
		const state = getState(ctx);
		state.pending = undefined;
		if (!state.once) return;
		state.once = false;
		updateStatus(ctx, state);
	});

	pi.on("before_provider_request", (event, ctx) => {
		const state = getState(ctx);
		updateStatus(ctx, state);

		if (!isFastEnabled(state)) return undefined;
		const { rule } = findRule(ctx, state);
		if (!rule?.body) return undefined;
		if (!isRecord(event.payload)) return undefined;
		if (event.payload.model !== ctx.model?.id) return undefined;

		// If any key is already set, someone else owns this knob. Leave it alone.
		const patch: Record<string, unknown> = {};
		for (const [key, value] of Object.entries(rule.body)) {
			if (key in event.payload) return undefined;
			patch[key] = value;
		}

		state.lastAppliedAt = Date.now();
		state.pending = rule;
		return { ...event.payload, ...patch };
	});

	pi.registerCommand("fast", {
		description: "Toggle provider fast mode (Anthropic speed, OpenAI priority tier)",
		getArgumentCompletions: (prefix: string) => {
			const items = ["on", "off", "once", "auto", "status", "rules"]
				.filter((v) => v.startsWith(prefix))
				.map((v) => ({ value: v, label: v }));
			return items.length > 0 ? items : null;
		},
		handler: async (args, ctx) => {
			const state = getState(ctx);
			const action = args.trim().toLowerCase();

			switch (action) {
				case "":
					state.override = isFastEnabled(state) ? "off" : "on";
					state.once = false;
					break;
				case "on":
				case "off":
				case "auto":
					state.once = false;
					state.override = action;
					break;
				case "once":
					state.once = true;
					break;
				case "status":
					break;
				case "rules": {
					let problems = 0;
					const lines = state.rules.map((r) => {
						const missing = new Set(unknownModels(ctx, r));
						problems += missing.size;
						const models = r.models.map((m) => (missing.has(m) ? `${m} (unknown?)` : m));
						const patch = Object.entries(r.body ?? {})
							.map(([k, v]) => `${k}=${JSON.stringify(v)}`)
							.join(" ");
						const beta = r.betaHeader ? ` +beta ${r.betaHeader}` : "";
						return `${r.provider}${r.api ? `/${r.api}` : ""}: ${models.join(", ")} → ${patch}${beta}`;
					});
					if (problems > 0) {
						lines.push(
							"",
							"(unknown?) = no such model in the registry — fast would silently do nothing",
						);
					}
					ctx.ui.notify(lines.join("\n"), problems > 0 ? "warning" : "info");
					return;
				}
				default:
					ctx.ui.notify("Usage: /fast [on|off|once|auto|status|rules]", "warning");
					return;
			}

			updateStatus(ctx, state);
			ctx.ui.notify(statusMessage(ctx, state), "info");
		},
	});
}