package net.minecraft.util.worldupdate;

import com.google.common.collect.Lists;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.mojang.datafixers.DataFixer;
import com.mojang.logging.LogUtils;
import it.unimi.dsi.fastutil.objects.Reference2FloatMap;
import it.unimi.dsi.fastutil.objects.Reference2FloatMaps;
import it.unimi.dsi.fastutil.objects.Reference2FloatOpenHashMap;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.ListIterator;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ThreadFactory;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import net.minecraft.ReportedException;
import net.minecraft.SharedConstants;
import net.minecraft.Util;
import net.minecraft.core.Registry;
import net.minecraft.core.RegistryAccess;
import net.minecraft.core.registries.Registries;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.network.chat.Component;
import net.minecraft.resources.ResourceKey;
import net.minecraft.util.datafix.DataFixTypes;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.chunk.ChunkGenerator;
import net.minecraft.world.level.chunk.storage.ChunkStorage;
import net.minecraft.world.level.chunk.storage.RecreatingChunkStorage;
import net.minecraft.world.level.chunk.storage.RecreatingSimpleRegionStorage;
import net.minecraft.world.level.chunk.storage.RegionFile;
import net.minecraft.world.level.chunk.storage.RegionStorageInfo;
import net.minecraft.world.level.chunk.storage.SimpleRegionStorage;
import net.minecraft.world.level.dimension.LevelStem;
import net.minecraft.world.level.saveddata.SavedData;
import net.minecraft.world.level.storage.DimensionDataStorage;
import net.minecraft.world.level.storage.LevelStorageSource;
import net.minecraft.world.level.storage.WorldData;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

public class WorldUpgrader implements AutoCloseable {
	static final Logger LOGGER = LogUtils.getLogger();
	private static final ThreadFactory THREAD_FACTORY = new ThreadFactoryBuilder().setDaemon(true).build();
	private static final String NEW_DIRECTORY_PREFIX = "new_";
	static final Component STATUS_UPGRADING_POI = Component.translatable("optimizeWorld.stage.upgrading.poi");
	static final Component STATUS_FINISHED_POI = Component.translatable("optimizeWorld.stage.finished.poi");
	static final Component STATUS_UPGRADING_ENTITIES = Component.translatable("optimizeWorld.stage.upgrading.entities");
	static final Component STATUS_FINISHED_ENTITIES = Component.translatable("optimizeWorld.stage.finished.entities");
	static final Component STATUS_UPGRADING_CHUNKS = Component.translatable("optimizeWorld.stage.upgrading.chunks");
	static final Component STATUS_FINISHED_CHUNKS = Component.translatable("optimizeWorld.stage.finished.chunks");
	final Registry<LevelStem> dimensions;
	final Set<ResourceKey<Level>> levels;
	final boolean eraseCache;
	final boolean recreateRegionFiles;
	final LevelStorageSource.LevelStorageAccess levelStorage;
	private final Thread thread;
	final DataFixer dataFixer;
	volatile boolean running = true;
	private volatile boolean finished;
	volatile float progress;
	volatile int totalChunks;
	volatile int totalFiles;
	volatile int converted;
	volatile int skipped;
	final Reference2FloatMap<ResourceKey<Level>> progressMap = Reference2FloatMaps.synchronize(new Reference2FloatOpenHashMap<>());
	volatile Component status = Component.translatable("optimizeWorld.stage.counting");
	static final Pattern REGEX = Pattern.compile("^r\\.(-?[0-9]+)\\.(-?[0-9]+)\\.mca$");
	final DimensionDataStorage overworldDataStorage;

