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

export const DEFAULT_THRESHOLDS = [3, 5, 8] as const;

export interface RepeatToolGuardOptions {
	/** Consecutive repeat counts that should receive a reminder. Defaults to 3, 5, and 8. */
	thresholds?: readonly number[];
	/** Tool name patterns to track. If omitted, all tools are tracked. Supports `*` wildcards. */
	include?: readonly string[];
	/** Tool name patterns to ignore. Supports `*` wildcards. */
	exclude?: readonly string[];
	/** Environment source for PI_REPEAT_TOOL_GUARD_* variables. Defaults to process.env. */
	env?: NodeJS.ProcessEnv;
}

type ToolResultContent = ToolResultEvent["content"];

export interface ToolResultPatch {
	content?: ToolResultContent;
	details?: ToolResultEvent["details"];
	isError?: boolean;
}

const REMINDER_TEXT_1 =
	"\n\n<system-reminder>\n" +
	"You are repeating the exact same tool call with identical parameters. " +
	"Please carefully analyze the previous result. If the task is not yet complete, " +
	"try a different method or parameters instead of repeating the same call.\n" +
	"</system-reminder>";

export function makeReminderText(toolName: string, repeatCount: number, canonicalInput: string): string {
	return (
		"\n\n<system-reminder>\n" +
		"You have repeatedly called the same tool with identical parameters many times.\n" +
		"Repeated tool call detected:\n" +
		`- tool: ${toolName}\n` +
		`- repeated_times: ${repeatCount}\n` +
		`- arguments: ${canonicalInput}\n` +
		"The previous repeated calls did not make progress. Do not call this exact same tool " +
		"with the exact same arguments again.\n" +
		"Carefully inspect the latest tool result and choose a different next action, " +
		"different parameters, or finish the task if enough evidence has been gathered.\n" +
		"</system-reminder>"
	);
}

export function parseListEnv(name: string, env: NodeJS.ProcessEnv = process.env): string[] {
	return (env[name] ?? "")
		.split(",")
		.map((item) => item.trim())
		.filter(Boolean);
}

export function normalizeThresholds(values: readonly number[] | undefined): number[] {
	if (!values || values.length === 0) return [...DEFAULT_THRESHOLDS];
	const parsed = values.filter((item) => Number.isInteger(item) && item > 1);
	return parsed.length === 0 ? [...DEFAULT_THRESHOLDS] : [...new Set(parsed)].sort((a, b) => a - b);
}

export function parseThresholds(env: NodeJS.ProcessEnv = process.env): number[] {
	const raw = parseListEnv("PI_REPEAT_TOOL_GUARD_THRESHOLDS", env);
	if (raw.length === 0) return [...DEFAULT_THRESHOLDS];
	const parsed = raw.map((item) => Number.parseInt(item, 10));
	return normalizeThresholds(parsed);
}

export function wildcardPatternToRegExp(pattern: string): RegExp {
	const escaped = pattern.replace(/[|\\{}()[\]^$+?.]/g, "\\$&");
	return new RegExp(`^${escaped.replaceAll("*", ".*")}$`);
}

export function compilePatterns(patterns: readonly string[]): RegExp[] {
	return patterns.map(wildcardPatternToRegExp);
}

export function matchesAny(value: string, patterns: readonly RegExp[]): boolean {
	return patterns.some((pattern) => pattern.test(value));
}

export function sortJsonValue(value: unknown, seen: WeakSet<object> = new WeakSet<object>()): unknown {
	if (Array.isArray(value)) {
		return value.map((item) => sortJsonValue(item, seen));
	}

	if (value && typeof value === "object") {
		if (seen.has(value)) return "[Circular]";
		seen.add(value);

		const sorted: Record<string, unknown> = {};
		for (const key of Object.keys(value).sort()) {
			const item = (value as Record<string, unknown>)[key];
			if (item !== undefined) {
				sorted[key] = sortJsonValue(item, seen);
			}
		}
		seen.delete(value);
		return sorted;
	}

	if (typeof value === "bigint") return value.toString();
	return value;
}

export function canonicalizeInput(input: unknown): string {
	try {
		return JSON.stringify(sortJsonValue(input)) ?? String(input);
	} catch {
		return String(input);
	}
}

export function appendReminderToContent(content: ToolResultContent, reminderText: string): ToolResultContent {
	const next = content.map((item) => ({ ...item }));

	for (let index = next.length - 1; index >= 0; index--) {
		const item = next[index];
		if (item.type === "text") {
			next[index] = { ...item, text: item.text + reminderText };
			return next;
		}
	}

	next.push({ type: "text", text: reminderText });
	return next;
}

export function createRepeatToolGuard(options: RepeatToolGuardOptions = {}): (pi: ExtensionAPI) => void {
	const env = options.env ?? process.env;
	const includePatterns = compilePatterns(options.include ?? parseListEnv("PI_REPEAT_TOOL_GUARD_INCLUDE", env));
	const excludePatterns = compilePatterns(options.exclude ?? parseListEnv("PI_REPEAT_TOOL_GUARD_EXCLUDE", env));
	const thresholds = normalizeThresholds(options.thresholds ?? parseThresholds(env));
	const thresholdSet = new Set(thresholds);
	const pendingReminders = new Map<string, string>();

	let consecutiveKey: string | undefined;
	let consecutiveCount = 0;

	function shouldTrackTool(toolName: string): boolean {
		if (includePatterns.length > 0 && !matchesAny(toolName, includePatterns)) return false;
		if (matchesAny(toolName, excludePatterns)) return false;
		return true;
	}

	function reset(): void {
		consecutiveKey = undefined;
		consecutiveCount = 0;
		pendingReminders.clear();
	}

	function handleToolCall(event: ToolCallEvent): undefined {
		if (!shouldTrackTool(event.toolName)) return undefined;

		const canonicalInput = canonicalizeInput(event.input);
		const key = JSON.stringify([event.toolName, canonicalInput]);
		if (key === consecutiveKey) {
			consecutiveCount++;
		} else {
			consecutiveKey = key;
			consecutiveCount = 1;
		}

		if (thresholdSet.has(consecutiveCount)) {
			pendingReminders.set(
				event.toolCallId,
				consecutiveCount === 3 ? REMINDER_TEXT_1 : makeReminderText(event.toolName, consecutiveCount, canonicalInput),
			);
		}

		return undefined;
	}

	function handleToolResult(event: ToolResultEvent): ToolResultPatch | undefined {
		const reminderText = pendingReminders.get(event.toolCallId);
		if (!reminderText) return undefined;

		pendingReminders.delete(event.toolCallId);
		return {
			content: appendReminderToContent(event.content, reminderText),
		};
	}

	return (pi: ExtensionAPI): void => {
		pi.on("session_start", reset);

		pi.on("message_end", (event) => {
			if (event.message.role === "user") {
				reset();
			}
			return undefined;
		});

		pi.on("tool_call", handleToolCall);
		pi.on("tool_result", handleToolResult);
	};
}

export default function repeatToolGuard(pi: ExtensionAPI): void {
	createRepeatToolGuard()(pi);
}