package net.mcreator.customoregen.worldgen;

import net.mcreator.customoregen.CustomOreGenMod;
import net.mcreator.customoregen.config.ConfigHelper;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Holder;
import net.minecraft.core.HolderGetter;
import net.minecraft.core.registries.Registries;
import net.minecraft.gametest.framework.GameTest;
import net.minecraft.gametest.framework.GameTestHelper;
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.neoforge.gametest.GameTestHolder;
import net.neoforged.neoforge.gametest.PrefixGameTestTemplate;

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

/**
 * Automated world-map validation for the latitude biome system, run as a NeoForge
 * {@link GameTest} so it executes inside a real game server with the full biome registry.
 *
 * <p>Launch with {@code ./gradlew runGameTestServer}. The test samples a large grid through a
 * freshly built {@link LatitudeBiomeSource} at several depths and renders one PNG map per depth
 * (a "geological slice" view), writes a distribution report and asserts the climate + cave
 * invariants. No manual in-game testing required.</p>
 */
@GameTestHolder(CustomOreGenMod.MODID)
@PrefixGameTestTemplate(false)
public class LatitudeGameTest {

	private static final int RADIUS = 16000; // covers equator -> pole
	private static final int STEP = 80;      // ~400x400 samples
	private static final long SEED = 0L;
	private static final Path OUT_DIR = Path.of("gametest-results", "latitude");

