package net.fabricmc.fabric.impl.client.rendering.hud;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Objects;
import java.util.SequencedCollection;
import java.util.SequencedSet;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.IntBinaryOperator;
import java.util.function.ToIntFunction;
import net.fabricmc.api.ClientModInitializer;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientLifecycleEvents;
import net.fabricmc.fabric.api.client.rendering.v1.hud.HudElementRegistry;
import net.fabricmc.fabric.api.client.rendering.v1.hud.StatusBarHeightProvider;
import net.fabricmc.fabric.api.client.rendering.v1.hud.VanillaHudElements;
import net.fabricmc.fabric.mixin.client.rendering.InGameHudAccessor;
import net.minecraft.class_1309;
import net.minecraft.class_1657;
import net.minecraft.class_2960;
import net.minecraft.class_310;
import net.minecraft.class_329;
import net.minecraft.class_3486;
import net.minecraft.class_3532;
import net.minecraft.class_5134;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.annotations.ApiStatus.NonExtendable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Environment(EnvType.CLIENT)
public final class HudStatusBarHeightRegistryImpl implements ClientModInitializer {
   public static final Logger LOGGER = LoggerFactory.getLogger("fabric-rendering-v1");
   static final int DEFAULT_HEIGHT = 39;
   static final int HELD_ITEM_TOOLTIP_HEIGHT = 20;
   static final int OVERLAY_MESSAGE_HEIGHT = 29;
   static final int TEXT_HEIGHT_DELTA = 9;
   static final StatusBarHeightProvider HEALTH_BAR = (player) -> {
      class_329 hud = class_310.method_1551().field_1705;
      int playerHealth = class_3532.method_15386(player.method_6032());
      int displayHealth = ((InGameHudAccessor)hud).fabric$getRenderHealthValue();
      float maxHealth = Math.max((float)player.method_45325(class_5134.field_23716), (float)Math.max(displayHealth, playerHealth));
      int absorptionAmount = class_3532.method_15386(player.method_6067());
      int healthRows = class_3532.method_15386((maxHealth + (float)absorptionAmount) / 2.0F / 10.0F);
      int rowShift = Math.max(10 - (healthRows - 2), 3);
      return 10 + (healthRows - 1) * rowShift;
   };
   static final StatusBarHeightProvider ARMOR_BAR = (player) -> player.method_6096() > 0 ? 10 : 0;
   static final StatusBarHeightProvider MOUNT_HEALTH = (player) -> {
      class_329 hud = class_310.method_1551().field_1705;
      class_1309 livingEntity = ((InGameHudAccessor)hud).fabric$callGetRiddenEntity();
      int vehicleMaxHearts = ((InGameHudAccessor)hud).fabric$callGetHeartCount(livingEntity);
      return ((InGameHudAccessor)hud).fabric$callGetHeartRows(vehicleMaxHearts) * 10;
   };
   static final StatusBarHeightProvider FOOD_BAR = (player) -> {
      class_329 hud = class_310.method_1551().field_1705;
      class_1309 livingEntity = ((InGameHudAccessor)hud).fabric$callGetRiddenEntity();
      return ((InGameHudAccessor)hud).fabric$callGetHeartCount(livingEntity) == 0 ? 10 : 0;
   };
   static final StatusBarHeightProvider AIR_BAR = (player) -> {
      int maxAirSupply = player.method_5748();
      int airSupply = Math.clamp((long)player.method_5669(), 0, maxAirSupply);
      boolean isInWater = player.method_5777(class_3486.field_15517);
      return !isInWater && airSupply >= maxAirSupply ? 0 : 10;
   };
   static final Map<class_2960, ResolvedHeightProvider> RESOLVED_VANILLA_HEIGHT_PROVIDERS;
   static final Map<class_2960, StatusBarHeightProvider> LEFT_VANILLA_HEIGHT_PROVIDERS;
   static final Map<class_2960, StatusBarHeightProvider> RIGHT_VANILLA_HEIGHT_PROVIDERS;
   static final Map<class_2960, StatusBarHeightProvider> LEFT_HEIGHT_PROVIDERS;
   static final Map<class_2960, StatusBarHeightProvider> RIGHT_HEIGHT_PROVIDERS;
   static @Nullable Map<class_2960, ResolvedHeightProvider> resolvedHeightProviders;

