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 javax.annotation.Nullable;
import net.minecraft.core.UUIDUtil;
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 net.minecraft.world.level.storage.ValueInput;
import net.minecraft.world.level.storage.ValueOutput;

public final class EntityReference<StoredEntityType extends Entity> { // Folia - region threading - force Entity Type
    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, org.bukkit.craftbukkit.entity.CraftEntity> entity; // Folia - region threading - force usage of CraftEntity

    public static <Type extends Entity> Codec<EntityReference<Type>> codec() { // Folia - region threading - force Entity Type
        return (Codec<EntityReference<Type>>)CODEC;
    }

    public static <Type extends Entity> StreamCodec<ByteBuf, EntityReference<Type>> streamCodec() { // Folia - region threading - force Entity Type
        return (StreamCodec<ByteBuf, EntityReference<Type>>)STREAM_CODEC;
    }

    public EntityReference(StoredEntityType entity) {
        this.entity = Either.right(entity.getBukkitEntity()); // Folia - region threading - force usage of CraftEntity
    }

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

    public UUID getUUID() {
        return this.entity.map(uuid -> (UUID)uuid, org.bukkit.craftbukkit.entity.CraftEntity::getUniqueId); // Folia - region threading - force usage of CraftEntity
    }

    @Nullable
    public StoredEntityType getEntity(UUIDLookup<? super StoredEntityType> uuidLookup, Class<StoredEntityType> entityClass) {
        Optional<org.bukkit.craftbukkit.entity.CraftEntity> optional = this.entity.right(); // Folia - region threading - force usage of CraftEntity
        if (optional.isPresent()) {
            StoredEntityType uniquelyIdentifyable = (StoredEntityType)optional.get().getHandleRaw(); // Folia - region threading - force usage of CraftEntity
            if (!uniquelyIdentifyable.isRemoved()) {
                return uniquelyIdentifyable;
            }

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

        Optional<UUID> optional1 = this.entity.left();
        if (optional1.isPresent()) {
            StoredEntityType uniquelyIdentifyable1 = this.resolve(uuidLookup.getEntity(optional1.get()), entityClass);
            if (uniquelyIdentifyable1 != null && !uniquelyIdentifyable1.isRemoved()) {
                this.entity = Either.right(uniquelyIdentifyable1.getBukkitEntity()); // Folia - region threading - force usage of CraftEntity
                return uniquelyIdentifyable1;
            }
        }

        return null;
    }

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

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

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

    public static void store(@Nullable EntityReference<?> key, ValueOutput output, String uuid) {
        if (key != null) {
            key.store(output, uuid);
        }
    }

    @Nullable
    public static <StoredEntityType extends Entity> StoredEntityType get( // Folia - region threading - force Entity Type
        @Nullable EntityReference<StoredEntityType> reference, UUIDLookup<? super StoredEntityType> uuidLookup, Class<StoredEntityType> entityClass
    ) {
        return reference != null ? reference.getEntity(uuidLookup, entityClass) : null;
    }

    @Nullable
    public static <StoredEntityType extends Entity> EntityReference<StoredEntityType> read(ValueInput input, String key) { // Folia - region threading - force Entity Type
        return input.read(key, EntityReference.<StoredEntityType>codec()).orElse(null);
    }

    @Nullable
    public static <StoredEntityType extends Entity> EntityReference<StoredEntityType> readWithOldOwnerConversion( // Folia - region threading - force Entity Type
        ValueInput input, String key, Level level
    ) {
        Optional<UUID> optional = input.read(key, UUIDUtil.CODEC);
        return optional.isPresent()
            ? new EntityReference<>(optional.get())
            : input.getString(key)
                .map(string -> OldUsersConverter.convertMobOwnerIfNecessary(level.getServer(), string))
                .map(EntityReference<StoredEntityType>::new)
                .orElse(null);
    }

    @Override
    public boolean equals(Object other) {
        return other == this || other instanceof EntityReference<?> entityReference && this.getUUID().equals(entityReference.getUUID());
    }

    @Override
    public int hashCode() {
        return this.getUUID().hashCode();
    }
}