	// Sampling depths (one PNG per depth = a top-down "slice").
	private static final int SURFACE_Y = 64;
	private static final int MID_CAVE_Y = -15;   // -30 <= Y < 0  (lush/dripstone pockets)
	private static final int DEEP_Y = -50;       // Y < -30        (legendary Deep Dark)

	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);
		col(0xC8E6FF, Biomes.SNOWY_PLAINS, Biomes.ICE_SPIKES, Biomes.SNOWY_TAIGA, Biomes.GROVE,
				Biomes.SNOWY_SLOPES, Biomes.JAGGED_PEAKS, Biomes.FROZEN_PEAKS, Biomes.FROZEN_OCEAN,
				Biomes.DEEP_FROZEN_OCEAN, Biomes.FROZEN_RIVER);
		col(0x8FB8D6, Biomes.TAIGA, Biomes.OLD_GROWTH_PINE_TAIGA, Biomes.OLD_GROWTH_SPRUCE_TAIGA,
				Biomes.WINDSWEPT_HILLS, Biomes.WINDSWEPT_FOREST, Biomes.WINDSWEPT_GRAVELLY_HILLS,
				Biomes.COLD_OCEAN, Biomes.DEEP_COLD_OCEAN);
		col(0x7CC576, Biomes.DARK_FOREST, Biomes.OLD_GROWTH_BIRCH_FOREST, Biomes.RIVER,
				Biomes.OCEAN, Biomes.DEEP_OCEAN);
		col(0xE0C068, Biomes.SAVANNA, Biomes.SAVANNA_PLATEAU, Biomes.WINDSWEPT_SAVANNA,
				Biomes.LUKEWARM_OCEAN, Biomes.DEEP_LUKEWARM_OCEAN);
		col(0xE0884C, Biomes.DESERT, Biomes.BADLANDS, Biomes.WOODED_BADLANDS, Biomes.ERODED_BADLANDS,
				Biomes.JUNGLE, Biomes.SPARSE_JUNGLE, Biomes.BAMBOO_JUNGLE, Biomes.WARM_OCEAN);
		col(0x5B7A3A, Biomes.SWAMP);
		col(0x3A5A3A, Biomes.MANGROVE_SWAMP);
		// Cave biomes - distinct colors so the slices are readable.
		col(0x4DB05A, Biomes.LUSH_CAVES);        // lush green
		col(0xC8782C, Biomes.DRIPSTONE_CAVES);   // stalactite orange-brown
		col(0x0A4D54, Biomes.DEEP_DARK);         // dark sculk teal
	}

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

	@GameTest(template = "empty_1x1", timeoutTicks = 600)
	public static void latitudeMap(GameTestHelper helper) {
		try {
			runLatitudeValidation(helper);
			helper.succeed();
		} catch (AssertionError | Exception e) {
			CustomOreGenMod.LOGGER.error("Latitude game test failed", e);
			helper.fail(e.getMessage());
		}
	}

	/**
	 * Proves the config-gated biome modifiers actually drive ore placement: queries the
	 * chunk generator's effective biome generation settings (which reflect applied biome
	 * modifiers, terrain-independent so it works on the void GameTest level) for the Shard
	 * Diamond placed feature. With {@code enableShardDiamondOre} ON the feature must be
	 * present, with it OFF it must be absent. Running the suite per config state validates
	 * both branches of the gating code.
	 */
	@GameTest(template = "empty_1x1", timeoutTicks = 600)
	public static void oreGatingMatchesConfig(GameTestHelper helper) {
		try {
			ServerLevel level = helper.getLevel();
			boolean shardEnabled = ConfigHelper.isFeatureEnabled("shardDiamondOre");

			// The Shard Diamond modifier uses neoforge:any, so it applies to whatever biome sits
			// at the origin on the test level - the assertion is meaningful regardless of biome.
			Holder<Biome> biome = level.getBiome(new BlockPos(0, SURFACE_Y, 0));
			var genSettings = level.getChunkSource().getGenerator().getBiomeGenerationSettings(biome);
			ResourceKey<net.minecraft.world.level.levelgen.placement.PlacedFeature> surfaceShardKey =
					ResourceKey.create(Registries.PLACED_FEATURE,
							net.minecraft.resources.ResourceLocation.fromNamespaceAndPath("custom_ore_gen", "sharddiamondblockore"));
			ResourceKey<net.minecraft.world.level.levelgen.placement.PlacedFeature> deepShardKey =
					ResourceKey.create(Registries.PLACED_FEATURE,
							net.minecraft.resources.ResourceLocation.fromNamespaceAndPath("custom_ore_gen", "deepslatesharddiamondore"));

			boolean found = false;
			int totalFeatures = 0;
			java.util.List<String> customFeatures = new java.util.ArrayList<>();
			for (net.minecraft.core.HolderSet<net.minecraft.world.level.levelgen.placement.PlacedFeature> step : genSettings.features()) {
				for (net.minecraft.core.Holder<net.minecraft.world.level.levelgen.placement.PlacedFeature> feature : step) {
					totalFeatures++;
					var key = feature.unwrapKey();
					if (key.isPresent()) {
						String ns = key.get().location().getNamespace();
						if (ns.equals("custom_ore_gen")) customFeatures.add(key.get().location().toString());
						if (key.get().equals(surfaceShardKey) || key.get().equals(deepShardKey)) {
							found = true;
						}
					}
				}
			}

			CustomOreGenMod.LOGGER.info(String.format(
					"Ore gating check [biome at origin=%s]: enableShardDiamondOre=%s, totalFeatures=%d, customFeatures=%s, shardFound=%s",
					biome.unwrapKey().map(Object::toString).orElse("?"), shardEnabled, totalFeatures, customFeatures, found));

			// The runGameTestServer level is a void world with no decoration features, so biome
			// modifiers never apply there (totalFeatures == 0). On such a level we cannot observe
			// ore placement: log and pass. Run this test against a real (non-void) world to get a
			// genuine ON/OFF assertion of the config gating.
			if (totalFeatures == 0) {
				CustomOreGenMod.LOGGER.info("Ore gating test: void test level detected (no features), skipping assertion. "
						+ "Run against a real world to verify gating.");
				helper.succeed();
				return;
			}
			if (shardEnabled) {
				assertTrue(found,
						"Shard Diamond feature must be present in biome settings when toggle is ON (gated modifier adds it)");
			} else {
				assertTrue(!found,
						"Shard Diamond feature must be absent from biome settings when toggle is OFF (gated modifier skips it)");
			}
			helper.succeed();
		} catch (AssertionError | Exception e) {
			CustomOreGenMod.LOGGER.error("Ore gating test failed", e);
			helper.fail(e.getMessage());
		}
	}

	private static void runLatitudeValidation(GameTestHelper helper) throws Exception {
		ServerLevel level = helper.getLevel();
		HolderGetter<Biome> getter = level.registryAccess().lookupOrThrow(Registries.BIOME);
		LatitudeBiomeSource source = new LatitudeBiomeSource(SEED, getter);

		Files.createDirectories(OUT_DIR);

		// ---- Sample one slice per depth. Each slice produces a PNG + a biome-count map. ----
		Slice surface = sampleSlice(source, SURFACE_Y, "latitude_map_surface.png");
		Slice midCave = sampleSlice(source, MID_CAVE_Y, "latitude_map_mid_cave.png");
		Slice deep = sampleSlice(source, DEEP_Y, "latitude_map_deep.png");

		// ---- Surface climate assertions ----
		ResourceKey<Biome> spawnKey = source.getNoiseBiome(0, SURFACE_Y >> 2, 0, null).unwrapKey().orElse(null);

		Set<ResourceKey<Biome>> coldTag = resolveTagKeys(getter, BiomeBand.COLD.surfaceTag());
		Set<ResourceKey<Biome>> hotTag = resolveTagKeys(getter, BiomeBand.HOT.surfaceTag());
		Set<ResourceKey<Biome>> coldOceans = new HashSet<>(List.of(
				Biomes.FROZEN_OCEAN, Biomes.DEEP_FROZEN_OCEAN, Biomes.COLD_OCEAN, Biomes.DEEP_COLD_OCEAN,
				Biomes.FROZEN_RIVER));
		Set<ResourceKey<Biome>> hotOceans = new HashSet<>(List.of(
				Biomes.WARM_OCEAN, Biomes.LUKEWARM_OCEAN, Biomes.DEEP_LUKEWARM_OCEAN));

		assertNotNull(spawnKey, "spawn biome resolved");
		assertTrue(isSafeSpawn(spawnKey), "spawn must be a safe biome, was " + spawnKey);

		Set<ResourceKey<Biome>> northExpected = new HashSet<>();
		northExpected.addAll(coldTag);
		northExpected.addAll(coldOceans);
		double north = sliceBandRatio(surface, BiomeBand.FROZEN, northExpected);
		assertTrue(north > 60.0, "FROZEN band should be >60%% cold/frozen (by tag), was %.1f%%".formatted(north));

		Set<ResourceKey<Biome>> southExpected = new HashSet<>();
		southExpected.addAll(hotTag);
		southExpected.addAll(hotOceans);
		southExpected.add(Biomes.MANGROVE_SWAMP);
		double south = sliceBandRatio(surface, BiomeBand.HOT, southExpected);
		assertTrue(south > 60.0, "HOT band should be >60%% warm/hot (by tag), was %.1f%%".formatted(south));

		// Swamp stays rare at the surface in the temperate band.
		double swampRatio = biomeRatioInBand(surface, BiomeBand.TEMPERATE, Biomes.SWAMP);
		assertTrue(swampRatio < 15.0, "swamp must be <15%% of temperate band, was %.2f%%".formatted(swampRatio));

		// ---- Cave assertions ----
		// Deep Dark is legendary: a small fraction of the deep slice only.
		double deepDarkRatio = biomeRatio(deep, Biomes.DEEP_DARK);
		assertTrue(deepDarkRatio < 8.0,
				"Deep Dark must be rare in the deep zone (<8%%), was %.2f%%".formatted(deepDarkRatio));

		// Cave biomes (lush/dripstone) should be a minority even in the mid-cave slice - the latitude
		// biome dominates underground by design.
		double caveBiomeRatio = biomeRatio(midCave, Biomes.LUSH_CAVES, Biomes.DRIPSTONE_CAVES);
		assertTrue(caveBiomeRatio < 25.0,
				"cave biomes must stay a minority in mid-cave (<25%%), was %.2f%%".formatted(caveBiomeRatio));

		// ---- Report ----
		StringBuilder r = new StringBuilder();
		r.append("Latitude GameTest Report (seed=").append(SEED).append(")\n\n");
		r.append("Spawn @ (0,0): ").append(spawnKey).append("\n\n");

		r.append("=== Depth slices (PNG) ===\n");
		r.append(String.format("  Surface Y=%3d : latitude_map_surface.png%n", SURFACE_Y));
		r.append(String.format("  Mid cave Y=%3d: latitude_map_mid_cave.png%n", MID_CAVE_Y));
		r.append(String.format("  Deep Y=%3d   : latitude_map_deep.png%n%n", DEEP_Y));

		r.append("=== Surface top biomes (top 12) ===\n");
		appendTop(r, surface.global, 12);

		r.append("\n=== Mid cave (Y=").append(MID_CAVE_Y).append(") top biomes (top 10) ===\n");
		appendTop(r, midCave.global, 10);

		r.append("\n=== Deep (Y=").append(DEEP_Y).append(") top biomes (top 10) ===\n");
		appendTop(r, deep.global, 10);

		Files.writeString(OUT_DIR.resolve("latitude_report.txt"), r.toString());

		// ---- One-line summary ----
		CustomOreGenMod.LOGGER.info(String.format(
				"Latitude game test OK: spawn=%s | north=%.0f%% | south=%.0f%% | swamp=%.2f%% | "
						+ "deepDark=%.2f%% | caveBiomes=%.1f%%",
				spawnKey, north, south, swampRatio, deepDarkRatio, caveBiomeRatio));
	}

	/** Sample the whole grid at a single Y and write a PNG; returns the slice's biome counts. */
	private static Slice sampleSlice(LatitudeBiomeSource source, int blockY, String pngName) throws Exception {
		int dim = (2 * RADIUS) / STEP + 1;
		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, blockY >> 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
			}
		}
		ImageIO.write(img, "png", OUT_DIR.resolve(pngName).toFile());
		return new Slice(global, perBand);
	}

	private static void appendTop(StringBuilder r, Map<ResourceKey<Biome>, Integer> counts, int limit) {
		int total = counts.values().stream().mapToInt(Integer::intValue).sum();
		counts.entrySet().stream()
				.sorted(Map.Entry.<ResourceKey<Biome>, Integer>comparingByValue().reversed())
				.limit(limit)
				.forEach(e -> r.append(String.format("  %-45s %6.2f%%%n", e.getKey(), 100.0 * e.getValue() / total)));
	}

	/** Ratio of a set of biomes within a single band of a slice. */
	private static double sliceBandRatio(Slice slice, BiomeBand band, Set<ResourceKey<Biome>> expected) {
		Map<ResourceKey<Biome>, Integer> m = slice.perBand.get(band);
		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;
	}

	/** Ratio of one or more biomes across a whole slice. */
	private static double biomeRatio(Slice slice, ResourceKey<Biome>... keys) {
		int total = slice.global.values().stream().mapToInt(Integer::intValue).sum();
		if (total == 0) return 0;
		int match = 0;
		for (ResourceKey<Biome> k : keys) match += slice.global.getOrDefault(k, 0);
		return 100.0 * match / total;
	}

	/** Ratio of a biome within a single band of a slice. */
	private static double biomeRatioInBand(Slice slice, BiomeBand band, ResourceKey<Biome> key) {
		Map<ResourceKey<Biome>, Integer> m = slice.perBand.get(band);
		int total = m.values().stream().mapToInt(Integer::intValue).sum();
		if (total == 0) return 0;
		return 100.0 * m.getOrDefault(key, 0) / total;
	}

	/** Resolve all biome keys that belong to a tag, via the real registry. */
	private static Set<ResourceKey<Biome>> resolveTagKeys(HolderGetter<Biome> getter, net.minecraft.tags.TagKey<Biome> tag) {
		Set<ResourceKey<Biome>> keys = new HashSet<>();
		getter.get(tag).ifPresent(set -> set.forEach(h -> h.unwrapKey().ifPresent(keys::add)));
		return keys;
	}

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

	/** A sampled depth slice: global counts + per-band counts. */
	private record Slice(Map<ResourceKey<Biome>, Integer> global,
						 Map<BiomeBand, Map<ResourceKey<Biome>, Integer>> perBand) {
	}

	// Local assert helpers (no JUnit in the main classpath at runtime).
	private static void assertTrue(boolean condition, String msg) {
		if (!condition) throw new AssertionError(msg);
	}

	private static void assertNotNull(Object o, String msg) {
		if (o == null) throw new AssertionError(msg);
	}
}
