import { type TSchema, Type } from "typebox";

type JsonSchema = {
	type?: string | string[];
	properties?: Record<string, JsonSchema>;
	items?: JsonSchema;
	$ref?: string;
	$defs?: Record<string, JsonSchema>;
	anyOf?: JsonSchema[];
	allOf?: JsonSchema[];
	format?: string;
};

function primaryType(schema: JsonSchema | undefined): string | undefined {
	if (!schema?.type) return undefined;
	if (typeof schema.type === "string") return schema.type;
	return schema.type.find((t) => t !== "null");
}

function isObjectSchema(schema: JsonSchema | undefined): boolean {
	if (!schema) return false;
	if (primaryType(schema) === "object") return true;
	return Boolean(schema.properties);
}

function isArraySchema(schema: JsonSchema | undefined): boolean {
	if (!schema) return false;
	if (primaryType(schema) === "array") return true;
	return Boolean(schema.items);
}

function resolveRef(schema: JsonSchema, root: JsonSchema): JsonSchema {
	const ref = schema.$ref;
	if (!ref?.startsWith("#/$defs/")) {
		return schema;
	}
	const defName = ref.slice("#/$defs/".length);
	const def = root.$defs?.[defName];
	if (!def) {
		return schema;
	}
	return def;
}

function resolveSchema(schema: JsonSchema | undefined, root: JsonSchema): JsonSchema | undefined {
	if (!schema) return undefined;
	if (schema.$ref) {
		return resolveSchema(resolveRef(schema, root), root);
	}
	return schema;
}

function parseJsonString(value: string): unknown {
	const trimmed = value.trim();
	if (!trimmed.startsWith("{") && !trimmed.startsWith("[")) {
		return value;
	}
	try {
		return JSON.parse(trimmed) as unknown;
	} catch {
		return value;
	}
}

function normalizeDateTime(value: unknown): unknown {
	if (typeof value !== "string") {
		return value;
	}
	// LLMs often emit +0000; JSON Schema date-time expects +00:00.
	return value.replace(/([+-]\d{2})(\d{2})$/, "$1:$2");
}

function normalizeForSchema(
	value: unknown,
	schema: JsonSchema | undefined,
	root: JsonSchema,
): unknown {
	const resolved = resolveSchema(schema, root);
	if (!resolved) {
		return value;
	}

	if (resolved.format === "date-time" && typeof value === "string") {
		value = normalizeDateTime(value);
	}

	if (Array.isArray(resolved.anyOf)) {
		const arrayBranch = resolved.anyOf.find((branch) => isArraySchema(branch));
		if (
			arrayBranch &&
			!Array.isArray(value) &&
			value !== null &&
			value !== undefined &&
			(typeof value === "number" || typeof value === "string" || typeof value === "boolean")
		) {
			return [normalizeForSchema(value, arrayBranch.items, root)];
		}
	}

	if (
		isArraySchema(resolved) &&
		!Array.isArray(value) &&
		value !== null &&
		value !== undefined &&
		(typeof value === "number" || typeof value === "string" || typeof value === "boolean")
	) {
		return [normalizeForSchema(value, resolved.items, root)];
	}

	return value;
}

function coerceValue(value: unknown, schema: JsonSchema | undefined, root: JsonSchema): unknown {
	const resolved = resolveSchema(schema, root);

	if (typeof value === "string" && (isObjectSchema(resolved) || isArraySchema(resolved))) {
		value = parseJsonString(value);
	}

	value = normalizeForSchema(value, schema, root);

	if (
		isObjectSchema(resolved) &&
		value !== null &&
		typeof value === "object" &&
		!Array.isArray(value)
	) {
		const record = value as Record<string, unknown>;
		if (!resolved?.properties) return record;

		const coerced: Record<string, unknown> = {};
		for (const [key, nested] of Object.entries(record)) {
			coerced[key] = coerceValue(nested, resolved.properties[key], root);
		}
		return coerced;
	}

	if (isArraySchema(resolved) && Array.isArray(value) && resolved?.items) {
		return value.map((item) => coerceValue(item, resolved.items, root));
	}

	return value;
}

/** Coerce JSON-stringified nested objects before TypeBox validation. */
export function coerceMcpArguments(args: unknown, inputSchema: unknown): unknown {
	if (!args || typeof args !== "object" || Array.isArray(args)) {
		return args;
	}
	if (!inputSchema || typeof inputSchema !== "object") {
		return args;
	}
	const root = inputSchema as JsonSchema;
	return coerceValue(args, root, root);
}

function relaxPropertySchema(property: JsonSchema, root: JsonSchema): JsonSchema {
	const resolved = resolveSchema(property, root);
	if (!resolved || !(isObjectSchema(resolved) || isArraySchema(resolved))) {
		return property;
	}

	// Keep $ref intact so $defs stay at the schema root.
	if (property.$ref) {
		return { anyOf: [property, { type: "string" }] };
	}

	return { anyOf: [relaxNestedSchema(property, root), { type: "string" }] };
}

function stripFormats(schema: JsonSchema): JsonSchema {
	const result: JsonSchema = { ...schema };
	delete result.format;

	if (result.properties) {
		result.properties = Object.fromEntries(
			Object.entries(result.properties).map(([key, property]) => [
				key,
				stripFormats(property),
			]),
		);
	}

	if (result.items && !Array.isArray(result.items)) {
		result.items = stripFormats(result.items);
	}

	if (result.anyOf) {
		result.anyOf = result.anyOf.map((branch) => stripFormats(branch));
	}

	if (result.allOf) {
		result.allOf = result.allOf.map((branch) => stripFormats(branch));
	}

	if (result.$defs) {
		result.$defs = Object.fromEntries(
			Object.entries(result.$defs).map(([key, def]) => [key, stripFormats(def)]),
		);
	}

	return result;
}

function relaxNestedSchema(schema: JsonSchema, root: JsonSchema): JsonSchema {
	const result: JsonSchema = stripFormats({ ...schema });
	if (result.properties) {
		result.properties = Object.fromEntries(
			Object.entries(result.properties).map(([key, property]) => [
				key,
				relaxPropertySchema(property, root),
			]),
		);
	}
	if (result.items && !Array.isArray(result.items)) {
		result.items = relaxPropertySchema(result.items, root);
	}
	return result;
}

function relaxSchema(schema: JsonSchema, root: JsonSchema): JsonSchema {
	const result: JsonSchema = stripFormats({
		...schema,
		properties: schema.properties
			? Object.fromEntries(
					Object.entries(schema.properties).map(([key, property]) => [
						key,
						relaxPropertySchema(property, root),
					]),
				)
			: schema.properties,
	});
	if (root.$defs) {
		result.$defs = Object.fromEntries(
			Object.entries(root.$defs).map(([key, def]) => [key, stripFormats(def)]),
		);
	}
	return result;
}

export function mcpSchemaToTypeBox(inputSchema: unknown): TSchema {
	if (!inputSchema || typeof inputSchema !== "object") {
		return Type.Object({});
	}
	return Type.Unsafe(relaxSchema(inputSchema as JsonSchema, inputSchema as JsonSchema));
}