SorceryCraft/src/main/java/com/thebrokenrail/sorcerycraft/command/SpellCommand.java

package com.thebrokenrail.sorcerycraft.command;

import com.mojang.brigadier.CommandDispatcher;
import com.mojang.brigadier.arguments.IntegerArgumentType;
import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
import com.thebrokenrail.sorcerycraft.SorceryCraft;
import com.thebrokenrail.sorcerycraft.spell.api.Spell;
import com.thebrokenrail.sorcerycraft.spell.api.registry.SpellRegistry;
import com.thebrokenrail.sorcerycraft.spell.util.SpellHelper;
import com.thebrokenrail.sorcerycraft.spell.util.SpellPlayerEntity;
import net.minecraft.command.arguments.EntityArgumentType;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.item.ItemStack;
import net.minecraft.server.command.CommandManager;
import net.minecraft.server.command.ServerCommandSource;
import net.minecraft.server.network.ServerPlayerEntity;
import net.minecraft.text.Texts;
import net.minecraft.text.TranslatableText;
import net.minecraft.util.Identifier;

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

public class SpellCommand {
    private static final DynamicCommandExceptionType NOT_HOLDING_SPELL_EXCEPTION = new DynamicCommandExceptionType(object -> new TranslatableText("command."+ SorceryCraft.NAMESPACE + ".spell.not_holding_spell", object));

