package net.mcreator.customoregen.worldgen;

import com.mojang.serialization.Codec;
import com.mojang.serialization.MapCodec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import net.minecraft.core.Holder;
import net.minecraft.core.HolderGetter;
import net.minecraft.core.HolderSet;
import net.minecraft.core.registries.Registries;
import net.minecraft.resources.RegistryOps;
import net.minecraft.util.RandomSource;
import net.minecraft.world.level.biome.Biome;
import net.minecraft.world.level.biome.BiomeSource;
import net.minecraft.world.level.biome.Biomes;
import net.minecraft.world.level.biome.Climate;
import net.minecraft.world.level.levelgen.synth.ImprovedNoise;

import java.util.ArrayList;
import java.util.EnumMap;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Stream;

/**
 * Custom {@link BiomeSource} that distributes biomes by latitude (Z axis).
 *
 * <ul>
 *   <li>Far north (large negative Z) = frozen / cold biomes</li>
 *   <li>Equator (Z &asymp; 0) = temperate biomes</li>
 *   <li>Far south (large positive Z) = hot biomes</li>
 * </ul>
 *
 * <p>Temperature is derived from the Z coordinate with a gentle boundary perturbation so the
 * climate bands are not perfectly straight lines. Inside each band, a very low-frequency
 * selector noise picks among the band's biomes, producing extremely large biomes that
 * encourage exploration and long-distance travel (railways, roads, nether hubs).</p>
 *
 * <p><b>Underground philosophy:</b> the surface latitude biome extends downward by default, so
 * mining under a desert feels like being in the desert. Cave biomes (lush / dripstone) only
 * appear as rare climatic pockets in the mid underground, and the Deep Dark is a legendary
 * (~1%) feature of the deep zone. This keeps the latitude logic coherent at every depth
 * instead of overriding it with a uniform cave-biome slab.</p>
 *
 * <p>Surface biomes come from dedicated climate tags (see {@code data/custom_ore_gen/tags/
 * worldgen/biome/latitude_*_surface.json}), which support optional Biomes O' Plenty biomes
 * via {@code "required": false} without ever creating unbound holders. Ocean pools are
 * vanilla-only.</p>
 */
public class LatitudeBiomeSource extends BiomeSource {

	public static final MapCodec<LatitudeBiomeSource> CODEC = RecordCodecBuilder.mapCodec(instance -> instance.group(
			Codec.LONG.fieldOf("seed").forGetter(src -> src.seed),
			RegistryOps.retrieveGetter(Registries.BIOME)
	).apply(instance, instance.stable(LatitudeBiomeSource::new)));

	// ------------------------------------------------------------------
	// Surface / climate tunables
	// ------------------------------------------------------------------

	/** Number of blocks for the temperature to go from 0 (equator) to +-1 (pole). */
	private static final double TEMPERATURE_SCALE = 16000.0;

	/** Noise frequency for the climate band boundary wobble. */
	private static final double BOUNDARY_NOISE_SCALE = 0.00015;

	/** Amplitude of the boundary wobble (in temperature units). */
	private static final double BOUNDARY_NOISE_AMPLITUDE = 0.15;

	/** Selector noise frequency for surface sub-biomes (very low = very large biomes). */
	private static final double SURFACE_SELECTOR_SCALE = 0.00033;

	/** Frequency of the land/ocean mask noise. */
	private static final double OCEAN_NOISE_SCALE = 0.0011;

	/** Above this ocean-noise value, the column is ocean. */
	private static final double OCEAN_THRESHOLD = 0.30;

	/** Frequency of the moisture noise that carves rare swamp/mangrove pockets. */
	private static final double MOISTURE_SCALE = 0.0009;

	/** Above this moisture value a wet biome (swamp/mangrove) overrides the surface. ~8% of land. */
	private static final double MOISTURE_THRESHOLD = 0.55;

	// ------------------------------------------------------------------
	// Underground tunables (3-zone model)
	// ------------------------------------------------------------------

	/** Top of the deep zone (Y &lt; this = deep underground). Deep Dark can live here. */
	private static final int DEEP_ZONE_TOP = -30;

	/** Top of the mid-cave zone (DEEP_ZONE_TOP &le; Y &lt; this = mid caves, lush/dripstone pockets). */
	private static final int MID_CAVE_TOP = 0;

	/** Frequency of the mid-cave pocket noise (lush/dripstone). */
	private static final double CAVE_SCALE = 0.003;