	public WorldUpgrader(
		LevelStorageSource.LevelStorageAccess levelStorage,
		DataFixer dataFixer,
		WorldData worldData,
		RegistryAccess registryAccess,
		boolean eraseCache,
		boolean recreateRegionFiles
	) {
		this.dimensions = registryAccess.lookupOrThrow(Registries.LEVEL_STEM);
		this.levels = (Set<ResourceKey<Level>>)this.dimensions.registryKeySet().stream().map(Registries::levelStemToLevel).collect(Collectors.toUnmodifiableSet());
		this.eraseCache = eraseCache;
		this.dataFixer = dataFixer;
		this.levelStorage = levelStorage;
		SavedData.Context context = new SavedData.Context(null, worldData.worldGenOptions().seed());
		this.overworldDataStorage = new DimensionDataStorage(context, this.levelStorage.getDimensionPath(Level.OVERWORLD).resolve("data"), dataFixer, registryAccess);
		this.recreateRegionFiles = recreateRegionFiles;
		this.thread = THREAD_FACTORY.newThread(this::work);
		this.thread.setUncaughtExceptionHandler((thread, throwable) -> {
			LOGGER.error("Error upgrading world", throwable);
			this.status = Component.translatable("optimizeWorld.stage.failed");
			this.finished = true;
		});
		this.thread.start();
	}

	public void cancel() {
		this.running = false;

		try {
			this.thread.join();
		} catch (InterruptedException var2) {
		}
	}

	private void work() {
		long l = Util.getMillis();
		LOGGER.info("Upgrading entities");
		new WorldUpgrader.EntityUpgrader().upgrade();
		LOGGER.info("Upgrading POIs");
		new WorldUpgrader.PoiUpgrader().upgrade();
		LOGGER.info("Upgrading blocks");
		new WorldUpgrader.ChunkUpgrader().upgrade();
		this.overworldDataStorage.saveAndJoin();
		l = Util.getMillis() - l;
		LOGGER.info("World optimizaton finished after {} seconds", l / 1000L);
		this.finished = true;
	}

	public boolean isFinished() {
		return this.finished;
	}

	public Set<ResourceKey<Level>> levels() {
		return this.levels;
	}

	public float dimensionProgress(ResourceKey<Level> level) {
		return this.progressMap.getFloat(level);
	}

	public float getProgress() {
		return this.progress;
	}

	public int getTotalChunks() {
		return this.totalChunks;
	}

	public int getConverted() {
		return this.converted;
	}

	public int getSkipped() {
		return this.skipped;
	}

	public Component getStatus() {
		return this.status;
	}

	public void close() {
		this.overworldDataStorage.close();
	}

	static Path resolveRecreateDirectory(Path path) {
		return path.resolveSibling("new_" + path.getFileName().toString());
	}

	abstract class AbstractUpgrader<T extends AutoCloseable> {
		private final Component upgradingStatus;
		private final Component finishedStatus;
		private final String type;
		private final String folderName;
		@Nullable
		protected CompletableFuture<Void> previousWriteFuture;
		protected final DataFixTypes dataFixType;

		AbstractUpgrader(final DataFixTypes dataFixType, final String type, final String folderName, final Component upgradingStatus, final Component finishedStatus) {
			this.dataFixType = dataFixType;
			this.type = type;
			this.folderName = folderName;
			this.upgradingStatus = upgradingStatus;
			this.finishedStatus = finishedStatus;
		}

		public void upgrade() {
			WorldUpgrader.this.totalFiles = 0;
			WorldUpgrader.this.totalChunks = 0;
			WorldUpgrader.this.converted = 0;
			WorldUpgrader.this.skipped = 0;
			List<WorldUpgrader.DimensionToUpgrade<T>> list = this.getDimensionsToUpgrade();
			if (WorldUpgrader.this.totalChunks != 0) {
				float f = WorldUpgrader.this.totalFiles;
				WorldUpgrader.this.status = this.upgradingStatus;

				while (WorldUpgrader.this.running) {
					boolean bl = false;
					float g = 0.0F;

					for (WorldUpgrader.DimensionToUpgrade<T> dimensionToUpgrade : list) {
						ResourceKey<Level> resourceKey = dimensionToUpgrade.dimensionKey;
						ListIterator<WorldUpgrader.FileToUpgrade> listIterator = dimensionToUpgrade.files;
						T autoCloseable = dimensionToUpgrade.storage;
						if (listIterator.hasNext()) {
							WorldUpgrader.FileToUpgrade fileToUpgrade = (WorldUpgrader.FileToUpgrade)listIterator.next();
							boolean bl2 = true;

							for (ChunkPos chunkPos : fileToUpgrade.chunksToUpgrade) {
								bl2 = bl2 && this.processOnePosition(resourceKey, autoCloseable, chunkPos);
								bl = true;
							}

							if (WorldUpgrader.this.recreateRegionFiles) {
								if (bl2) {
									this.onFileFinished(fileToUpgrade.file);
								} else {
									WorldUpgrader.LOGGER.error("Failed to convert region file {}", fileToUpgrade.file.getPath());
								}
							}
						}

						float h = listIterator.nextIndex() / f;
						WorldUpgrader.this.progressMap.put(resourceKey, h);
						g += h;
					}

					WorldUpgrader.this.progress = g;
					if (!bl) {
						break;
					}
				}

				WorldUpgrader.this.status = this.finishedStatus;

				for (WorldUpgrader.DimensionToUpgrade<T> dimensionToUpgrade2 : list) {
					try {
						dimensionToUpgrade2.storage.close();
					} catch (Exception var14) {
						WorldUpgrader.LOGGER.error("Error upgrading chunk", (Throwable)var14);
					}
				}
			}
		}