   public void onInitializeClient() {
      ClientLifecycleEvents.CLIENT_STARTED.register((ClientLifecycleEvents.ClientStarted)(minecraft) -> init());
   }

   public static void addLeft(class_2960 id, StatusBarHeightProvider heightProvider) {
      if (resolvedHeightProviders == null) {
         LEFT_HEIGHT_PROVIDERS.put(id, heightProvider);
      } else {
         throw new IllegalStateException("Height provider registry already frozen!");
      }
   }

   public static void addRight(class_2960 id, StatusBarHeightProvider heightProvider) {
      if (resolvedHeightProviders == null) {
         RIGHT_HEIGHT_PROVIDERS.put(id, heightProvider);
      } else {
         throw new IllegalStateException("Height provider registry already frozen!");
      }
   }

   public static int getHeight(class_2960 id) {
      if (resolvedHeightProviders == null) {
         throw new IllegalStateException("Trying to get status bar height for " + String.valueOf(id) + " too early");
      } else if (!resolvedHeightProviders.containsKey(id)) {
         throw new IllegalArgumentException("Unknown status bar: " + String.valueOf(id));
      } else {
         class_1657 player = ((InGameHudAccessor)class_310.method_1551().field_1705).fabric$callGetCameraPlayer();
         if (player == null) {
            throw new IllegalStateException("Trying to get status bar height for " + String.valueOf(id) + " without a camera player");
         } else {
            return 39 + ((ResolvedHeightProvider)resolvedHeightProviders.get(id)).getResolvedHeight(player);
         }
      }
   }

   static void init() {
      if (LEFT_VANILLA_HEIGHT_PROVIDERS.equals(LEFT_HEIGHT_PROVIDERS) && RIGHT_VANILLA_HEIGHT_PROVIDERS.equals(RIGHT_HEIGHT_PROVIDERS)) {
         HudStatusBarHeightRegistryImpl.resolvedHeightProviders = RESOLVED_VANILLA_HEIGHT_PROVIDERS;
      } else {
         Map<class_2960, ResolvedHeightProvider> resolvedHeightProviders = new LinkedHashMap();
         Map var10000 = LEFT_HEIGHT_PROVIDERS;
         Objects.requireNonNull(resolvedHeightProviders);
         ResolvedHeightProvider maxLeftHeightProvider = resolveHeightProviders(var10000, resolvedHeightProviders::put);
         var10000 = RIGHT_HEIGHT_PROVIDERS;
         Objects.requireNonNull(resolvedHeightProviders);
         ResolvedHeightProvider maxRightHeightProvider = resolveHeightProviders(var10000, resolvedHeightProviders::put);
         applyVanillaHeightProviders(resolvedHeightProviders, reduceToIntFunctions(maxLeftHeightProvider, maxRightHeightProvider, Math::max));
         HudStatusBarHeightRegistryImpl.resolvedHeightProviders = ImmutableMap.copyOf(resolvedHeightProviders);
      }

   }

   private static ResolvedHeightProvider resolveHeightProviders(Map<class_2960, StatusBarHeightProvider> heightProviderLookup, BiConsumer<class_2960, ResolvedHeightProvider> heightProviderConsumer) {
      SequencedSet<class_2960> orderedHeightProviders = getOrderedHeightProviders(heightProviderLookup);
      Set<class_2960> unregisteredHudElements = Sets.difference(heightProviderLookup.keySet(), orderedHeightProviders);
      if (!unregisteredHudElements.isEmpty()) {
         throw new IllegalStateException("Unregistered hud elements: " + String.valueOf(unregisteredHudElements));
      } else {
         for(class_2960 id : heightProviderLookup.keySet()) {
            ResolvedHeightProvider heightProvider = resolveHeightProvider(id, heightProviderLookup, orderedHeightProviders);
            heightProviderConsumer.accept(id, heightProvider);
         }

         return resolveMaximumHeightProvider((class_2960)orderedHeightProviders.getLast(), heightProviderLookup, orderedHeightProviders);
      }
   }