    public static void register(CommandDispatcher<ServerCommandSource> dispatcher) {
        dispatcher.register(CommandManager.literal("spell")
                .requires(source -> source.hasPermissionLevel(4))
                .then(CommandManager.literal("list")
                        .then(CommandManager.argument("player", EntityArgumentType.players())
                            .executes(ctx -> {
                                Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                for (PlayerEntity player : players) {
                                    SpellPlayerEntity spellPlayer = (SpellPlayerEntity) player;
                                    Map<Identifier, Integer> spellMap = spellPlayer.getDiscoveredSpells();
                                    ctx.getSource().sendFeedback(new TranslatableText("command." + SorceryCraft.NAMESPACE + ".spell.listing_spells", player.getDisplayName(), Texts.join(spellMap.entrySet(), spell -> SpellHelper.getTranslatedSpell(spell.getKey(), spell.getValue()))), false);
                                }
                                return 0;
                            })
                        )
                )
                .then(CommandManager.literal("forget")
                        .then(CommandManager.argument("player", EntityArgumentType.players())
                                .executes(ctx -> {
                                    Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                    for (PlayerEntity player : players) {
                                        SpellPlayerEntity spellPlayer = (SpellPlayerEntity) player;
                                        Map<Identifier, Integer> spells = spellPlayer.getDiscoveredSpells();
                                        for (Map.Entry<Identifier, Integer> entry : spells.entrySet()) {
                                            ctx.getSource().sendFeedback(new TranslatableText("command." + SorceryCraft.NAMESPACE + ".spell.forgotten_spell", player.getDisplayName(), SpellHelper.getTranslatedSpellChat(entry.getKey(), entry.getValue())), true);
                                        }
                                        spellPlayer.setDiscoveredSpells(new HashMap<>());
                                    }
                                    return 1;
                                })
                                .then(CommandManager.argument("spell", SpellArgumentType.spell())
                                        .executes(ctx -> {
                                            Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                            for (PlayerEntity player : players) {
                                                Identifier spell = SpellArgumentType.getSpell(ctx, "spell");
                                                SpellPlayerEntity spellPlayer = (SpellPlayerEntity) player;
                                                Map<Identifier, Integer> spells = spellPlayer.getDiscoveredSpells();
                                                if (spells.containsKey(spell)) {
                                                    ctx.getSource().sendFeedback(new TranslatableText("command." + SorceryCraft.NAMESPACE + ".spell.forgotten_spell", player.getDisplayName(), SpellHelper.getTranslatedSpellChat(spell, spells.get(spell))), true);
                                                    spells.remove(spell);
                                                }
                                                spellPlayer.setDiscoveredSpells(spells);
                                            }
                                            return 1;
                                        })
                                )
                        )
                )
                .then(CommandManager.literal("discover")
                        .then(CommandManager.argument("player", EntityArgumentType.players())
                                .executes(ctx -> {
                                    Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                    for (PlayerEntity player : players) {
                                        Map<Identifier, Integer> spellMap = new HashMap<>();
                                        Spell[] maxSpells = SpellRegistry.getMaxSpells();
                                        for (Spell spell : maxSpells) {
                                            spellMap.put(spell.getID(), spell.getLevel());
                                        }
                                        SpellHelper.learnSpells(player, spellMap);
                                    }
                                    return 1;
                                })
                                .then(CommandManager.argument("spell", SpellArgumentType.spell())
                                        .then(CommandManager.argument("level", IntegerArgumentType.integer())
                                                .executes(ctx -> {
                                                    Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                                    for (PlayerEntity player : players) {
                                                        Identifier spell = SpellArgumentType.getSpell(ctx, "spell");
                                                        int level = IntegerArgumentType.getInteger(ctx, "level") - 1;
                                                        Map<Identifier, Integer> spellMap = new HashMap<>();
                                                        spellMap.put(spell, level);
                                                        SpellHelper.learnSpells(player, spellMap);
                                                    }
                                                    return 1;
                                                })
                                        )
                                )
                        )
                )
                .then(CommandManager.literal("apply")
                        .then(CommandManager.argument("player", EntityArgumentType.players())
                                .then(CommandManager.argument("spell", SpellArgumentType.spell())
                                        .then(CommandManager.argument("level", IntegerArgumentType.integer())
                                                .executes(ctx -> {
                                                    Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                                    for (PlayerEntity player : players) {
                                                        Identifier spell = SpellArgumentType.getSpell(ctx, "spell");
                                                        int level = IntegerArgumentType.getInteger(ctx, "level") - 1;

                                                        ItemStack stack = player.getMainHandStack();

                                                        if (stack.getItem() != SorceryCraft.SPELL_ITEM) {
                                                            throw NOT_HOLDING_SPELL_EXCEPTION.create(player);
                                                        }

                                                        Map<Identifier, Integer> spellMap = SpellHelper.getSpells(stack);
                                                        spellMap.put(spell, level);
                                                        SpellHelper.setSpells(stack, spellMap);

                                                        ctx.getSource().sendFeedback(new TranslatableText("command." + SorceryCraft.NAMESPACE + ".spell.applied_spell", SpellHelper.getTranslatedSpell(spell, level)), true);
                                                    }
                                                    return 1;
                                                })
                                        )
                                )
                        )
                )
                .then(CommandManager.literal("remove")
                        .then(CommandManager.argument("player", EntityArgumentType.players())
                                .then(CommandManager.argument("spell", SpellArgumentType.spell())
                                        .executes(ctx -> {
                                            Collection<ServerPlayerEntity> players = EntityArgumentType.getPlayers(ctx, "player");
                                            for (PlayerEntity player : players) {
                                                Identifier spell = SpellArgumentType.getSpell(ctx, "spell");

                                                ItemStack stack = player.getMainHandStack();

                                                if (stack.getItem() != SorceryCraft.SPELL_ITEM) {
                                                    throw NOT_HOLDING_SPELL_EXCEPTION.create(player);
                                                }

                                                Map<Identifier, Integer> spellMap = SpellHelper.getSpells(stack);
                                                if (spellMap.containsKey(spell)) {
                                                    ctx.getSource().sendFeedback(new TranslatableText("command." + SorceryCraft.NAMESPACE + ".spell.removed_spell", SpellHelper.getTranslatedSpell(spell, spellMap.get(spell))), true);
                                                    spellMap.remove(spell);
                                                }
                                                SpellHelper.setSpells(stack, spellMap);
                                            }
                                            return 1;
                                        })
                                )
                        )
                )
        );
    }
}