package net.minecraft.server.commands;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.mojang.brigadier.CommandDispatcher;
import com.mojang.brigadier.arguments.BoolArgumentType;
import com.mojang.brigadier.arguments.FloatArgumentType;
import com.mojang.brigadier.arguments.IntegerArgumentType;
import com.mojang.brigadier.exceptions.CommandSyntaxException;
import com.mojang.brigadier.exceptions.Dynamic2CommandExceptionType;
import com.mojang.brigadier.exceptions.Dynamic4CommandExceptionType;
import java.util.Collection;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.Commands;
import net.minecraft.commands.arguments.EntityArgument;
import net.minecraft.commands.arguments.coordinates.Vec2Argument;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Direction;
import net.minecraft.network.chat.Component;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.tags.BlockTags;
import net.minecraft.util.Mth;
import net.minecraft.util.RandomSource;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.player.Player;
import net.minecraft.world.level.BlockGetter;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.phys.Vec2;
import net.minecraft.world.scores.Team;

public class SpreadPlayersCommand {
	private static final int MAX_ITERATION_COUNT = 10000;
	private static final Dynamic4CommandExceptionType ERROR_FAILED_TO_SPREAD_TEAMS = new Dynamic4CommandExceptionType(
		(object, object2, object3, object4) -> Component.translatableEscape("commands.spreadplayers.failed.teams", object, object2, object3, object4)
	);
	private static final Dynamic4CommandExceptionType ERROR_FAILED_TO_SPREAD_ENTITIES = new Dynamic4CommandExceptionType(
		(object, object2, object3, object4) -> Component.translatableEscape("commands.spreadplayers.failed.entities", object, object2, object3, object4)
	);
	private static final Dynamic2CommandExceptionType ERROR_INVALID_MAX_HEIGHT = new Dynamic2CommandExceptionType(
		(object, object2) -> Component.translatableEscape("commands.spreadplayers.failed.invalid.height", object, object2)
	);

	public static void register(CommandDispatcher<CommandSourceStack> dispatcher) {
		dispatcher.register(
			Commands.literal("spreadplayers")
				.requires(commandSourceStack -> commandSourceStack.hasPermission(2))
				.then(
					Commands.argument("center", Vec2Argument.vec2())
						.then(
							Commands.argument("spreadDistance", FloatArgumentType.floatArg(0.0F))
								.then(
									Commands.argument("maxRange", FloatArgumentType.floatArg(1.0F))
										.then(
											Commands.argument("respectTeams", BoolArgumentType.bool())
												.then(
													Commands.argument("targets", EntityArgument.entities())
														.executes(
															commandContext -> spreadPlayers(
																commandContext.getSource(),
																Vec2Argument.getVec2(commandContext, "center"),
																FloatArgumentType.getFloat(commandContext, "spreadDistance"),
																FloatArgumentType.getFloat(commandContext, "maxRange"),
																commandContext.getSource().getLevel().getMaxY() + 1,
																BoolArgumentType.getBool(commandContext, "respectTeams"),
																EntityArgument.getEntities(commandContext, "targets")
															)
														)
												)
										)
										.then(
											Commands.literal("under")
												.then(
													Commands.argument("maxHeight", IntegerArgumentType.integer())
														.then(
															Commands.argument("respectTeams", BoolArgumentType.bool())
																.then(
																	Commands.argument("targets", EntityArgument.entities())
																		.executes(
																			commandContext -> spreadPlayers(
																				commandContext.getSource(),
																				Vec2Argument.getVec2(commandContext, "center"),
																				FloatArgumentType.getFloat(commandContext, "spreadDistance"),
																				FloatArgumentType.getFloat(commandContext, "maxRange"),
																				IntegerArgumentType.getInteger(commandContext, "maxHeight"),
																				BoolArgumentType.getBool(commandContext, "respectTeams"),
																				EntityArgument.getEntities(commandContext, "targets")
																			)
																		)
																)
														)
												)
										)
								)
						)
				)
		);
	}