   private static SequencedSet<class_2960> getOrderedHeightProviders(Map<class_2960, StatusBarHeightProvider> heightProviderLookup) {
      SequencedSet<class_2960> orderedHeightProviders = new LinkedHashSet();

      for(class_2960 id : RESOLVED_VANILLA_HEIGHT_PROVIDERS.keySet()) {
         for(HudLayer hudLayer : ((HudElementRegistryImpl.RootLayer)HudElementRegistryImpl.ROOT_ELEMENTS.get(id)).layers()) {
            Objects.requireNonNull(orderedHeightProviders);
            addOrderedHeightProvider(hudLayer, heightProviderLookup, orderedHeightProviders::add);
         }
      }

      for(Map.Entry<class_2960, HudElementRegistryImpl.RootLayer> entry : HudElementRegistryImpl.ROOT_ELEMENTS.entrySet()) {
         if (!RESOLVED_VANILLA_HEIGHT_PROVIDERS.containsKey(entry.getKey())) {
            for(HudLayer hudLayer : ((HudElementRegistryImpl.RootLayer)entry.getValue()).layers()) {
               Objects.requireNonNull(orderedHeightProviders);
               addOrderedHeightProvider(hudLayer, heightProviderLookup, orderedHeightProviders::add);
            }
         }
      }

      return orderedHeightProviders;
   }

   private static void addOrderedHeightProvider(HudLayer hudLayer, Map<class_2960, StatusBarHeightProvider> heightProviderLookup, Consumer<class_2960> heightProviderConsumer) {
      if (!hudLayer.isRemoved() && heightProviderLookup.containsKey(hudLayer.id())) {
         heightProviderConsumer.accept(hudLayer.id());
      }

   }

   private static ResolvedHeightProvider resolveHeightProvider(class_2960 id, Map<class_2960, StatusBarHeightProvider> heightProviderLookup, SequencedCollection<class_2960> orderedHeightProviders) {
      ResolvedHeightProvider heightProvider = HudStatusBarHeightRegistryImpl.ResolvedHeightProvider.ZERO;

      for(class_2960 heightProviderLocation : orderedHeightProviders) {
         if (heightProviderLocation.equals(id)) {
            return heightProvider;
         }

         if (heightProviderLookup.containsKey(heightProviderLocation)) {
            heightProvider = reduceToIntFunctions(heightProvider, (ToIntFunction)heightProviderLookup.get(heightProviderLocation), Integer::sum);
         }
      }

      throw new IllegalStateException("Unknown height provider: " + String.valueOf(id));
   }

   private static ResolvedHeightProvider resolveMaximumHeightProvider(class_2960 id, Map<class_2960, StatusBarHeightProvider> heightProviderLookup, SequencedCollection<class_2960> orderedHeightProviders) {
      ResolvedHeightProvider heightProvider = resolveHeightProvider(id, heightProviderLookup, orderedHeightProviders);
      return reduceToIntFunctions((ToIntFunction)heightProviderLookup.get(id), heightProvider, Integer::sum);
   }

   private static ResolvedHeightProvider reduceToIntFunctions(ToIntFunction<class_1657> first, ToIntFunction<class_1657> second, IntBinaryOperator operator) {
      return (player) -> operator.applyAsInt(first.applyAsInt(player), second.applyAsInt(player));
   }

   private static void applyVanillaHeightProviders(Map<class_2960, ResolvedHeightProvider> resolvedHeightProviders, ResolvedHeightProvider maxHeightProvider) {
      for(Map.Entry<class_2960, ResolvedHeightProvider> entry : RESOLVED_VANILLA_HEIGHT_PROVIDERS.entrySet()) {
         if (isVanillaHeightProvider((class_2960)entry.getKey())) {
            ResolvedHeightProvider expectedHeightProvider = (ResolvedHeightProvider)entry.getValue();
            ResolvedHeightProvider actualHeightProvider = (ResolvedHeightProvider)resolvedHeightProviders.put((class_2960)entry.getKey(), expectedHeightProvider);
            Objects.requireNonNull(actualHeightProvider, () -> "resolved height provider " + String.valueOf(entry.getKey()) + " is null");
            replaceVanillaElement((class_2960)entry.getKey(), reduceToIntFunctions(expectedHeightProvider, actualHeightProvider, (i1, i2) -> i1 - i2));
         } else {
            LOGGER.debug("Skipped wrapping hud element {} for applying height provider offsets", entry.getKey());
         }
      }

      replaceVanillaElement(VanillaHudElements.HELD_ITEM_TOOLTIP, (player) -> 20 - Math.max(20, maxHeightProvider.getResolvedHeight(player)));
      replaceVanillaElement(VanillaHudElements.OVERLAY_MESSAGE, (player) -> 29 - Math.max(29, maxHeightProvider.getResolvedHeight(player) + 9));
   }

