SorceryCraft/src/main/java/com/thebrokenrail/sorcerycraft/mixin/MixinPlayerEntity.java

package com.thebrokenrail.sorcerycraft.mixin;

import com.thebrokenrail.sorcerycraft.gui.CastingTableScreenHandler;
import com.thebrokenrail.sorcerycraft.spell.util.SpellHelper;
import com.thebrokenrail.sorcerycraft.spell.util.SpellPlayerEntity;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.screen.ScreenHandler;
import net.minecraft.util.Identifier;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.Unique;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;

import java.util.HashMap;
import java.util.Map;

@SuppressWarnings("unused")
@Mixin(PlayerEntity.class)
public class MixinPlayerEntity implements SpellPlayerEntity {
    @Shadow
    public ScreenHandler currentScreenHandler;
    @Unique
    private Map<Identifier, Integer> discoveredSpells = new HashMap<>();

    @Inject(at = @At("HEAD"), method = "readCustomDataFromTag")
    public void readCustomDataFromTag(CompoundTag tag, CallbackInfo info) {
        discoveredSpells = SpellHelper.getSpells(tag);
    }

    @Inject(at = @At("HEAD"), method = "writeCustomDataToTag")
    public void writeCustomDataToTag(CompoundTag tag, CallbackInfo info) {
        tag.put(SpellHelper.SPELL_TAG, SpellHelper.createSpellsTag(discoveredSpells));
    }

    @Override
    public void setDiscoveredSpells(Map<Identifier, Integer> spells) {
        discoveredSpells = spells;
        if (currentScreenHandler instanceof CastingTableScreenHandler) {
            //noinspection ConstantConditions
            ((CastingTableScreenHandler) currentScreenHandler).setSpells((PlayerEntity) (Object) this);
        }
    }

    @Override
    public Map<Identifier, Integer> getDiscoveredSpells() {
        return discoveredSpells;
    }
}