package net.minecraft.util.datafix.fixes;

import java.util.Map;
import com.google.common.collect.ImmutableMap;
import com.mojang.datafixers.util.Pair;
import com.mojang.serialization.Dynamic;
import java.util.function.Function;
import com.mojang.datafixers.DSL;
import com.mojang.datafixers.Typed;
import com.mojang.datafixers.schemas.Schema;

public class MemoryExpiryDataFix extends NamedEntityFix {
    public MemoryExpiryDataFix(final Schema schema, final String string) {
        super(schema, false, "Memory expiry data fix (" + string + ")", References.ENTITY, string);
    }
    
    @Override
    protected Typed<?> fix(final Typed<?> typed) {
        return typed.update(DSL.remainderFinder(), (Function)this::fixTag);
    }
    
    public Dynamic<?> fixTag(final Dynamic<?> dynamic) {
        return dynamic.update("Brain", (Function)this::updateBrain);
    }
    
    private Dynamic<?> updateBrain(final Dynamic<?> dynamic) {
        return dynamic.update("memories", (Function)this::updateMemories);
    }
    
    private Dynamic<?> updateMemories(final Dynamic<?> dynamic) {
        return dynamic.updateMapValues((Function)this::updateMemoryEntry);
    }
    
    private Pair<Dynamic<?>, Dynamic<?>> updateMemoryEntry(final Pair<Dynamic<?>, Dynamic<?>> pair) {
        return (Pair<Dynamic<?>, Dynamic<?>>)pair.mapSecond((Function)this::wrapMemoryValue);
    }
    
    private Dynamic<?> wrapMemoryValue(final Dynamic<?> dynamic) {
        return dynamic.createMap((Map)ImmutableMap.<Dynamic, Dynamic<?>>of(dynamic.createString("value"), dynamic));
    }
}