	/** Above this value a lush/dripstone pocket overrides the surface biome. ~5-8% of the mid zone. */
	private static final double CAVE_THRESHOLD = 0.38;

	/** Frequency of the Deep Dark noise (very low = large, rare regions). */
	private static final double DEEP_DARK_SCALE = 0.0011;

	/** Vertical frequency of the Deep Dark noise (keeps it coherent in tall sections). */
	private static final double DEEP_DARK_SCALE_Y = 0.012;

	/** Above this value the Deep Dark overrides. Tuned for ~1-2% of the deep zone (legendary). */
	private static final double DEEP_DARK_THRESHOLD = 0.55;

	// ------------------------------------------------------------------
	// Spawn safe zone
	// ------------------------------------------------------------------

	/** Half-size of the guaranteed safe spawn square around the origin (plains/forest). */
	private static final int SPAWN_SAFE_RADIUS = 96;

	/** Selector frequency inside the spawn safe zone (finer, for gentle variety). */
	private static final double SPAWN_SELECTOR_SCALE = 0.02;

	// ------------------------------------------------------------------
	// Fields
	// ------------------------------------------------------------------

	private final long seed;
	private final HolderGetter<Biome> biomeGetter;

	private final ImprovedNoise boundaryNoise;
	private final ImprovedNoise selectorNoise;
	private final ImprovedNoise oceanNoise;
	private final ImprovedNoise moistureNoise;
	private final ImprovedNoise caveNoise;
	private final ImprovedNoise deepDarkNoise;

	/** Vanilla-only pools resolved eagerly (safe: vanilla biomes always exist and are bound). */
	private final EnumMap<BiomeBand, ResolvedBand> resolvedVanilla;
	private final Set<Holder<Biome>> possibleBiomes;
	private final Holder<Biome> fallback;

	/** Safe spawn pool (plains / forests) resolved eagerly. Always near the origin. */
	private final List<Holder<Biome>> spawnSafeBiomes;

	/** Rare wet biomes (swamp = temperate, mangrove = warm/hot), placed by moisture noise. */
	private final Holder<Biome> swampBiome;
	private final Holder<Biome> mangroveBiome;

	/** Cave biomes resolved eagerly (vanilla, always present). Lush = humid bands, Dripstone = dry bands. */
	private final Holder<Biome> lushCavesBiome;
	private final Holder<Biome> dripstoneCavesBiome;
	private final Holder<Biome> deepDarkBiome;

	/** Mod-aware surface pools resolved lazily from climate tags (safe: tags are bound by then). */
	private volatile EnumMap<BiomeBand, List<Holder<Biome>>> surfaceFromTag;

	public LatitudeBiomeSource(long seed, HolderGetter<Biome> biomeGetter) {
		this.seed = seed;
		this.biomeGetter = biomeGetter;

		RandomSource rng = RandomSource.create(seed);
		this.boundaryNoise = new ImprovedNoise(rng);
		this.selectorNoise = new ImprovedNoise(RandomSource.create(seed ^ 0x4C415449L));
		this.oceanNoise = new ImprovedNoise(RandomSource.create(seed ^ 0x4F434541L));
		this.moistureNoise = new ImprovedNoise(RandomSource.create(seed ^ 0x57455421L));
		this.caveNoise = new ImprovedNoise(RandomSource.create(seed ^ 0x43415645L));
		this.deepDarkNoise = new ImprovedNoise(RandomSource.create(seed ^ 0x44454550L));

		this.resolvedVanilla = new EnumMap<>(BiomeBand.class);
		Set<Holder<Biome>> all = new HashSet<>();
		for (BiomeBand band : BiomeBand.values()) {
			List<Holder<Biome>> surface = BiomeBand.resolve(biomeGetter, band.surfaceVanilla());
			List<Holder<Biome>> ocean = BiomeBand.resolve(biomeGetter, band.ocean());

			resolvedVanilla.put(band, new ResolvedBand(surface, ocean));
			all.addAll(surface);
			all.addAll(ocean);
		}

		// Cave biomes (vanilla) declared so the source advertises them as possible.
		this.lushCavesBiome = resolveAndAdd(biomeGetter, Biomes.LUSH_CAVES, all);
		this.dripstoneCavesBiome = resolveAndAdd(biomeGetter, Biomes.DRIPSTONE_CAVES, all);
		this.deepDarkBiome = resolveAndAdd(biomeGetter, Biomes.DEEP_DARK, all);

		this.possibleBiomes = all;

		List<Holder<Biome>> safe = new ArrayList<>();
		safe.addAll(BiomeBand.resolve(biomeGetter, List.of(
				Biomes.PLAINS, Biomes.SUNFLOWER_PLAINS,
				Biomes.FOREST, Biomes.BIRCH_FOREST, Biomes.FLOWER_FOREST,
				Biomes.MEADOW
		)));
		this.spawnSafeBiomes = safe.isEmpty() ? null : safe;

		this.swampBiome = biomeGetter.get(Biomes.SWAMP).orElse(null);
		this.mangroveBiome = biomeGetter.get(Biomes.MANGROVE_SWAMP).orElse(null);

		Holder<Biome> plains = biomeGetter.get(Biomes.PLAINS).orElse(null);
		this.fallback = plains != null ? plains : (possibleBiomes.isEmpty() ? null : possibleBiomes.iterator().next());
	}

