package net.minecraft.client.gui.screen.world;

import com.mojang.datafixers.DataFixer;
import com.mojang.logging.LogUtils;
import it.unimi.dsi.fastutil.booleans.BooleanConsumer;
import it.unimi.dsi.fastutil.objects.Object2IntMap;
import it.unimi.dsi.fastutil.objects.Object2IntOpenCustomHashMap;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.minecraft.client.MinecraftClient;
import net.minecraft.client.gui.DrawContext;
import net.minecraft.client.gui.screen.Screen;
import net.minecraft.client.gui.widget.ButtonWidget;
import net.minecraft.registry.Registry;
import net.minecraft.registry.RegistryKey;
import net.minecraft.registry.RegistryKeys;
import net.minecraft.registry.DynamicRegistryManager.Immutable;
import net.minecraft.screen.ScreenTexts;
import net.minecraft.server.SaveLoader;
import net.minecraft.text.Text;
import net.minecraft.util.Util;
import net.minecraft.util.math.MathHelper;
import net.minecraft.world.SaveProperties;
import net.minecraft.world.World;
import net.minecraft.world.dimension.DimensionOptions;
import net.minecraft.world.level.LevelInfo;
import net.minecraft.world.level.storage.LevelStorage.Session;
import net.minecraft.world.updater.WorldUpdater;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

@Environment(EnvType.CLIENT)
public class OptimizeWorldScreen extends Screen {
	private static final Logger LOGGER = LogUtils.getLogger();
	private static final Object2IntMap<RegistryKey<World>> DIMENSION_COLORS = (Object2IntMap<RegistryKey<World>>)Util.make(
		new Object2IntOpenCustomHashMap(Util.identityHashStrategy()), colors -> {
			colors.put(World.OVERWORLD, -13408734);
			colors.put(World.NETHER, -10075085);
			colors.put(World.END, -8943531);
			colors.defaultReturnValue(-2236963);
		}
	);
	private final BooleanConsumer callback;
	private final WorldUpdater updater;

	@Nullable
	public static OptimizeWorldScreen create(MinecraftClient client, BooleanConsumer callback, DataFixer dataFixer, Session storageSession, boolean eraseCache) {
		try {
			SaveLoader saveLoader = client.createIntegratedServerLoader().createSaveLoader(storageSession, false);

			OptimizeWorldScreen var8;
			try {
				SaveProperties saveProperties = saveLoader.saveProperties();
				Immutable immutable = saveLoader.combinedDynamicRegistries().getCombinedRegistryManager();
				storageSession.backupLevelDataFile(immutable, saveProperties);
				var8 = new OptimizeWorldScreen(callback, dataFixer, storageSession, saveProperties.getLevelInfo(), eraseCache, immutable.get(RegistryKeys.DIMENSION));
			} catch (Throwable var10) {
				if (saveLoader != null) {
					try {
						saveLoader.close();
					} catch (Throwable var9) {
						var10.addSuppressed(var9);
					}
				}

				throw var10;
			}

			if (saveLoader != null) {
				saveLoader.close();
			}

			return var8;
		} catch (Exception var11) {
			LOGGER.warn("Failed to load datapacks, can't optimize world", (Throwable)var11);
			return null;
		}
	}

	private OptimizeWorldScreen(
		BooleanConsumer callback,
		DataFixer dataFixer,
		Session storageSession,
		LevelInfo levelInfo,
		boolean eraseCache,
		Registry<DimensionOptions> dimensionOptionsRegistry
	) {
		super(Text.translatable("optimizeWorld.title", new Object[]{levelInfo.getLevelName()}));
		this.callback = callback;
		this.updater = new WorldUpdater(storageSession, dataFixer, dimensionOptionsRegistry, eraseCache);
	}

	@Override
	protected void init() {
		super.init();
		this.addDrawableChild(ButtonWidget.builder(ScreenTexts.CANCEL, button -> {
			this.updater.cancel();
			this.callback.accept(false);
		}).dimensions(this.width / 2 - 100, this.height / 4 + 150, 200, 20).build());
	}

	@Override
	public void tick() {
		if (this.updater.isDone()) {
			this.callback.accept(true);
		}
	}

	@Override
	public void close() {
		this.callback.accept(false);
	}

	@Override
	public void removed() {
		this.updater.cancel();
	}

	@Override
	public void render(DrawContext context, int mouseX, int mouseY, float delta) {
		this.renderBackground(context);
		context.drawCenteredTextWithShadow(this.textRenderer, this.title, this.width / 2, 20, 16777215);
		int i = this.width / 2 - 150;
		int j = this.width / 2 + 150;
		int k = this.height / 4 + 100;
		int l = k + 10;
		context.drawCenteredTextWithShadow(this.textRenderer, this.updater.getStatus(), this.width / 2, k - 9 - 2, 10526880);
		if (this.updater.getTotalChunkCount() > 0) {
			context.fill(i - 1, k - 1, j + 1, l + 1, -16777216);
			context.drawTextWithShadow(
				this.textRenderer, Text.translatable("optimizeWorld.info.converted", new Object[]{this.updater.getUpgradedChunkCount()}), i, 40, 10526880
			);
			context.drawTextWithShadow(
				this.textRenderer, Text.translatable("optimizeWorld.info.skipped", new Object[]{this.updater.getSkippedChunkCount()}), i, 40 + 9 + 3, 10526880
			);
			context.drawTextWithShadow(
				this.textRenderer, Text.translatable("optimizeWorld.info.total", new Object[]{this.updater.getTotalChunkCount()}), i, 40 + (9 + 3) * 2, 10526880
			);
			int m = 0;

			for (RegistryKey<World> registryKey : this.updater.getWorlds()) {
				int n = MathHelper.floor(this.updater.getProgress(registryKey) * (j - i));
				context.fill(i + m, k, i + m + n, l, DIMENSION_COLORS.getInt(registryKey));
				m += n;
			}

			int o = this.updater.getUpgradedChunkCount() + this.updater.getSkippedChunkCount();
			context.drawCenteredTextWithShadow(this.textRenderer, o + " / " + this.updater.getTotalChunkCount(), this.width / 2, k + 2 * 9 + 2, 10526880);
			context.drawCenteredTextWithShadow(
				this.textRenderer, MathHelper.floor(this.updater.getProgress() * 100.0F) + "%", this.width / 2, k + (l - k) / 2 - 9 / 2, 10526880
			);
		}

		super.render(context, mouseX, mouseY, delta);
	}
}
