package net.minecraft.world.level.levelgen.structure.pools;

import com.google.common.collect.Lists;
import com.mojang.logging.LogUtils;
import java.util.List;
import java.util.Optional;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Direction;
import net.minecraft.core.Holder;
import net.minecraft.core.Registry;
import net.minecraft.core.RegistryAccess;
import net.minecraft.core.Vec3i;
import net.minecraft.core.registries.Registries;
import net.minecraft.data.worldgen.Pools;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.util.RandomSource;
import net.minecraft.util.SequencedPriorityIterator;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.LevelHeightAccessor;
import net.minecraft.world.level.StructureManager;
import net.minecraft.world.level.block.JigsawBlock;
import net.minecraft.world.level.block.Rotation;
import net.minecraft.world.level.chunk.ChunkGenerator;
import net.minecraft.world.level.levelgen.Heightmap;
import net.minecraft.world.level.levelgen.RandomState;
import net.minecraft.world.level.levelgen.WorldgenRandom;
import net.minecraft.world.level.levelgen.structure.BoundingBox;
import net.minecraft.world.level.levelgen.structure.PoolElementStructurePiece;
import net.minecraft.world.level.levelgen.structure.StructurePiece;
import net.minecraft.world.level.levelgen.structure.Structure.GenerationContext;
import net.minecraft.world.level.levelgen.structure.Structure.GenerationStub;
import net.minecraft.world.level.levelgen.structure.pieces.StructurePiecesBuilder;
import net.minecraft.world.level.levelgen.structure.pools.alias.PoolAliasLookup;
import net.minecraft.world.level.levelgen.structure.structures.JigsawStructure;
import net.minecraft.world.level.levelgen.structure.templatesystem.LiquidSettings;
import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplateManager;
import net.minecraft.world.phys.AABB;
import net.minecraft.world.phys.shapes.BooleanOp;
import net.minecraft.world.phys.shapes.Shapes;
import net.minecraft.world.phys.shapes.VoxelShape;
import org.apache.commons.lang3.mutable.MutableObject;
import org.slf4j.Logger;

public class JigsawPlacement {
	static final Logger LOGGER = LogUtils.getLogger();
	private static final int UNSET_HEIGHT = Integer.MIN_VALUE;