		private List<WorldUpgrader.DimensionToUpgrade<T>> getDimensionsToUpgrade() {
			List<WorldUpgrader.DimensionToUpgrade<T>> list = Lists.<WorldUpgrader.DimensionToUpgrade<T>>newArrayList();

			for (ResourceKey<Level> resourceKey : WorldUpgrader.this.levels) {
				RegionStorageInfo regionStorageInfo = new RegionStorageInfo(WorldUpgrader.this.levelStorage.getLevelId(), resourceKey, this.type);
				Path path = WorldUpgrader.this.levelStorage.getDimensionPath(resourceKey).resolve(this.folderName);
				T autoCloseable = this.createStorage(regionStorageInfo, path);
				ListIterator<WorldUpgrader.FileToUpgrade> listIterator = this.getFilesToProcess(regionStorageInfo, path);
				list.add(new WorldUpgrader.DimensionToUpgrade(resourceKey, autoCloseable, listIterator));
			}

			return list;
		}

		protected abstract T createStorage(RegionStorageInfo regionStorageInfo, Path path);

		private ListIterator<WorldUpgrader.FileToUpgrade> getFilesToProcess(RegionStorageInfo regionStorageInfo, Path path) {
			List<WorldUpgrader.FileToUpgrade> list = getAllChunkPositions(regionStorageInfo, path);
			WorldUpgrader.this.totalFiles = WorldUpgrader.this.totalFiles + list.size();
			WorldUpgrader.this.totalChunks = WorldUpgrader.this.totalChunks + list.stream().mapToInt(fileToUpgrade -> fileToUpgrade.chunksToUpgrade.size()).sum();
			return list.listIterator();
		}

		private static List<WorldUpgrader.FileToUpgrade> getAllChunkPositions(RegionStorageInfo regionStorageInfo, Path path) {
			File[] files = path.toFile().listFiles((filex, string) -> string.endsWith(".mca"));
			if (files == null) {
				return List.of();
			} else {
				List<WorldUpgrader.FileToUpgrade> list = Lists.<WorldUpgrader.FileToUpgrade>newArrayList();

				for (File file : files) {
					Matcher matcher = WorldUpgrader.REGEX.matcher(file.getName());
					if (matcher.matches()) {
						int i = Integer.parseInt(matcher.group(1)) << 5;
						int j = Integer.parseInt(matcher.group(2)) << 5;
						List<ChunkPos> list2 = Lists.<ChunkPos>newArrayList();

						try (RegionFile regionFile = new RegionFile(regionStorageInfo, file.toPath(), path, true)) {
							for (int k = 0; k < 32; k++) {
								for (int l = 0; l < 32; l++) {
									ChunkPos chunkPos = new ChunkPos(k + i, l + j);
									if (regionFile.doesChunkExist(chunkPos)) {
										list2.add(chunkPos);
									}
								}
							}

							if (!list2.isEmpty()) {
								list.add(new WorldUpgrader.FileToUpgrade(regionFile, list2));
							}
						} catch (Throwable var18) {
							WorldUpgrader.LOGGER.error("Failed to read chunks from region file {}", file.toPath(), var18);
						}
					}
				}

				return list;
			}
		}

