package net.fabricmc.fabric.mixin.client.rendering;

import com.llamalad7.mixinextras.injector.ModifyExpressionValue;
import it.unimi.dsi.fastutil.objects.ObjectIterator;
import it.unimi.dsi.fastutil.objects.Reference2ReferenceMap;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.fabricmc.fabric.impl.client.rendering.ColorResolverRegistryImpl;
import net.minecraft.class_1923;
import net.minecraft.class_2338;
import net.minecraft.class_4700;
import net.minecraft.class_638;
import net.minecraft.class_6539;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.Unique;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;

@Environment(EnvType.CLIENT)
@Mixin({class_638.class})
public abstract class ClientWorldMixin {
   @Unique
   private final Reference2ReferenceMap<class_6539, class_4700> customColorCache = ColorResolverRegistryImpl.createCustomCacheMap((resolver) -> new class_4700((pos) -> this.method_23780(pos, resolver)));

   @Shadow
   public abstract int method_23780(class_2338 var1, class_6539 var2);

   @Inject(
      method = {"method_23782(Lnet/minecraft/class_1923;)V"},
      at = {@At("RETURN")}
   )
   private void onResetChunkColor(class_1923 chunkPos, CallbackInfo ci) {
      ObjectIterator var3 = this.customColorCache.values().iterator();

      while(var3.hasNext()) {
         class_4700 cache = (class_4700)var3.next();
         cache.method_23769(chunkPos.field_9181, chunkPos.field_9180);
      }

   }

   @Inject(
      method = {"method_23784()V"},
      at = {@At("RETURN")}
   )
   private void onReloadColor(CallbackInfo ci) {
      ObjectIterator var2 = this.customColorCache.values().iterator();

      while(var2.hasNext()) {
         class_4700 cache = (class_4700)var2.next();
         cache.method_23768();
      }

   }

   @ModifyExpressionValue(
      method = {"method_23752(Lnet/minecraft/class_2338;Lnet/minecraft/class_6539;)I"},
      at = {@At(
   value = "INVOKE",
   target = "Lit/unimi/dsi/fastutil/objects/Object2ObjectArrayMap;get(Ljava/lang/Object;)Ljava/lang/Object;"
)}
   )
   private Object modifyNullCache(Object cache, class_2338 pos, class_6539 resolver) {
      if (cache == null) {
         cache = this.customColorCache.get(resolver);
         if (cache == null) {
            throw new UnsupportedOperationException("ClientWorld.getColor called with unregistered ColorResolver " + String.valueOf(resolver));
         }
      }

      return cache;
   }
}