	private static int spreadPlayers(
		CommandSourceStack source, Vec2 center, float spreadDistance, float maxRange, int maxHeight, boolean respectTeams, Collection<? extends Entity> targets
	) throws CommandSyntaxException {
		ServerLevel serverLevel = source.getLevel();
		int i = serverLevel.getMinY();
		if (maxHeight < i) {
			throw ERROR_INVALID_MAX_HEIGHT.create(maxHeight, i);
		} else {
			RandomSource randomSource = RandomSource.create();
			double d = center.x - maxRange;
			double e = center.y - maxRange;
			double f = center.x + maxRange;
			double g = center.y + maxRange;
			SpreadPlayersCommand.Position[] positions = createInitialPositions(randomSource, respectTeams ? getNumberOfTeams(targets) : targets.size(), d, e, f, g);
			spreadPositions(center, spreadDistance, serverLevel, randomSource, d, e, f, g, maxHeight, positions, respectTeams);
			double h = setPlayerPositions(targets, serverLevel, positions, maxHeight, respectTeams);
			source.sendSuccess(
				() -> Component.translatable(
					"commands.spreadplayers.success." + (respectTeams ? "teams" : "entities"), positions.length, center.x, center.y, String.format(Locale.ROOT, "%.2f", h)
				),
				true
			);
			return positions.length;
		}
	}

	/**
	 * Gets the number of unique teams for the given list of entities.
	 */
	private static int getNumberOfTeams(Collection<? extends Entity> entities) {
		Set<Team> set = Sets.<Team>newHashSet();

		for (Entity entity : entities) {
			if (entity instanceof Player) {
				set.add(entity.getTeam());
			} else {
				set.add(null);
			}
		}

		return set.size();
	}

	private static void spreadPositions(
		Vec2 center,
		double spreadDistance,
		ServerLevel level,
		RandomSource random,
		double minX,
		double minZ,
		double maxX,
		double maxZ,
		int maxHeight,
		SpreadPlayersCommand.Position[] positions,
		boolean respectTeams
	) throws CommandSyntaxException {
		boolean bl = true;
		double d = Float.MAX_VALUE;

		int i;
		for (i = 0; i < 10000 && bl; i++) {
			bl = false;
			d = Float.MAX_VALUE;

			for (int j = 0; j < positions.length; j++) {
				SpreadPlayersCommand.Position position = positions[j];
				int k = 0;
				SpreadPlayersCommand.Position position2 = new SpreadPlayersCommand.Position();

				for (int l = 0; l < positions.length; l++) {
					if (j != l) {
						SpreadPlayersCommand.Position position3 = positions[l];
						double e = position.dist(position3);
						d = Math.min(e, d);
						if (e < spreadDistance) {
							k++;
							position2.x = position2.x + (position3.x - position.x);
							position2.z = position2.z + (position3.z - position.z);
						}
					}
				}

				if (k > 0) {
					position2.x /= k;
					position2.z /= k;
					double f = position2.getLength();
					if (f > 0.0) {
						position2.normalize();
						position.moveAway(position2);
					} else {
						position.randomize(random, minX, minZ, maxX, maxZ);
					}

					bl = true;
				}

				if (position.clamp(minX, minZ, maxX, maxZ)) {
					bl = true;
				}
			}

			if (!bl) {
				for (SpreadPlayersCommand.Position position2 : positions) {
					if (!position2.isSafe(level, maxHeight)) {
						position2.randomize(random, minX, minZ, maxX, maxZ);
						bl = true;
					}
				}
			}
		}

		if (d == Float.MAX_VALUE) {
			d = 0.0;
		}

		if (i >= 10000) {
			if (respectTeams) {
				throw ERROR_FAILED_TO_SPREAD_TEAMS.create(positions.length, center.x, center.y, String.format(Locale.ROOT, "%.2f", d));
			} else {
				throw ERROR_FAILED_TO_SPREAD_ENTITIES.create(positions.length, center.x, center.y, String.format(Locale.ROOT, "%.2f", d));
			}
		}
	}

