package net.minecraft.world.entity.animal;

import net.minecraft.core.Holder;
import net.minecraft.core.component.DataComponentGetter;
import net.minecraft.core.component.DataComponentType;
import net.minecraft.core.component.DataComponents;
import net.minecraft.core.registries.Registries;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.syncher.EntityDataAccessor;
import net.minecraft.network.syncher.EntityDataSerializers;
import net.minecraft.network.syncher.SynchedEntityData;
import net.minecraft.network.syncher.SynchedEntityData.Builder;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.DifficultyInstance;
import net.minecraft.world.entity.AgeableMob;
import net.minecraft.world.entity.EntitySpawnReason;
import net.minecraft.world.entity.EntityType;
import net.minecraft.world.entity.SpawnGroupData;
import net.minecraft.world.entity.variant.SpawnContext;
import net.minecraft.world.entity.variant.VariantUtils;
import net.minecraft.world.level.Level;
import net.minecraft.world.level.ServerLevelAccessor;
import org.jetbrains.annotations.Nullable;

public class Cow extends AbstractCow {
	private static final EntityDataAccessor<Holder<CowVariant>> DATA_VARIANT_ID = SynchedEntityData.defineId(Cow.class, EntityDataSerializers.COW_VARIANT);

	public Cow(EntityType<? extends Cow> entityType, Level level) {
		super(entityType, level);
	}

	@Override
	protected void defineSynchedData(Builder builder) {
		super.defineSynchedData(builder);
		builder.define(DATA_VARIANT_ID, VariantUtils.getDefaultOrAny(this.registryAccess(), CowVariants.TEMPERATE));
	}

	@Override
	public void addAdditionalSaveData(CompoundTag tag) {
		super.addAdditionalSaveData(tag);
		VariantUtils.writeVariant(tag, this.getVariant());
	}

	@Override
	public void readAdditionalSaveData(CompoundTag tag) {
		super.readAdditionalSaveData(tag);
		VariantUtils.readVariant(tag, this.registryAccess(), Registries.COW_VARIANT).ifPresent(this::setVariant);
	}

	@Nullable
	public Cow getBreedOffspring(ServerLevel serverLevel, AgeableMob ageableMob) {
		Cow cow = EntityType.COW.create(serverLevel, EntitySpawnReason.BREEDING);
		if (cow != null && ageableMob instanceof Cow cow2) {
			cow.setVariant(this.random.nextBoolean() ? this.getVariant() : cow2.getVariant());
		}

		return cow;
	}

	@Override
	public SpawnGroupData finalizeSpawn(
		ServerLevelAccessor level, DifficultyInstance difficulty, EntitySpawnReason spawnReason, @Nullable SpawnGroupData spawnGroupData
	) {
		CowVariants.selectVariantToSpawn(this.random, this.registryAccess(), SpawnContext.create(level, this.blockPosition())).ifPresent(this::setVariant);
		return super.finalizeSpawn(level, difficulty, spawnReason, spawnGroupData);
	}

	public void setVariant(Holder<CowVariant> variant) {
		this.entityData.set(DATA_VARIANT_ID, variant);
	}

	public Holder<CowVariant> getVariant() {
		return this.entityData.get(DATA_VARIANT_ID);
	}

	@Nullable
	@Override
	public <T> T get(DataComponentType<? extends T> component) {
		return component == DataComponents.COW_VARIANT ? castComponentValue((DataComponentType<T>)component, this.getVariant()) : super.get(component);
	}

	@Override
	protected void applyImplicitComponents(DataComponentGetter componentGetter) {
		this.applyImplicitComponentIfPresent(componentGetter, DataComponents.COW_VARIANT);
		super.applyImplicitComponents(componentGetter);
	}

	@Override
	protected <T> boolean applyImplicitComponent(DataComponentType<T> component, T value) {
		if (component == DataComponents.COW_VARIANT) {
			this.setVariant(castComponentValue(DataComponents.COW_VARIANT, value));
			return true;
		} else {
			return super.applyImplicitComponent(component, value);
		}
	}
}