		private boolean processOnePosition(ResourceKey<Level> dimesion, T storage, ChunkPos chunkPos) {
			boolean bl = false;

			try {
				bl = this.tryProcessOnePosition(storage, chunkPos, dimesion);
			} catch (CompletionException | ReportedException var7) {
				Throwable throwable = var7.getCause();
				if (!(throwable instanceof IOException)) {
					throw var7;
				}

				WorldUpgrader.LOGGER.error("Error upgrading chunk {}", chunkPos, throwable);
			}

			if (bl) {
				WorldUpgrader.this.converted++;
			} else {
				WorldUpgrader.this.skipped++;
			}

			return bl;
		}

		protected abstract boolean tryProcessOnePosition(T chunkStorage, ChunkPos chunkPos, ResourceKey<Level> dimension);

		private void onFileFinished(RegionFile regionFile) {
			if (WorldUpgrader.this.recreateRegionFiles) {
				if (this.previousWriteFuture != null) {
					this.previousWriteFuture.join();
				}

				Path path = regionFile.getPath();
				Path path2 = path.getParent();
				Path path3 = WorldUpgrader.resolveRecreateDirectory(path2).resolve(path.getFileName().toString());

				try {
					if (path3.toFile().exists()) {
						Files.delete(path);
						Files.move(path3, path);
					} else {
						WorldUpgrader.LOGGER.error("Failed to replace an old region file. New file {} does not exist.", path3);
					}
				} catch (IOException var6) {
					WorldUpgrader.LOGGER.error("Failed to replace an old region file", (Throwable)var6);
				}
			}
		}
	}

	class ChunkUpgrader extends WorldUpgrader.AbstractUpgrader<ChunkStorage> {
		ChunkUpgrader() {
			super(DataFixTypes.CHUNK, "chunk", "region", WorldUpgrader.STATUS_UPGRADING_CHUNKS, WorldUpgrader.STATUS_FINISHED_CHUNKS);
		}

		protected boolean tryProcessOnePosition(ChunkStorage chunkStorage, ChunkPos chunkPos, ResourceKey<Level> resourceKey) {
			CompoundTag compoundTag = (CompoundTag)((Optional)chunkStorage.read(chunkPos).join()).orElse(null);
			if (compoundTag != null) {
				int i = ChunkStorage.getVersion(compoundTag);
				ChunkGenerator chunkGenerator = WorldUpgrader.this.dimensions.getValueOrThrow(Registries.levelToLevelStem(resourceKey)).generator();
				CompoundTag compoundTag2 = chunkStorage.upgradeChunkTag(
					resourceKey, () -> WorldUpgrader.this.overworldDataStorage, compoundTag, chunkGenerator.getTypeNameForDataFixer()
				);
				ChunkPos chunkPos2 = new ChunkPos(compoundTag2.getIntOr("xPos", 0), compoundTag2.getIntOr("zPos", 0));
				if (!chunkPos2.equals(chunkPos)) {
					WorldUpgrader.LOGGER.warn("Chunk {} has invalid position {}", chunkPos, chunkPos2);
				}

				boolean bl = i < SharedConstants.getCurrentVersion().getDataVersion().getVersion();
				if (WorldUpgrader.this.eraseCache) {
					bl = bl || compoundTag2.contains("Heightmaps");
					compoundTag2.remove("Heightmaps");
					bl = bl || compoundTag2.contains("isLightOn");
					compoundTag2.remove("isLightOn");
					ListTag listTag = compoundTag2.getListOrEmpty("sections");

					for (int j = 0; j < listTag.size(); j++) {
						Optional<CompoundTag> optional = listTag.getCompound(j);
						if (!optional.isEmpty()) {
							CompoundTag compoundTag3 = (CompoundTag)optional.get();
							bl = bl || compoundTag3.contains("BlockLight");
							compoundTag3.remove("BlockLight");
							bl = bl || compoundTag3.contains("SkyLight");
							compoundTag3.remove("SkyLight");
						}
					}
				}

				if (bl || WorldUpgrader.this.recreateRegionFiles) {
					if (this.previousWriteFuture != null) {
						this.previousWriteFuture.join();
					}

					this.previousWriteFuture = chunkStorage.write(chunkPos, () -> compoundTag2);
					return true;
				}
			}

			return false;
		}