	public static Optional<GenerationStub> addPieces(
		GenerationContext context,
		Holder<StructureTemplatePool> startPool,
		Optional<ResourceLocation> startJigsawName,
		int maxDepth,
		BlockPos pos,
		boolean useExpansionHack,
		Optional<Heightmap.Types> projectStartToHeightmap,
		int maxDistanceFromCenter,
		PoolAliasLookup aliasLookup,
		DimensionPadding dimensionPadding,
		LiquidSettings liquidSettings
	) {
		RegistryAccess registryAccess = context.registryAccess();
		ChunkGenerator chunkGenerator = context.chunkGenerator();
		StructureTemplateManager structureTemplateManager = context.structureTemplateManager();
		LevelHeightAccessor levelHeightAccessor = context.heightAccessor();
		WorldgenRandom worldgenRandom = context.random();
		Registry<StructureTemplatePool> registry = registryAccess.lookupOrThrow(Registries.TEMPLATE_POOL);
		Rotation rotation = Rotation.getRandom(worldgenRandom);
		StructureTemplatePool structureTemplatePool = (StructureTemplatePool)startPool.unwrapKey()
			.flatMap(resourceKey -> registry.getOptional(aliasLookup.lookup(resourceKey)))
			.orElse(startPool.value());
		StructurePoolElement structurePoolElement = structureTemplatePool.getRandomTemplate(worldgenRandom);
		if (structurePoolElement == EmptyPoolElement.INSTANCE) {
			return Optional.empty();
		} else {
			BlockPos blockPos;
			if (startJigsawName.isPresent()) {
				ResourceLocation resourceLocation = (ResourceLocation)startJigsawName.get();
				Optional<BlockPos> optional = getRandomNamedJigsaw(structurePoolElement, resourceLocation, pos, rotation, structureTemplateManager, worldgenRandom);
				if (optional.isEmpty()) {
					LOGGER.error(
						"No starting jigsaw {} found in start pool {}",
						resourceLocation,
						startPool.unwrapKey().map(resourceKey -> resourceKey.location().toString()).orElse("<unregistered>")
					);
					return Optional.empty();
				}

				blockPos = (BlockPos)optional.get();
			} else {
				blockPos = pos;
			}

			Vec3i vec3i = blockPos.subtract(pos);
			BlockPos blockPos2 = pos.subtract(vec3i);
			PoolElementStructurePiece poolElementStructurePiece = new PoolElementStructurePiece(
				structureTemplateManager,
				structurePoolElement,
				blockPos2,
				structurePoolElement.getGroundLevelDelta(),
				rotation,
				structurePoolElement.getBoundingBox(structureTemplateManager, blockPos2, rotation),
				liquidSettings
			);
			BoundingBox boundingBox = poolElementStructurePiece.getBoundingBox();
			int i = (boundingBox.maxX() + boundingBox.minX()) / 2;
			int j = (boundingBox.maxZ() + boundingBox.minZ()) / 2;
			int k = projectStartToHeightmap.isEmpty()
				? blockPos2.getY()
				: pos.getY() + chunkGenerator.getFirstFreeHeight(i, j, (Heightmap.Types)projectStartToHeightmap.get(), levelHeightAccessor, context.randomState());
			int l = boundingBox.minY() + poolElementStructurePiece.getGroundLevelDelta();
			poolElementStructurePiece.move(0, k - l, 0);
			if (isStartTooCloseToWorldHeightLimits(levelHeightAccessor, dimensionPadding, poolElementStructurePiece.getBoundingBox())) {
				LOGGER.debug(
					"Center piece {} with bounding box {} does not fit dimension padding {}",
					structurePoolElement,
					poolElementStructurePiece.getBoundingBox(),
					dimensionPadding
				);
				return Optional.empty();
			} else {
				int m = k + vec3i.getY();
				return Optional.of(
					new GenerationStub(
						new BlockPos(i, m, j),
						structurePiecesBuilder -> {
							List<PoolElementStructurePiece> list = Lists.<PoolElementStructurePiece>newArrayList();
							list.add(poolElementStructurePiece);
							if (maxDepth > 0) {
								AABB aABB = new AABB(
									i - maxDistanceFromCenter,
									Math.max(m - maxDistanceFromCenter, levelHeightAccessor.getMinY() + dimensionPadding.bottom()),
									j - maxDistanceFromCenter,
									i + maxDistanceFromCenter + 1,
									Math.min(m + maxDistanceFromCenter + 1, levelHeightAccessor.getMaxY() + 1 - dimensionPadding.top()),
									j + maxDistanceFromCenter + 1
								);
								VoxelShape voxelShape = Shapes.join(Shapes.create(aABB), Shapes.create(AABB.of(boundingBox)), BooleanOp.ONLY_FIRST);
								addPieces(
									context.randomState(),
									maxDepth,
									useExpansionHack,
									chunkGenerator,
									structureTemplateManager,
									levelHeightAccessor,
									worldgenRandom,
									registry,
									poolElementStructurePiece,
									list,
									voxelShape,
									aliasLookup,
									liquidSettings
								);
								list.forEach(structurePiecesBuilder::addPiece);
							}
						}
					)
				);
			}
		}
	}

	private static boolean isStartTooCloseToWorldHeightLimits(LevelHeightAccessor level, DimensionPadding padding, BoundingBox boundingBox) {
		if (padding == DimensionPadding.ZERO) {
			return false;
		} else {
			int i = level.getMinY() + padding.bottom();
			int j = level.getMaxY() - padding.top();
			return boundingBox.minY() < i || boundingBox.maxY() > j;
		}
	}

	private static Optional<BlockPos> getRandomNamedJigsaw(
		StructurePoolElement element,
		ResourceLocation startJigsawName,
		BlockPos pos,
		Rotation rotation,
		StructureTemplateManager structureTemplateManager,
		WorldgenRandom random
	) {
		for (StructureTemplate.JigsawBlockInfo jigsawBlockInfo : element.getShuffledJigsawBlocks(structureTemplateManager, pos, rotation, random)) {
			if (startJigsawName.equals(jigsawBlockInfo.name())) {
				return Optional.of(jigsawBlockInfo.info().pos());
			}
		}

		return Optional.empty();
	}

