package net.minecraft.util.datafix.fixes;

import com.mojang.datafixers.DSL;
import com.mojang.datafixers.DataFix;
import com.mojang.datafixers.DataFixUtils;
import com.mojang.datafixers.TypeRewriteRule;
import com.mojang.datafixers.schemas.Schema;
import com.mojang.serialization.Dynamic;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public class DebugProfileLookingAtSplitFix extends DataFix {
	public DebugProfileLookingAtSplitFix(final Schema outputSchema) {
		super(outputSchema, false);
	}

	@Override
	public TypeRewriteRule makeRule() {
		return this.fixTypeEverywhereTyped(
			"DebugProfileLookingAtSplitFix",
			this.getInputSchema().getType(References.DEBUG_PROFILE),
			input -> input.update(DSL.remainderFinder(), tag -> tag.update("custom", DebugProfileLookingAtSplitFix::updateOptions))
		);
	}

	private static <T> Dynamic<T> updateOptions(final Dynamic<T> custom) {
		return DataFixUtils.orElse(custom.getMapValues().map(map -> {
			Map<Dynamic<T>, Dynamic<T>> newOptions = new HashMap();
			map.forEach((key, value) -> renamedKey(key).ifPresentOrElse(newKey -> newOptions.putIfAbsent(newKey, value), () -> newOptions.put(key, value)));
			return custom.createMap(newOptions);
		}).result(), custom);
	}

	private static <T> Optional<Dynamic<T>> renamedKey(final Dynamic<T> keyDynamic) {
		return keyDynamic.asString().result().flatMap(key -> {
			return switch (key) {
				case "minecraft:looking_at_block" -> Optional.of("minecraft:looking_at_block_state");
				case "minecraft:looking_at_fluid" -> Optional.of("minecraft:looking_at_fluid_state");
				default -> Optional.empty();
			};
		}).map(keyDynamic::createString);
	}
}