import * as z from "zod/mini";

export type ProtocolJsonValue =
	| string
	| number
	| boolean
	| null
	| ProtocolJsonValue[]
	| { [key: string]: ProtocolJsonValue };

export type ProtocolJsonObject = Record<string, ProtocolJsonValue>;

export interface ProtocolJsonBounds {
	maxCollectionSize: number;
	maxDepth: number;
	maxKeyLength: number;
	maxStringLength: number;
	maxTotalNodes: number;
}

const protocolJsonNumberSchema = z.number().check(
	z.refine(Number.isFinite, {
		message: "Expected a finite JSON number.",
	}),
);

export const jsonObjectWithoutPrototypeKeySchema = z.unknown().check(
	z.refine(
		(value) =>
			!(
				typeof value === "object" &&
				value !== null &&
				Object.hasOwn(value, "__proto__")
			),
		{
			message: 'JSON objects cannot contain an own "__proto__" key.',
		},
	),
);

export const protocolJsonValueSchema: z.ZodMiniType<ProtocolJsonValue> = z.lazy(
	() =>
		z.union([
			z.string(),
			protocolJsonNumberSchema,
			z.boolean(),
			z.null(),
			z.array(protocolJsonValueSchema),
			z.pipe(
				jsonObjectWithoutPrototypeKeySchema,
				z.record(z.string(), protocolJsonValueSchema),
			),
		]),
);

export const protocolJsonObjectSchema = z.pipe(
	jsonObjectWithoutPrototypeKeySchema,
	z.record(z.string(), protocolJsonValueSchema),
);

const hasAtMostJsonNodes = (
	value: ProtocolJsonValue,
	maxTotalNodes: number,
): boolean => {
	const pending: ProtocolJsonValue[] = [value];
	let visited = 0;

	while (pending.length > 0) {
		visited += 1;
		if (visited > maxTotalNodes) {
			return false;
		}

		const current = pending.pop();
		if (Array.isArray(current)) {
			if (visited + pending.length + current.length > maxTotalNodes) {
				return false;
			}
			for (const child of current) {
				pending.push(child);
			}
		} else if (current !== null && typeof current === "object") {
			const children = Object.values(current);
			if (visited + pending.length + children.length > maxTotalNodes) {
				return false;
			}
			for (const child of children) {
				pending.push(child);
			}
		}
	}

	return true;
};

const createBoundedProtocolJsonValueSchema = (
	bounds: ProtocolJsonBounds,
	remainingDepth: number,
): z.ZodMiniType<ProtocolJsonValue> => {
	const scalarSchema = z.union([
		z.string().check(z.maxLength(bounds.maxStringLength)),
		protocolJsonNumberSchema,
		z.boolean(),
		z.null(),
	]);

	if (remainingDepth <= 0) {
		return scalarSchema;
	}

	const childSchema = createBoundedProtocolJsonValueSchema(
		bounds,
		remainingDepth - 1,
	);
	const objectSchema = z.pipe(
		jsonObjectWithoutPrototypeKeySchema,
		z
			.record(z.string().check(z.maxLength(bounds.maxKeyLength)), childSchema)
			.check(
				z.refine(
					(value) => Object.keys(value).length <= bounds.maxCollectionSize,
					{
						message: `Expected at most ${bounds.maxCollectionSize} object keys.`,
					},
				),
			),
	);

	return z.union([
		scalarSchema,
		z.array(childSchema).check(z.maxLength(bounds.maxCollectionSize)),
		objectSchema,
	]);
};

export const createBoundedProtocolJsonObjectSchema = (
	bounds: ProtocolJsonBounds,
): z.ZodMiniType<ProtocolJsonObject> => {
	const valueSchema = createBoundedProtocolJsonValueSchema(
		bounds,
		bounds.maxDepth,
	);

	return z.pipe(
		jsonObjectWithoutPrototypeKeySchema,
		z
			.record(z.string().check(z.maxLength(bounds.maxKeyLength)), valueSchema)
			.check(
				z.refine(
					(value) => Object.keys(value).length <= bounds.maxCollectionSize,
					{
						message: `Expected at most ${bounds.maxCollectionSize} object keys.`,
					},
				),
				z.refine((value) => hasAtMostJsonNodes(value, bounds.maxTotalNodes), {
					message: `Expected at most ${bounds.maxTotalNodes} JSON nodes.`,
				}),
			),
	);
};