	private static void addPieces(
		RandomState randomState,
		int maxDepth,
		boolean useExpansionHack,
		ChunkGenerator chunkGenerator,
		StructureTemplateManager structureTemplateManager,
		LevelHeightAccessor level,
		RandomSource random,
		Registry<StructureTemplatePool> pools,
		PoolElementStructurePiece startPiece,
		List<PoolElementStructurePiece> pieces,
		VoxelShape free,
		PoolAliasLookup aliasLookup,
		LiquidSettings liquidSettings
	) {
		JigsawPlacement.Placer placer = new JigsawPlacement.Placer(pools, maxDepth, chunkGenerator, structureTemplateManager, pieces, random);
		placer.tryPlacingChildren(startPiece, new MutableObject<>(free), 0, useExpansionHack, level, randomState, aliasLookup, liquidSettings);

		while (placer.placing.hasNext()) {
			JigsawPlacement.PieceState pieceState = placer.placing.next();
			placer.tryPlacingChildren(pieceState.piece, pieceState.free, pieceState.depth, useExpansionHack, level, randomState, aliasLookup, liquidSettings);
		}
	}

	public static boolean generateJigsaw(
		ServerLevel level, Holder<StructureTemplatePool> startPool, ResourceLocation startJigsawName, int maxDepth, BlockPos pos, boolean keepJigsaws
	) {
		ChunkGenerator chunkGenerator = level.getChunkSource().getGenerator();
		StructureTemplateManager structureTemplateManager = level.getStructureManager();
		StructureManager structureManager = level.structureManager();
		RandomSource randomSource = level.getRandom();
		GenerationContext generationContext = new GenerationContext(
			level.registryAccess(),
			chunkGenerator,
			chunkGenerator.getBiomeSource(),
			level.getChunkSource().randomState(),
			structureTemplateManager,
			level.getSeed(),
			new ChunkPos(pos),
			level,
			holder -> true
		);
		Optional<GenerationStub> optional = addPieces(
			generationContext,
			startPool,
			Optional.of(startJigsawName),
			maxDepth,
			pos,
			false,
			Optional.empty(),
			128,
			PoolAliasLookup.EMPTY,
			JigsawStructure.DEFAULT_DIMENSION_PADDING,
			JigsawStructure.DEFAULT_LIQUID_SETTINGS
		);
		if (optional.isPresent()) {
			StructurePiecesBuilder structurePiecesBuilder = ((GenerationStub)optional.get()).getPiecesBuilder();

			for (StructurePiece structurePiece : structurePiecesBuilder.build().pieces()) {
				if (structurePiece instanceof PoolElementStructurePiece poolElementStructurePiece) {
					poolElementStructurePiece.place(level, structureManager, chunkGenerator, randomSource, BoundingBox.infinite(), pos, keepJigsaws);
				}
			}

			return true;
		} else {
			return false;
		}
	}

	record PieceState(PoolElementStructurePiece piece, MutableObject<VoxelShape> free, int depth) {
	}

	static final class Placer {
		private final Registry<StructureTemplatePool> pools;
		private final int maxDepth;
		private final ChunkGenerator chunkGenerator;
		private final StructureTemplateManager structureTemplateManager;
		private final List<? super PoolElementStructurePiece> pieces;
		private final RandomSource random;
		final SequencedPriorityIterator<JigsawPlacement.PieceState> placing = new SequencedPriorityIterator<>();

		Placer(
			Registry<StructureTemplatePool> pools,
			int maxDepth,
			ChunkGenerator chunkGenerator,
			StructureTemplateManager structureTemplateManager,
			List<? super PoolElementStructurePiece> pieces,
			RandomSource random
		) {
			this.pools = pools;
			this.maxDepth = maxDepth;
			this.chunkGenerator = chunkGenerator;
			this.structureTemplateManager = structureTemplateManager;
			this.pieces = pieces;
			this.random = random;
		}

