package net.minecraft.util.datafix.fixes;

import com.mojang.datafixers.TypeRewriteRule;
import com.mojang.datafixers.schemas.Schema;
import com.mojang.serialization.Dynamic;

public class BlockEntityUUIDFix extends AbstractUUIDFix {
	public BlockEntityUUIDFix(Schema outputSchema) {
		super(outputSchema, References.BLOCK_ENTITY);
	}

	@Override
	protected TypeRewriteRule makeRule() {
		return this.fixTypeEverywhereTyped("BlockEntityUUIDFix", this.getInputSchema().getType(this.typeReference), typed -> {
			typed = this.updateNamedChoice(typed, "minecraft:conduit", this::updateConduit);
			return this.updateNamedChoice(typed, "minecraft:skull", this::updateSkull);
		});
	}

	private Dynamic<?> updateSkull(Dynamic<?> skullTag) {
		return (Dynamic<?>)skullTag.get("Owner")
			.get()
			.map(dynamic -> (Dynamic)replaceUUIDString(dynamic, "Id", "Id").orElse(dynamic))
			.map(dynamic2 -> skullTag.remove("Owner").set("SkullOwner", dynamic2))
			.result()
			.orElse(skullTag);
	}

	private Dynamic<?> updateConduit(Dynamic<?> conduitTag) {
		return (Dynamic<?>)replaceUUIDMLTag(conduitTag, "target_uuid", "Target").orElse(conduitTag);
	}
}