package net.minecraft.util.datafix.fixes;

import com.mojang.serialization.DataResult;
import com.mojang.datafixers.types.Type;
import com.mojang.datafixers.util.Pair;
import com.mojang.datafixers.Typed;
import com.mojang.datafixers.DSL;
import com.mojang.datafixers.TypeRewriteRule;
import java.util.Optional;
import java.util.stream.Stream;
import com.mojang.datafixers.DataFixUtils;
import com.mojang.serialization.Dynamic;
import com.mojang.datafixers.schemas.Schema;
import com.mojang.datafixers.DataFix;

public class MobSpawnerEntityIdentifiersFix extends DataFix {
    public MobSpawnerEntityIdentifiersFix(final Schema schema, final boolean boolean2) {
        super(schema, boolean2);
    }
    
    private Dynamic<?> fix(Dynamic<?> dynamic) {
        if (!"MobSpawner".equals(dynamic.get("id").asString(""))) {
            return dynamic;
        }
        final Optional<String> optional3 = (Optional<String>)dynamic.get("EntityId").asString().result();
        if (optional3.isPresent()) {
            Dynamic<?> dynamic2 = DataFixUtils.orElse(dynamic.get("SpawnData").result(), dynamic.emptyMap());
            dynamic2 = dynamic2.set("id", dynamic2.createString(optional3.get().isEmpty() ? "Pig" : optional3.get()));
            dynamic = dynamic.set("SpawnData", (Dynamic)dynamic2);
            dynamic = dynamic.remove("EntityId");
        }
        final Optional<? extends Stream<? extends Dynamic<?>>> optional4 = dynamic.get("SpawnPotentials").asStreamOpt().result();
        if (optional4.isPresent()) {
            final Optional<String> optional5;
            Dynamic<?> dynamic3;
            dynamic = dynamic.set("SpawnPotentials", dynamic.createList((Stream)((Stream)optional4.get()).map(dynamic -> {
                optional5 = (Optional<String>)dynamic.get("Type").asString().result();
                if (optional5.isPresent()) {
                    dynamic3 = ((Dynamic)DataFixUtils.orElse(dynamic.get("Properties").result(), dynamic.emptyMap())).set("id", dynamic.createString((String)optional5.get()));
                    return dynamic.set("Entity", (Dynamic)dynamic3).remove("Type").remove("Properties");
                }
                else {
                    return dynamic;
                }
            })));
        }
        return dynamic;
    }
    
    public TypeRewriteRule makeRule() {
        final Type<?> type2 = this.getOutputSchema().getType(References.UNTAGGED_SPAWNER);
        final Dynamic<?> dynamic4;
        final Dynamic<?> dynamic5;
        final Type type3;
        final DataResult<? extends Pair<? extends Typed<?>, ?>> dataResult5;
        return this.fixTypeEverywhereTyped("MobSpawnerEntityIdentifiersFix", this.getInputSchema().getType(References.UNTAGGED_SPAWNER), (Type)type2, typed -> {
            dynamic4 = typed.get(DSL.remainderFinder());
            dynamic5 = dynamic4.set("id", dynamic4.createString("MobSpawner"));
            dataResult5 = type3.readTyped((Dynamic)this.fix(dynamic5));
            if (!dataResult5.result().isPresent()) {
                return typed;
            }
            else {
                return (Typed)dataResult5.result().get().getFirst();
            }
        });
    }
}