   private static boolean isVanillaHeightProvider(class_2960 id) {
      if (LEFT_HEIGHT_PROVIDERS.containsKey(id) && LEFT_HEIGHT_PROVIDERS.get(id) == LEFT_VANILLA_HEIGHT_PROVIDERS.get(id)) {
         return true;
      } else {
         return RIGHT_HEIGHT_PROVIDERS.containsKey(id) && RIGHT_HEIGHT_PROVIDERS.get(id) == RIGHT_VANILLA_HEIGHT_PROVIDERS.get(id);
      }
   }

   private static void replaceVanillaElement(class_2960 id, ResolvedHeightProvider heightProvider) {
      HudElementRegistry.replaceElement(id, (layer) -> (context, tickCounter) -> {
            class_1657 player = ((InGameHudAccessor)class_310.method_1551().field_1705).fabric$callGetCameraPlayer();
            int height = player != null ? heightProvider.getResolvedHeight(player) : 0;
            if (height != 0) {
               context.method_51448().pushMatrix();
               context.method_51448().translate(0.0F, (float)height);
            }

            layer.render(context, tickCounter);
            if (height != 0) {
               context.method_51448().popMatrix();
            }

         });
   }

   static {
      class_2960 var10000 = VanillaHudElements.HEALTH_BAR;
      ResolvedHeightProvider var10001 = HudStatusBarHeightRegistryImpl.ResolvedHeightProvider.ZERO;
      class_2960 var10002 = VanillaHudElements.ARMOR_BAR;
      StatusBarHeightProvider var10003 = HEALTH_BAR;
      Objects.requireNonNull(var10003);
      RESOLVED_VANILLA_HEIGHT_PROVIDERS = ImmutableMap.of(var10000, var10001, var10002, var10003::getStatusBarHeight, VanillaHudElements.MOUNT_HEALTH, HudStatusBarHeightRegistryImpl.ResolvedHeightProvider.ZERO, VanillaHudElements.FOOD_BAR, HudStatusBarHeightRegistryImpl.ResolvedHeightProvider.ZERO, VanillaHudElements.AIR_BAR, reduceToIntFunctions(MOUNT_HEALTH, FOOD_BAR, Integer::sum));
      LEFT_VANILLA_HEIGHT_PROVIDERS = ImmutableMap.of(VanillaHudElements.HEALTH_BAR, HEALTH_BAR, VanillaHudElements.ARMOR_BAR, ARMOR_BAR);
      RIGHT_VANILLA_HEIGHT_PROVIDERS = ImmutableMap.of(VanillaHudElements.MOUNT_HEALTH, MOUNT_HEALTH, VanillaHudElements.FOOD_BAR, FOOD_BAR, VanillaHudElements.AIR_BAR, AIR_BAR);
      LEFT_HEIGHT_PROVIDERS = new HashMap(LEFT_VANILLA_HEIGHT_PROVIDERS);
      RIGHT_HEIGHT_PROVIDERS = new HashMap(RIGHT_VANILLA_HEIGHT_PROVIDERS);
   }

   @FunctionalInterface
   @Environment(EnvType.CLIENT)
   public interface ResolvedHeightProvider extends ToIntFunction<class_1657> {
      ResolvedHeightProvider ZERO = (player) -> 0;

      int getResolvedHeight(class_1657 var1);

      @NonExtendable
      default int applyAsInt(class_1657 player) {
         return this.getResolvedHeight(player);
      }
   }
}