package net.minecraft.world.entity;

import com.mojang.datafixers.util.Either;
import com.mojang.serialization.Codec;
import io.netty.buffer.ByteBuf;
import java.util.Optional;
import java.util.UUID;
import net.minecraft.core.UUIDUtil;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.server.players.OldUsersConverter;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.entity.UUIDLookup;
import net.minecraft.world.level.entity.UniquelyIdentifyable;
import org.jetbrains.annotations.Nullable;

public class EntityReference<StoredEntityType extends UniquelyIdentifyable> {
	private static final Codec<? extends EntityReference<?>> CODEC = UUIDUtil.CODEC.xmap(EntityReference::new, EntityReference::getUUID);
	private static final StreamCodec<ByteBuf, ? extends EntityReference<?>> STREAM_CODEC = UUIDUtil.STREAM_CODEC
		.map(EntityReference::new, EntityReference::getUUID);
	private Either<UUID, StoredEntityType> entity;

	public static <Type extends UniquelyIdentifyable> Codec<EntityReference<Type>> codec() {
		return (Codec<EntityReference<Type>>)CODEC;
	}

	public static <Type extends UniquelyIdentifyable> StreamCodec<ByteBuf, EntityReference<Type>> streamCodec() {
		return (StreamCodec<ByteBuf, EntityReference<Type>>)STREAM_CODEC;
	}

	public EntityReference(StoredEntityType entity) {
		this.entity = Either.right(entity);
	}

	public EntityReference(UUID uuid) {
		this.entity = Either.left(uuid);
	}

	public UUID getUUID() {
		return this.entity.map(uUID -> uUID, UniquelyIdentifyable::getUUID);
	}

	@Nullable
	public StoredEntityType getEntity(UUIDLookup<? super StoredEntityType> uuidLookup, Class<StoredEntityType> entityClass) {
		Optional<StoredEntityType> optional = this.entity.right();
		if (optional.isPresent()) {
			StoredEntityType uniquelyIdentifyable = (StoredEntityType)optional.get();
			if (!uniquelyIdentifyable.isRemoved()) {
				return uniquelyIdentifyable;
			}

			this.entity = Either.left(uniquelyIdentifyable.getUUID());
		}

		Optional<UUID> optional2 = this.entity.left();
		if (optional2.isPresent()) {
			StoredEntityType uniquelyIdentifyable2 = this.resolve(uuidLookup.getEntity((UUID)optional2.get()), entityClass);
			if (uniquelyIdentifyable2 != null && !uniquelyIdentifyable2.isRemoved()) {
				this.entity = Either.right(uniquelyIdentifyable2);
				return uniquelyIdentifyable2;
			}
		}

		return null;
	}

	@Nullable
	private StoredEntityType resolve(@Nullable UniquelyIdentifyable entity, Class<StoredEntityType> entityClass) {
		return (StoredEntityType)(entity != null && entityClass.isAssignableFrom(entity.getClass()) ? entityClass.cast(entity) : null);
	}

	public boolean matches(StoredEntityType entity) {
		return this.getUUID().equals(entity.getUUID());
	}

	public void store(CompoundTag tag, String key) {
		tag.store(key, UUIDUtil.CODEC, this.getUUID());
	}

	@Nullable
	public static <StoredEntityType extends UniquelyIdentifyable> StoredEntityType get(
		@Nullable EntityReference<StoredEntityType> reference, UUIDLookup<? super StoredEntityType> uuidLookup, Class<StoredEntityType> entityClass
	) {
		return reference != null ? reference.getEntity(uuidLookup, entityClass) : null;
	}

	@Nullable
	public static <StoredEntityType extends UniquelyIdentifyable> EntityReference<StoredEntityType> read(CompoundTag tag, String key) {
		return (EntityReference<StoredEntityType>)tag.read(key, codec()).orElse(null);
	}

	@Nullable
	public static <StoredEntityType extends UniquelyIdentifyable> EntityReference<StoredEntityType> readWithOldOwnerConversion(
		CompoundTag tag, String key, Level level
	) {
		Optional<UUID> optional = tag.read(key, UUIDUtil.CODEC);
		return optional.isPresent()
			? new EntityReference<>((UUID)optional.get())
			: (EntityReference)tag.getString(key)
				.map(string -> OldUsersConverter.convertMobOwnerIfNecessary(level.getServer(), string))
				.map(EntityReference::new)
				.orElse(null);
	}
}