import { shortHash } from "../utils/hash.js";

const MAX_TOOL_ARGUMENT_KEY_LENGTH = 256;
const LONG_KEY_PREFIX = "tool_arg_";

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

function clampKey(base: string, maxLength: number): string {
	return base.length <= maxLength ? base : base.slice(0, maxLength);
}

function makeSafeKey(
	key: string,
	maxLength: number,
	usedKeys: Set<string>,
	seen: Map<string, string>,
): string {
	if (key.length <= maxLength && !usedKeys.has(key)) {
		return key;
	}

	if (usedKeys.has(key)) {
		const base = `${LONG_KEY_PREFIX}${shortHash(key)}`;
		const safeBase = clampKey(base, maxLength);
		let next = 0;
		let candidate = safeBase;
		while (usedKeys.has(candidate)) {
			candidate = clampKey(`${safeBase}_${next}`, maxLength);
			next += 1;
		}
		seen.set(key, candidate);
		return candidate;
	}

	const existing = seen.get(key);
	if (existing) {
		let next = 0;
		let candidate = existing;
		while (usedKeys.has(candidate)) {
			candidate = clampKey(`${existing}_${next}`, maxLength);
			next += 1;
		}
		return candidate;
	}

	const base = `${LONG_KEY_PREFIX}${shortHash(key)}`;
	const safeBase = clampKey(base, maxLength);
	let next = 0;
	let candidate = safeBase;
	while (usedKeys.has(candidate)) {
		candidate = clampKey(`${safeBase}_${next}`, maxLength);
		next += 1;
	}
	seen.set(key, candidate);
	return candidate;
}

export function sanitizeToolCallArgumentsForSerialization(
	args: unknown,
	maxKeyLength = MAX_TOOL_ARGUMENT_KEY_LENGTH,
): unknown {
	if (isObject(args)) {
		const output: Record<string, unknown> = {};
		const usedKeys = new Set<string>();
		const replacements = new Map<string, string>();

		for (const [key, value] of Object.entries(args)) {
			const safeKey = makeSafeKey(key, maxKeyLength, usedKeys, replacements);
			output[safeKey] = sanitizeToolCallArgumentsForSerialization(
				value,
				maxKeyLength,
			);
			usedKeys.add(safeKey);
		}
		return output;
	}

	if (Array.isArray(args)) {
		return args.map((entry) =>
			sanitizeToolCallArgumentsForSerialization(entry, maxKeyLength),
		);
	}

	return args;
}