		void tryPlacingChildren(
			PoolElementStructurePiece piece,
			MutableObject<VoxelShape> free,
			int depth,
			boolean useExpansionHack,
			LevelHeightAccessor level,
			RandomState random,
			PoolAliasLookup poolAliasLookup,
			LiquidSettings liquidSettings
		) {
			StructurePoolElement structurePoolElement = piece.getElement();
			BlockPos blockPos = piece.getPosition();
			Rotation rotation = piece.getRotation();
			StructureTemplatePool.Projection projection = structurePoolElement.getProjection();
			boolean bl = projection == StructureTemplatePool.Projection.RIGID;
			MutableObject<VoxelShape> mutableObject = new MutableObject<>();
			BoundingBox boundingBox = piece.getBoundingBox();
			int i = boundingBox.minY();

			label129:
			for (StructureTemplate.JigsawBlockInfo jigsawBlockInfo : structurePoolElement.getShuffledJigsawBlocks(
				this.structureTemplateManager, blockPos, rotation, this.random
			)) {
				StructureTemplate.StructureBlockInfo structureBlockInfo = jigsawBlockInfo.info();
				Direction direction = JigsawBlock.getFrontFacing(structureBlockInfo.state());
				BlockPos blockPos2 = structureBlockInfo.pos();
				BlockPos blockPos3 = blockPos2.relative(direction);
				int j = blockPos2.getY() - i;
				int k = Integer.MIN_VALUE;
				ResourceKey<StructureTemplatePool> resourceKey = poolAliasLookup.lookup(jigsawBlockInfo.pool());
				Optional<? extends Holder<StructureTemplatePool>> optional = this.pools.get(resourceKey);
				if (optional.isEmpty()) {
					JigsawPlacement.LOGGER.warn("Empty or non-existent pool: {}", resourceKey.location());
				} else {
					Holder<StructureTemplatePool> holder = (Holder<StructureTemplatePool>)optional.get();
					if (holder.value().size() == 0 && !holder.is(Pools.EMPTY)) {
						JigsawPlacement.LOGGER.warn("Empty or non-existent pool: {}", resourceKey.location());
					} else {
						Holder<StructureTemplatePool> holder2 = holder.value().getFallback();
						if (holder2.value().size() == 0 && !holder2.is(Pools.EMPTY)) {
							JigsawPlacement.LOGGER
								.warn("Empty or non-existent fallback pool: {}", holder2.unwrapKey().map(resourceKeyx -> resourceKeyx.location().toString()).orElse("<unregistered>"));
						} else {
							boolean bl2 = boundingBox.isInside(blockPos3);
							MutableObject<VoxelShape> mutableObject2;
							if (bl2) {
								mutableObject2 = mutableObject;
								if (mutableObject.getValue() == null) {
									mutableObject.setValue(Shapes.create(AABB.of(boundingBox)));
								}
							} else {
								mutableObject2 = free;
							}

							List<StructurePoolElement> list = Lists.<StructurePoolElement>newArrayList();
							if (depth != this.maxDepth) {
								list.addAll(holder.value().getShuffledTemplates(this.random));
							}

							list.addAll(holder2.value().getShuffledTemplates(this.random));
							int l = jigsawBlockInfo.placementPriority();

							for (StructurePoolElement structurePoolElement2 : list) {
								if (structurePoolElement2 == EmptyPoolElement.INSTANCE) {
									break;
								}

								for (Rotation rotation2 : Rotation.getShuffled(this.random)) {
									List<StructureTemplate.JigsawBlockInfo> list2 = structurePoolElement2.getShuffledJigsawBlocks(
										this.structureTemplateManager, BlockPos.ZERO, rotation2, this.random
									);
									BoundingBox boundingBox2 = structurePoolElement2.getBoundingBox(this.structureTemplateManager, BlockPos.ZERO, rotation2);
									int m;
									if (useExpansionHack && boundingBox2.getYSpan() <= 16) {
										m = list2.stream().mapToInt(jigsawBlockInfox -> {
											StructureTemplate.StructureBlockInfo structureBlockInfox = jigsawBlockInfox.info();
											if (!boundingBox2.isInside(structureBlockInfox.pos().relative(JigsawBlock.getFrontFacing(structureBlockInfox.state())))) {
												return 0;
											} else {
												ResourceKey<StructureTemplatePool> resourceKeyx = poolAliasLookup.lookup(jigsawBlockInfox.pool());
												Optional<? extends Holder<StructureTemplatePool>> optionalx = this.pools.get(resourceKeyx);
												Optional<Holder<StructureTemplatePool>> optional2 = optionalx.map(holderx -> ((StructureTemplatePool)holderx.value()).getFallback());
												int ix = (Integer)optionalx.map(holderx -> ((StructureTemplatePool)holderx.value()).getMaxSize(this.structureTemplateManager)).orElse(0);
												int jx = (Integer)optional2.map(holderx -> ((StructureTemplatePool)holderx.value()).getMaxSize(this.structureTemplateManager)).orElse(0);
												return Math.max(ix, jx);
											}
										}).max().orElse(0);
									} else {
										m = 0;
									}

									for (StructureTemplate.JigsawBlockInfo jigsawBlockInfo2 : list2) {
										if (JigsawBlock.canAttach(jigsawBlockInfo, jigsawBlockInfo2)) {
											BlockPos blockPos4 = jigsawBlockInfo2.info().pos();
											BlockPos blockPos5 = blockPos3.subtract(blockPos4);
											BoundingBox boundingBox3 = structurePoolElement2.getBoundingBox(this.structureTemplateManager, blockPos5, rotation2);
											int n = boundingBox3.minY();
											StructureTemplatePool.Projection projection2 = structurePoolElement2.getProjection();
											boolean bl3 = projection2 == StructureTemplatePool.Projection.RIGID;
											int o = blockPos4.getY();
											int p = j - o + JigsawBlock.getFrontFacing(structureBlockInfo.state()).getStepY();
											int q;
											if (bl && bl3) {
												q = i + p;
											} else {
												if (k == Integer.MIN_VALUE) {
													k = this.chunkGenerator.getFirstFreeHeight(blockPos2.getX(), blockPos2.getZ(), Heightmap.Types.WORLD_SURFACE_WG, level, random);
												}

												q = k - o;
											}

											int r = q - n;
											BoundingBox boundingBox4 = boundingBox3.moved(0, r, 0);
											BlockPos blockPos6 = blockPos5.offset(0, r, 0);
											if (m > 0) {
												int s = Math.max(m + 1, boundingBox4.maxY() - boundingBox4.minY());
												boundingBox4.encapsulate(new BlockPos(boundingBox4.minX(), boundingBox4.minY() + s, boundingBox4.minZ()));
											}

											if (!Shapes.joinIsNotEmpty(mutableObject2.getValue(), Shapes.create(AABB.of(boundingBox4).deflate(0.25)), BooleanOp.ONLY_SECOND)) {
												mutableObject2.setValue(Shapes.joinUnoptimized(mutableObject2.getValue(), Shapes.create(AABB.of(boundingBox4)), BooleanOp.ONLY_FIRST));
												int s = piece.getGroundLevelDelta();
												int t;
												if (bl3) {
													t = s - p;
												} else {
													t = structurePoolElement2.getGroundLevelDelta();
												}

												PoolElementStructurePiece poolElementStructurePiece = new PoolElementStructurePiece(
													this.structureTemplateManager, structurePoolElement2, blockPos6, t, rotation2, boundingBox4, liquidSettings
												);
												int u;
												if (bl) {
													u = i + j;
												} else if (bl3) {
													u = q + o;
												} else {
													if (k == Integer.MIN_VALUE) {
														k = this.chunkGenerator.getFirstFreeHeight(blockPos2.getX(), blockPos2.getZ(), Heightmap.Types.WORLD_SURFACE_WG, level, random);
													}

													u = k + p / 2;
												}

												piece.addJunction(new JigsawJunction(blockPos3.getX(), u - j + s, blockPos3.getZ(), p, projection2));
												poolElementStructurePiece.addJunction(new JigsawJunction(blockPos2.getX(), u - o + t, blockPos2.getZ(), -p, projection));
												this.pieces.add(poolElementStructurePiece);
												if (depth + 1 <= this.maxDepth) {
													JigsawPlacement.PieceState pieceState = new JigsawPlacement.PieceState(poolElementStructurePiece, mutableObject2, depth + 1);
													this.placing.add(pieceState, l);
												}
												continue label129;
											}
										}
									}
								}
							}
						}
					}
				}
			}
		}
	}
}