		protected ChunkStorage createStorage(RegionStorageInfo regionStorageInfo, Path path) {
			return (ChunkStorage)(WorldUpgrader.this.recreateRegionFiles
				? new RecreatingChunkStorage(
					regionStorageInfo.withTypeSuffix("source"),
					path,
					regionStorageInfo.withTypeSuffix("target"),
					WorldUpgrader.resolveRecreateDirectory(path),
					WorldUpgrader.this.dataFixer,
					true
				)
				: new ChunkStorage(regionStorageInfo, path, WorldUpgrader.this.dataFixer, true));
		}
	}

	record DimensionToUpgrade<T>(ResourceKey<Level> dimensionKey, T storage, ListIterator<WorldUpgrader.FileToUpgrade> files) {
	}

	class EntityUpgrader extends WorldUpgrader.SimpleRegionStorageUpgrader {
		EntityUpgrader() {
			super(DataFixTypes.ENTITY_CHUNK, "entities", WorldUpgrader.STATUS_UPGRADING_ENTITIES, WorldUpgrader.STATUS_FINISHED_ENTITIES);
		}

		@Override
		protected CompoundTag upgradeTag(SimpleRegionStorage regionStorage, CompoundTag chunkTag) {
			return regionStorage.upgradeChunkTag(chunkTag, -1);
		}
	}

	record FileToUpgrade(RegionFile file, List<ChunkPos> chunksToUpgrade) {
	}

	class PoiUpgrader extends WorldUpgrader.SimpleRegionStorageUpgrader {
		PoiUpgrader() {
			super(DataFixTypes.POI_CHUNK, "poi", WorldUpgrader.STATUS_UPGRADING_POI, WorldUpgrader.STATUS_FINISHED_POI);
		}

		@Override
		protected CompoundTag upgradeTag(SimpleRegionStorage regionStorage, CompoundTag chunkTag) {
			return regionStorage.upgradeChunkTag(chunkTag, 1945);
		}
	}

	abstract class SimpleRegionStorageUpgrader extends WorldUpgrader.AbstractUpgrader<SimpleRegionStorage> {
		SimpleRegionStorageUpgrader(final DataFixTypes dataFixType, final String type, final Component upgradingStatus, final Component finishedStatus) {
			super(dataFixType, type, type, upgradingStatus, finishedStatus);
		}

		protected SimpleRegionStorage createStorage(RegionStorageInfo regionStorageInfo, Path path) {
			return (SimpleRegionStorage)(WorldUpgrader.this.recreateRegionFiles
				? new RecreatingSimpleRegionStorage(
					regionStorageInfo.withTypeSuffix("source"),
					path,
					regionStorageInfo.withTypeSuffix("target"),
					WorldUpgrader.resolveRecreateDirectory(path),
					WorldUpgrader.this.dataFixer,
					true,
					this.dataFixType
				)
				: new SimpleRegionStorage(regionStorageInfo, path, WorldUpgrader.this.dataFixer, true, this.dataFixType));
		}

		protected boolean tryProcessOnePosition(SimpleRegionStorage simpleRegionStorage, ChunkPos chunkPos, ResourceKey<Level> resourceKey) {
			CompoundTag compoundTag = (CompoundTag)((Optional)simpleRegionStorage.read(chunkPos).join()).orElse(null);
			if (compoundTag != null) {
				int i = ChunkStorage.getVersion(compoundTag);
				CompoundTag compoundTag2 = this.upgradeTag(simpleRegionStorage, compoundTag);
				boolean bl = i < SharedConstants.getCurrentVersion().getDataVersion().getVersion();
				if (bl || WorldUpgrader.this.recreateRegionFiles) {
					if (this.previousWriteFuture != null) {
						this.previousWriteFuture.join();
					}

					this.previousWriteFuture = simpleRegionStorage.write(chunkPos, compoundTag2);
					return true;
				}
			}

			return false;
		}

		protected abstract CompoundTag upgradeTag(SimpleRegionStorage regionStorage, CompoundTag chunkTag);
	}
}