	private static Holder<Biome> resolveAndAdd(HolderGetter<Biome> getter, net.minecraft.resources.ResourceKey<Biome> key, Set<Holder<Biome>> into) {
		Holder<Biome> holder = getter.get(key).orElse(null);
		if (holder != null) {
			into.add(holder);
		}
		return holder;
	}

	@Override
	protected MapCodec<? extends BiomeSource> codec() {
		return CODEC;
	}

	@Override
	protected Stream<Holder<Biome>> collectPossibleBiomes() {
		return possibleBiomes.stream();
	}

	@Override
	public Holder<Biome> getNoiseBiome(int quartX, int quartY, int quartZ, Climate.Sampler sampler) {
		int blockX = quartX * 4;
		int blockZ = quartZ * 4;
		int blockY = quartY * 4;

		double latitudeTemp = blockZ / TEMPERATURE_SCALE;
		double wobble = boundaryNoise.noise(blockX * BOUNDARY_NOISE_SCALE, 0.0, blockZ * BOUNDARY_NOISE_SCALE) * BOUNDARY_NOISE_AMPLITUDE;
		double temperature = clamp(latitudeTemp + wobble, -1.0, 1.0);

		BiomeBand band = BiomeBand.fromTemperature(temperature);
		ResolvedBand pool = resolvedVanilla.get(band);
		if (pool == null) {
			pool = resolvedVanilla.get(BiomeBand.TEMPERATE);
		}

		// Guaranteed safe spawn zone (open, buildable biomes) around the origin.
		if (spawnSafeBiomes != null && Math.abs(blockX) < SPAWN_SAFE_RADIUS && Math.abs(blockZ) < SPAWN_SAFE_RADIUS) {
			return pickBiome(spawnSafeBiomes, blockX, blockZ, SPAWN_SELECTOR_SCALE);
		}

		// Determine the column's base biome: ocean mask, rare wet pocket, or latitude surface biome.
		Holder<Biome> base = baseSurfaceBiome(band, pool, blockX, blockZ);

		// Underground cave overrides (3-zone model) extend latitude rather than replace it.
		return applyCaveOverrides(base, band, blockX, blockY, blockZ);
	}

	/** Resolve the surface biome for a column: ocean, swamp/mangrove pocket, or the band's surface tag. */
	private Holder<Biome> baseSurfaceBiome(BiomeBand band, ResolvedBand pool, int blockX, int blockZ) {
		double oceanValue = oceanNoise.noise(blockX * OCEAN_NOISE_SCALE, 0.0, blockZ * OCEAN_NOISE_SCALE);
		if (oceanValue > OCEAN_THRESHOLD && !pool.ocean().isEmpty()) {
			return pickBiome(pool.ocean(), blockX, blockZ, 0.002);
		}

		double moisture = moistureNoise.noise(blockX * MOISTURE_SCALE, 0.0, blockZ * MOISTURE_SCALE);
		if (moisture > MOISTURE_THRESHOLD) {
			if (band == BiomeBand.TEMPERATE && swampBiome != null) {
				return swampBiome;
			}
			if ((band == BiomeBand.WARM || band == BiomeBand.HOT) && mangroveBiome != null) {
				return mangroveBiome;
			}
		}

		List<Holder<Biome>> surface = resolveSurfaceTag(band);
		if (surface == null || surface.isEmpty()) {
			surface = pool.surface();
			if (surface.isEmpty()) {
				surface = resolvedVanilla.get(BiomeBand.TEMPERATE).surface();
			}
		}
		return pickBiome(surface, blockX, blockZ, SURFACE_SELECTOR_SCALE);
	}