	private static double setPlayerPositions(
		Collection<? extends Entity> targets, ServerLevel level, SpreadPlayersCommand.Position[] positions, int maxHeight, boolean respectTeams
	) {
		double d = 0.0;
		int i = 0;
		Map<Team, SpreadPlayersCommand.Position> map = Maps.<Team, SpreadPlayersCommand.Position>newHashMap();

		for (Entity entity : targets) {
			SpreadPlayersCommand.Position position;
			if (respectTeams) {
				Team team = entity instanceof Player ? entity.getTeam() : null;
				if (!map.containsKey(team)) {
					map.put(team, positions[i++]);
				}

				position = (SpreadPlayersCommand.Position)map.get(team);
			} else {
				position = positions[i++];
			}

			entity.teleportTo(
				level, Mth.floor(position.x) + 0.5, position.getSpawnY(level, maxHeight), Mth.floor(position.z) + 0.5, Set.of(), entity.getYRot(), entity.getXRot(), true
			);
			double e = Double.MAX_VALUE;

			for (SpreadPlayersCommand.Position position2 : positions) {
				if (position != position2) {
					double f = position.dist(position2);
					e = Math.min(f, e);
				}
			}

			d += e;
		}

		return targets.size() < 2 ? 0.0 : d / targets.size();
	}

	private static SpreadPlayersCommand.Position[] createInitialPositions(RandomSource random, int count, double minX, double minZ, double maxX, double maxZ) {
		SpreadPlayersCommand.Position[] positions = new SpreadPlayersCommand.Position[count];

		for (int i = 0; i < positions.length; i++) {
			SpreadPlayersCommand.Position position = new SpreadPlayersCommand.Position();
			position.randomize(random, minX, minZ, maxX, maxZ);
			positions[i] = position;
		}

		return positions;
	}

	static class Position {
		double x;
		double z;

		double dist(SpreadPlayersCommand.Position other) {
			double d = this.x - other.x;
			double e = this.z - other.z;
			return Math.sqrt(d * d + e * e);
		}

		void normalize() {
			double d = this.getLength();
			this.x /= d;
			this.z /= d;
		}

		double getLength() {
			return Math.sqrt(this.x * this.x + this.z * this.z);
		}

		public void moveAway(SpreadPlayersCommand.Position other) {
			this.x = this.x - other.x;
			this.z = this.z - other.z;
		}

		public boolean clamp(double minX, double minZ, double maxX, double maxZ) {
			boolean bl = false;
			if (this.x < minX) {
				this.x = minX;
				bl = true;
			} else if (this.x > maxX) {
				this.x = maxX;
				bl = true;
			}

			if (this.z < minZ) {
				this.z = minZ;
				bl = true;
			} else if (this.z > maxZ) {
				this.z = maxZ;
				bl = true;
			}

			return bl;
		}

		public int getSpawnY(BlockGetter level, int y) {
			BlockPos.MutableBlockPos mutableBlockPos = new BlockPos.MutableBlockPos(this.x, (double)(y + 1), this.z);
			boolean bl = level.getBlockState(mutableBlockPos).isAir();
			mutableBlockPos.move(Direction.DOWN);
			boolean bl2 = level.getBlockState(mutableBlockPos).isAir();

			while (mutableBlockPos.getY() > level.getMinY()) {
				mutableBlockPos.move(Direction.DOWN);
				boolean bl3 = level.getBlockState(mutableBlockPos).isAir();
				if (!bl3 && bl2 && bl) {
					return mutableBlockPos.getY() + 1;
				}

				bl = bl2;
				bl2 = bl3;
			}

			return y + 1;
		}

		public boolean isSafe(BlockGetter level, int y) {
			BlockPos blockPos = BlockPos.containing(this.x, this.getSpawnY(level, y) - 1, this.z);
			BlockState blockState = level.getBlockState(blockPos);
			return blockPos.getY() < y && !blockState.liquid() && !blockState.is(BlockTags.FIRE);
		}

		public void randomize(RandomSource random, double minX, double minZ, double maxX, double maxZ) {
			this.x = Mth.nextDouble(random, minX, maxX);
			this.z = Mth.nextDouble(random, minZ, maxZ);
		}
	}
}