package net.mcreator.customoregen.worldgen;

import com.mojang.brigadier.CommandDispatcher;
import com.mojang.brigadier.arguments.IntegerArgumentType;
import com.mojang.brigadier.context.CommandContext;
import net.mcreator.customoregen.CustomOreGenMod;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.Commands;
import net.minecraft.core.Holder;
import net.minecraft.core.HolderGetter;
import net.minecraft.core.registries.Registries;
import net.minecraft.network.chat.Component;
import net.minecraft.resources.ResourceKey;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.level.biome.Biome;
import net.minecraft.world.level.biome.Biomes;
import net.neoforged.bus.api.SubscribeEvent;
import net.neoforged.fml.common.EventBusSubscriber;
import net.neoforged.neoforge.event.RegisterCommandsEvent;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * {@code /latitude map [radius] [step]} — renders a top-down biome map of the latitude world.
 *
 * <p>Samples {@link LatitudeBiomeSource} on a large square grid (independent of the current
 * world's biome source, so it works in any world) and writes:
 * <ul>
 *   <li>{@code run/latitude/latitude_map.png} — one colour per biome category (the "map view"),</li>
 *   <li>{@code run/latitude/latitude_report.txt} — per-band distribution + validation invariants.</li>
 * </ul>
 *
 * <p>Use this to validate the climate distribution without exploring by hand.</p>
 */
@EventBusSubscriber(modid = CustomOreGenMod.MODID, bus = EventBusSubscriber.Bus.GAME)
public class LatitudeMapCommand {

	private static final int DEFAULT_RADIUS = 24000;
	private static final int DEFAULT_STEP = 60;

	// Biome key -> colour category (kept in sync with BiomeBand pools)
	private static final Map<ResourceKey<Biome>, Integer> COLORS = new HashMap<>();

	static {
		col(0xB6E388, Biomes.PLAINS, Biomes.SUNFLOWER_PLAINS, Biomes.FOREST, Biomes.BIRCH_FOREST,
				Biomes.FLOWER_FOREST, Biomes.MEADOW, Biomes.CHERRY_GROVE); // safe spawn / temperate
		col(0xC8E6FF, Biomes.SNOWY_PLAINS, Biomes.ICE_SPIKES, Biomes.SNOWY_TAIGA, Biomes.GROVE,
				Biomes.SNOWY_SLOPES, Biomes.JAGGED_PEAKS, Biomes.FROZEN_PEAKS); // frozen
		col(0x8FB8D6, Biomes.TAIGA, Biomes.OLD_GROWTH_PINE_TAIGA, Biomes.OLD_GROWTH_SPRUCE_TAIGA,
				Biomes.WINDSWEPT_HILLS, Biomes.WINDSWEPT_FOREST, Biomes.WINDSWEPT_GRAVELLY_HILLS); // cold
		col(0x7CC576, Biomes.DARK_FOREST, Biomes.OLD_GROWTH_BIRCH_FOREST); // temperate
		col(0xE0C068, Biomes.SAVANNA, Biomes.SAVANNA_PLATEAU, Biomes.WINDSWEPT_SAVANNA); // warm
		col(0xE0884C, Biomes.DESERT, Biomes.BADLANDS, Biomes.WOODED_BADLANDS, Biomes.ERODED_BADLANDS,
				Biomes.JUNGLE, Biomes.SPARSE_JUNGLE, Biomes.BAMBOO_JUNGLE); // hot
		col(0x2B5C8A, Biomes.OCEAN, Biomes.DEEP_OCEAN, Biomes.LUKEWARM_OCEAN, Biomes.DEEP_LUKEWARM_OCEAN,
				Biomes.WARM_OCEAN, Biomes.FROZEN_OCEAN, Biomes.DEEP_FROZEN_OCEAN, Biomes.COLD_OCEAN,
				Biomes.DEEP_COLD_OCEAN, Biomes.RIVER, Biomes.FROZEN_RIVER); // water
		col(0x5B7A3A, Biomes.SWAMP); // swamp (should be rare pockets)
		col(0x3A5A3A, Biomes.MANGROVE_SWAMP); // mangrove
		col(0x6B4E3A, Biomes.LUSH_CAVES, Biomes.DRIPSTONE_CAVES, Biomes.DEEP_DARK); // underground
	}

	private static void col(int rgb, ResourceKey<Biome>... keys) {
		for (ResourceKey<Biome> k : keys) COLORS.put(k, rgb);
	}

	@SubscribeEvent
	public static void onRegisterCommands(RegisterCommandsEvent event) {
		CommandDispatcher<CommandSourceStack> dispatcher = event.getDispatcher();
		dispatcher.register(Commands.literal("latitude")
				.requires(source -> source.hasPermission(2))
				.then(Commands.literal("map")
						.executes(ctx -> run(ctx, DEFAULT_RADIUS, DEFAULT_STEP))
						.then(Commands.argument("radius", IntegerArgumentType.integer(1000, 100000))
								.executes(ctx -> run(ctx, IntegerArgumentType.getInteger(ctx, "radius"), DEFAULT_STEP))
								.then(Commands.argument("step", IntegerArgumentType.integer(4, 1000))
										.executes(ctx -> run(ctx,
												IntegerArgumentType.getInteger(ctx, "radius"),
												IntegerArgumentType.getInteger(ctx, "step")))))));
	}

	private static int run(CommandContext<CommandSourceStack> ctx, int radius, int step) {
		CommandSourceStack src = ctx.getSource();
		if (!(src.getLevel() instanceof ServerLevel serverLevel)) {
			src.sendFailure(Component.literal("Must be run on a server level."));
			return 0;
		}

		long seed = serverLevel.getSeed();
		HolderGetter<Biome> getter = serverLevel.registryAccess().lookupOrThrow(Registries.BIOME);
		LatitudeBiomeSource source = new LatitudeBiomeSource(seed, getter);

		int dim = (2 * radius) / step + 1;
		src.sendSystemMessage(Component.literal(
				"\u00A7b[Latitude]\u00A7r Sampling " + dim + "x" + dim + " grid (radius=" + radius + ", step=" + step + ")..."));

		BufferedImage img = new BufferedImage(dim, dim, BufferedImage.TYPE_INT_RGB);
		Map<ResourceKey<Biome>, Integer> global = new HashMap<>();
		Map<BiomeBand, Map<ResourceKey<Biome>, Integer>> perBand = new EnumMap<>(BiomeBand.class);
		for (BiomeBand b : BiomeBand.values()) perBand.put(b, new HashMap<>());

		for (int zi = 0; zi < dim; zi++) {
			for (int xi = 0; xi < dim; xi++) {
				int blockX = -radius + xi * step;
				int blockZ = -radius + zi * step;
				Holder<Biome> holder = source.getNoiseBiome(blockX >> 2, 64 >> 2, blockZ >> 2, null);
				ResourceKey<Biome> key = holder.unwrapKey().orElse(null);

				global.merge(key, 1, Integer::sum);
				perBand.get(BiomeBand.fromTemperature(blockZ / 16000.0)).merge(key, 1, Integer::sum);

				int rgb = key == null ? 0x888888 : COLORS.getOrDefault(key, 0x888888);
				img.setRGB(xi, dim - 1 - zi, rgb); // flip Z so north is at the top
			}
		}

		// ---- Spawn biome ----
		ResourceKey<Biome> spawnKey = source.getNoiseBiome(0, 64 >> 2, 0, null).unwrapKey().orElse(null);

		// ---- Write PNG ----
		try {
			Files.createDirectories(Path.of("latitude"));
			ImageIO.write(img, "png", Path.of("latitude", "latitude_map.png").toFile());
		} catch (Exception e) {
			CustomOreGenMod.LOGGER.error("Latitude map: failed to write PNG", e);
			src.sendFailure(Component.literal("Failed to write PNG: " + e.getMessage()));
			return 0;
		}

		// ---- Build report ----
		StringBuilder r = new StringBuilder();
		r.append("Latitude World Map Report\n");
		r.append("seed=").append(seed).append("  radius=").append(radius).append("  step=").append(step)
				.append("  grid=").append(dim).append("x").append(dim).append("\n\n");
		r.append("Spawn biome @ (0,0): ").append(spawnKey).append("\n\n");

		int total = global.values().stream().mapToInt(Integer::intValue).sum();
		r.append("Global biome distribution (top 15):\n");
		global.entrySet().stream()
				.sorted(Map.Entry.<ResourceKey<Biome>, Integer>comparingByValue().reversed())
				.limit(15)
				.forEach(e -> r.append(String.format("  %-45s %6.2f%%%n", e.getKey(), 100.0 * e.getValue() / total)));

		List<String> warnings = new ArrayList<>();
		for (BiomeBand band : BiomeBand.values()) {
			Map<ResourceKey<Biome>, Integer> m = perBand.get(band);
			int bt = m.values().stream().mapToInt(Integer::intValue).sum();
			if (bt == 0) continue;
			int swamp = m.getOrDefault(Biomes.SWAMP, 0) + m.getOrDefault(Biomes.MANGROVE_SWAMP, 0);
			double swampPct = 100.0 * swamp / bt;
			r.append("\nBand ").append(band).append(" (").append(bt).append(" samples, swamp+mangrove ")
					.append(String.format("%.2f%%", swampPct)).append("):\n");
			m.entrySet().stream()
					.sorted(Map.Entry.<ResourceKey<Biome>, Integer>comparingByValue().reversed())
					.limit(6)
					.forEach(e -> r.append(String.format("    %-43s %6.2f%%%n", e.getKey(), 100.0 * e.getValue() / bt)));
			if (swampPct > 15.0) {
				warnings.add(band + " has " + String.format("%.1f%%", swampPct) + " swamp (expected <15%)");
			}
		}

		// Invariant checks
		r.append("\n=== Validation ===\n");
		boolean spawnOk = spawnKey != null && isSafeSpawn(spawnKey);
		r.append("Spawn on safe biome: ").append(spawnOk ? "PASS" : "FAIL (" + spawnKey + ")").append("\n");
		double frozenCold = bandRatio(perBand.get(BiomeBand.FROZEN), Set.of(
				Biomes.SNOWY_PLAINS, Biomes.ICE_SPIKES, Biomes.SNOWY_TAIGA, Biomes.GROVE, Biomes.SNOWY_SLOPES,
				Biomes.JAGGED_PEAKS, Biomes.FROZEN_PEAKS, Biomes.TAIGA, Biomes.OLD_GROWTH_PINE_TAIGA,
				Biomes.OLD_GROWTH_SPRUCE_TAIGA, Biomes.WINDSWEPT_HILLS, Biomes.WINDSWEPT_FOREST,
				Biomes.FROZEN_OCEAN, Biomes.DEEP_FROZEN_OCEAN, Biomes.COLD_OCEAN, Biomes.DEEP_COLD_OCEAN,
				Biomes.FROZEN_RIVER));
		r.append(String.format("North (FROZEN) cold/frozen ratio: %.1f%% (want >60%%)%n", frozenCold));
		double hotRatio = bandRatio(perBand.get(BiomeBand.HOT), Set.of(
				Biomes.DESERT, Biomes.BADLANDS, Biomes.WOODED_BADLANDS, Biomes.ERODED_BADLANDS,
				Biomes.JUNGLE, Biomes.SPARSE_JUNGLE, Biomes.BAMBOO_JUNGLE, Biomes.SAVANNA,
				Biomes.SAVANNA_PLATEAU, Biomes.WARM_OCEAN));
		r.append(String.format("South (HOT) warm/hot ratio: %.1f%% (want >60%%)%n", hotRatio));
		r.append("Warnings: ").append(warnings.isEmpty() ? "none" : warnings).append("\n");

		try {
			Files.writeString(Path.of("latitude", "latitude_report.txt"), r.toString());
		} catch (Exception e) {
			CustomOreGenMod.LOGGER.error("Latitude map: failed to write report", e);
		}

		// ---- Summary to chat ----
		src.sendSystemMessage(Component.literal("\u00A7a[Latitude]\u00A7r Map + report saved to /latitude/"));
		src.sendSystemMessage(Component.literal("  Spawn: " + spawnKey + (spawnOk ? " \u00A7aOK\u00A7r" : " \u00A7cBAD\u00A7r")));
		src.sendSystemMessage(Component.literal(String.format(
				"  North cold/frozen: %.0f%% | South hot: %.0f%% | Swamp warnings: %d",
				frozenCold, hotRatio, warnings.size())));
		return 1;
	}

	private static double bandRatio(Map<ResourceKey<Biome>, Integer> m, Set<ResourceKey<Biome>> expected) {
		int total = m.values().stream().mapToInt(Integer::intValue).sum();
		if (total == 0) return 0;
		int match = 0;
		for (ResourceKey<Biome> k : expected) match += m.getOrDefault(k, 0);
		return 100.0 * match / total;
	}

	private static boolean isSafeSpawn(ResourceKey<Biome> key) {
		Set<ResourceKey<Biome>> safe = new HashSet<>(List.of(
				Biomes.PLAINS, Biomes.SUNFLOWER_PLAINS, Biomes.FOREST, Biomes.BIRCH_FOREST,
				Biomes.FLOWER_FOREST, Biomes.MEADOW, Biomes.CHERRY_GROVE));
		return safe.contains(key);
	}
}