	/**
	 * Underground overrides keep the latitude biome by default and only carve rare cave features:
	 * <ul>
	 *   <li>Deep zone (Y &lt; {@value #DEEP_ZONE_TOP}): legendary Deep Dark pockets (~1%), else surface biome.</li>
	 *   <li>Mid caves (DEEP_ZONE_TOP &le; Y &lt; {@value #MID_CAVE_TOP}): lush/dripstone pockets (~8%) matching the climate.</li>
	 *   <li>Near surface: no override, latitude biome as-is.</li>
	 * </ul>
	 */
	private Holder<Biome> applyCaveOverrides(Holder<Biome> surfaceBiome, BiomeBand band, int blockX, int blockY, int blockZ) {
		// Deep zone: legendary Deep Dark.
		if (blockY < DEEP_ZONE_TOP) {
			if (deepDarkBiome != null) {
				double dd = deepDarkNoise.noise(blockX * DEEP_DARK_SCALE, blockY * DEEP_DARK_SCALE_Y, blockZ * DEEP_DARK_SCALE);
				if (dd > DEEP_DARK_THRESHOLD) {
					return deepDarkBiome;
				}
			}
			return surfaceBiome;
		}

		// Mid caves: rare lush/dripstone pockets by climate.
		if (blockY < MID_CAVE_TOP) {
			Holder<Biome> caveBiome = caveBiomeFor(band);
			if (caveBiome != null) {
				double cave = caveNoise.noise(blockX * CAVE_SCALE, 0.0, blockZ * CAVE_SCALE);
				if (cave > CAVE_THRESHOLD) {
					return caveBiome;
				}
			}
		}

		return surfaceBiome;
	}

	/** Lush caves for humid/warm bands, dripstone for dry/cold bands. */
	private Holder<Biome> caveBiomeFor(BiomeBand band) {
		if (band == BiomeBand.FROZEN || band == BiomeBand.COLD) {
			return dripstoneCavesBiome;
		}
		return lushCavesBiome;
	}

	/** Lazily resolve surface biomes from climate tags. Tags are bound after the registry freezes,
	 * so this is safe and never creates unbound holders. */
	private List<Holder<Biome>> resolveSurfaceTag(BiomeBand band) {
		EnumMap<BiomeBand, List<Holder<Biome>>> cache = surfaceFromTag;
		if (cache != null) {
			return cache.get(band);
		}
		synchronized (this) {
			if (surfaceFromTag == null) {
				EnumMap<BiomeBand, List<Holder<Biome>>> built = new EnumMap<>(BiomeBand.class);
				for (BiomeBand b : BiomeBand.values()) {
					List<Holder<Biome>> holders = new ArrayList<>();
					biomeGetter.get(b.surfaceTag()).map(HolderSet::stream).ifPresent(stream -> stream.forEach(holders::add));
					built.put(b, holders);
				}
				surfaceFromTag = built;
			}
			return surfaceFromTag.get(band);
		}
	}

	private Holder<Biome> pickBiome(List<Holder<Biome>> biomes, int blockX, int blockZ, double scale) {
		if (biomes == null || biomes.isEmpty()) {
			return fallback;
		}
		// Dual-octave noise for a flatter, near-uniform index distribution: a single ImprovedNoise
		// value is bell-curved around 0, which makes the middle of the biome list dominate. Adding a
		// second octave at an offset flattens the curve so no biome is disproportionately common.
		double n1 = selectorNoise.noise(blockX * scale, blockZ * scale, 1000.0);
		double n2 = selectorNoise.noise(blockX * scale * 1.9 + 137.0, blockZ * scale * 1.9 - 211.0, 2000.0);
		double combined = (n1 + n2 * 0.5) / 1.5;
		double normalized = combined * 0.5 + 0.5;
		int idx = (int) (normalized * biomes.size());
		if (idx >= biomes.size()) idx = biomes.size() - 1;
		if (idx < 0) idx = 0;
		return biomes.get(idx);
	}

	private static double clamp(double value, double min, double max) {
		return value < min ? min : (value > max ? max : value);
	}

	private record ResolvedBand(List<Holder<Biome>> surface, List<Holder<Biome>> ocean) {
	}
}
