minecraft-src/net/minecraft/server/commands/ScheduleCommand.java

package net.minecraft.server.commands;

import com.mojang.brigadier.CommandDispatcher;
import com.mojang.brigadier.arguments.IntegerArgumentType;
import com.mojang.brigadier.arguments.StringArgumentType;
import com.mojang.brigadier.exceptions.CommandSyntaxException;
import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
import com.mojang.brigadier.exceptions.SimpleCommandExceptionType;
import com.mojang.brigadier.suggestion.SuggestionProvider;
import com.mojang.datafixers.util.Either;
import com.mojang.datafixers.util.Pair;
import java.util.Collection;
import java.util.Optional;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.Commands;
import net.minecraft.commands.SharedSuggestionProvider;
import net.minecraft.commands.arguments.TimeArgument;
import net.minecraft.commands.arguments.item.FunctionArgument;
import net.minecraft.commands.functions.CommandFunction;
import net.minecraft.commands.functions.MacroFunction;
import net.minecraft.network.chat.Component;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.MinecraftServer;
import net.minecraft.world.level.timers.FunctionCallback;
import net.minecraft.world.level.timers.FunctionTagCallback;
import net.minecraft.world.level.timers.TimerQueue;

public class ScheduleCommand {
	private static final SimpleCommandExceptionType ERROR_SAME_TICK = new SimpleCommandExceptionType(Component.translatable("commands.schedule.same_tick"));
	private static final DynamicCommandExceptionType ERROR_CANT_REMOVE = new DynamicCommandExceptionType(
		object -> Component.translatableEscape("commands.schedule.cleared.failure", object)
	);
	private static final SimpleCommandExceptionType ERROR_MACRO = new SimpleCommandExceptionType(Component.translatableEscape("commands.schedule.macro"));
	private static final SuggestionProvider<CommandSourceStack> SUGGEST_SCHEDULE = (commandContext, suggestionsBuilder) -> SharedSuggestionProvider.suggest(
		commandContext.getSource().getServer().getWorldData().overworldData().getScheduledEvents().getEventsIds(), suggestionsBuilder
	);

	public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
		dispatcher.register(
			Commands.literal("schedule")
				.requires(Commands.hasPermission(2))
				.then(
					Commands.literal("function")
						.then(
							Commands.argument("function", FunctionArgument.functions())
								.suggests(FunctionCommand.SUGGEST_FUNCTION)
								.then(
									Commands.argument("time", TimeArgument.time())
										.executes(
											commandContext -> schedule(
												commandContext.getSource(),
												FunctionArgument.getFunctionOrTag(commandContext, "function"),
												IntegerArgumentType.getInteger(commandContext, "time"),
												true
											)
										)
										.then(
											Commands.literal("append")
												.executes(
													commandContext -> schedule(
														commandContext.getSource(),
														FunctionArgument.getFunctionOrTag(commandContext, "function"),
														IntegerArgumentType.getInteger(commandContext, "time"),
														false
													)
												)
										)
										.then(
											Commands.literal("replace")
												.executes(
													commandContext -> schedule(
														commandContext.getSource(),
														FunctionArgument.getFunctionOrTag(commandContext, "function"),
														IntegerArgumentType.getInteger(commandContext, "time"),
														true
													)
												)
										)
								)
						)
				)
				.then(
					Commands.literal("clear")
						.then(
							Commands.argument("function", StringArgumentType.greedyString())
								.suggests(SUGGEST_SCHEDULE)
								.executes(commandContext -> remove(commandContext.getSource(), StringArgumentType.getString(commandContext, "function")))
						)
				)
		);
	}

	private static int schedule(
		CommandSourceStack source,
		Pair<ResourceLocation, Either<CommandFunction<CommandSourceStack>, Collection<CommandFunction<CommandSourceStack>>>> function,
		int time,
		boolean append
	) throws CommandSyntaxException {
		if (time == 0) {
			throw ERROR_SAME_TICK.create();
		} else {
			long l = source.getLevel().getGameTime() + time;
			ResourceLocation resourceLocation = function.getFirst();
			TimerQueue<MinecraftServer> timerQueue = source.getServer().getWorldData().overworldData().getScheduledEvents();
			Optional<CommandFunction<CommandSourceStack>> optional = function.getSecond().left();
			if (optional.isPresent()) {
				if (optional.get() instanceof MacroFunction) {
					throw ERROR_MACRO.create();
				}

				String string = resourceLocation.toString();
				if (append) {
					timerQueue.remove(string);
				}

				timerQueue.schedule(string, l, new FunctionCallback(resourceLocation));
				source.sendSuccess(() -> Component.translatable("commands.schedule.created.function", Component.translationArg(resourceLocation), time, l), true);
			} else {
				String string = "#" + resourceLocation;
				if (append) {
					timerQueue.remove(string);
				}

				timerQueue.schedule(string, l, new FunctionTagCallback(resourceLocation));
				source.sendSuccess(() -> Component.translatable("commands.schedule.created.tag", Component.translationArg(resourceLocation), time, l), true);
			}

			return Math.floorMod(l, Integer.MAX_VALUE);
		}
	}

	private static int remove(CommandSourceStack source, String function) throws CommandSyntaxException {
		int i = source.getServer().getWorldData().overworldData().getScheduledEvents().remove(function);
		if (i == 0) {
			throw ERROR_CANT_REMOVE.create(function);
		} else {
			source.sendSuccess(() -> Component.translatable("commands.schedule.cleared.success", i, function), true);
			return i;
		}
	}
}