import { dbg } from "./debug.js";
import type { KlausRequest } from "./protocol.js";

const SUBSTITUTIONS = [
	["asks about pi itself", "asks about this harness"],
	["pi packages", "harness packages"],
] as const;
const PI_SYSTEM_PROMPT_MARKERS = [
	"You are an expert coding assistant operating inside pi",
	"Pi documentation",
] as const;

export function rewriteSubscriptionPrompt(
	request: KlausRequest,
	warn: (message: string) => void = () => undefined,
): KlausRequest {
	let systemPrompt = request.systemPrompt;
	const isPiSystemPrompt = PI_SYSTEM_PROMPT_MARKERS.some((marker) =>
		systemPrompt.includes(marker),
	);
	const substitutions: Array<(typeof SUBSTITUTIONS)[number]> = [];
	const missing: string[] = [];
	for (const substitution of SUBSTITUTIONS) {
		const [before, after] = substitution;
		const occurrences = systemPrompt.split(before).length - 1;
		if (occurrences > 1) {
			throw new Error(`Klaus found repeated ${before}.`);
		}
		if (occurrences === 1) {
			systemPrompt = systemPrompt.replace(before, after);
			substitutions.push(substitution);
		} else {
			missing.push(before);
		}
	}
	if (missing.length > 0 && isPiSystemPrompt) {
		warn(
			`Klaus did not find expected Pi system-prompt phrase${missing.length === 1 ? "" : "s"}: ${missing.map((phrase) => JSON.stringify(phrase)).join(", ")}. Subscription routing may fail after upstream prompt changes.`,
		);
	}
	dbg?.("subscriptionPrompt.rewrite", { isPiSystemPrompt });
	return { ...request, systemPrompt };
}