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

import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;
import net.fabricmc.api.EnvType;
import net.fabricmc.api.Environment;
import net.fabricmc.fabric.api.client.rendering.v1.hud.HudElement;
import net.fabricmc.fabric.api.client.rendering.v1.hud.VanillaHudElements;
import net.minecraft.class_2960;
import net.minecraft.class_332;
import net.minecraft.class_9779;
import org.apache.commons.lang3.mutable.MutableBoolean;
import org.jetbrains.annotations.VisibleForTesting;

@Environment(EnvType.CLIENT)
public class HudElementRegistryImpl {
   @VisibleForTesting
   static final List<class_2960> VANILLA_ELEMENT_IDS;
   @VisibleForTesting
   public static final Map<class_2960, RootLayer> ROOT_ELEMENTS;
   private static final RootLayer FIRST;
   private static final RootLayer LAST;

   public static RootLayer getRoot(class_2960 id) {
      return (RootLayer)ROOT_ELEMENTS.get(id);
   }

   public static void addFirst(class_2960 id, HudElement element) {
      validateUnique(id);
      FIRST.layers().addFirst(HudLayer.ofElement(id, element));
   }

   public static void addLast(class_2960 id, HudElement element) {
      validateUnique(id);
      LAST.layers().addLast(HudLayer.ofElement(id, element));
   }

   public static void attachElementBefore(class_2960 beforeThis, class_2960 id, HudElement element) {
      validateUnique(id);
      boolean didChange = findLayer(beforeThis, (l, iterator) -> {
         iterator.previous();
         iterator.add(HudLayer.ofElement(id, element));
         iterator.next();
         return true;
      });
      if (!didChange) {
         throw new IllegalArgumentException("Layer with identifier " + String.valueOf(beforeThis) + " not found");
      }
   }

   public static void attachElementAfter(class_2960 afterThis, class_2960 id, HudElement element) {
      validateUnique(id);
      boolean didChange = findLayer(afterThis, (l, iterator) -> {
         iterator.add(HudLayer.ofElement(id, element));
         return true;
      });
      if (!didChange) {
         throw new IllegalArgumentException("Layer with identifier " + String.valueOf(afterThis) + " not found");
      }
   }

   public static void removeElement(class_2960 identifier) {
      boolean didChange = findLayer(identifier, (l, iterator) -> {
         class_2960 var10001 = l.id();
         Objects.requireNonNull(l);
         iterator.set(HudLayer.of(var10001, l::element, true));
         return true;
      });
      if (!didChange) {
         throw new IllegalArgumentException("Layer with identifier " + String.valueOf(identifier) + " not found");
      }
   }

   public static void replaceElement(class_2960 identifier, Function<HudElement, HudElement> replacer) {
      boolean didChange = findLayer(identifier, (l, iterator) -> {
         class_2960 var10001 = l.id();
         Objects.requireNonNull(l);
         iterator.set(HudLayer.of(var10001, replacer.compose(l::element), l.isRemoved()));
         return true;
      });
      if (!didChange) {
         throw new IllegalArgumentException("Layer with identifier " + String.valueOf(identifier) + " not found");
      }
   }

   @VisibleForTesting
   static void validateUnique(class_2960 id) {
      visitLayers((l, iterator) -> {
         if (l.id().equals(id)) {
            throw new IllegalArgumentException("Layer with identifier " + String.valueOf(id) + " already exists");
         } else {
            return false;
         }
      });
   }

   @VisibleForTesting
   static boolean findLayer(class_2960 identifier, LayerVisitor visitor) {
      MutableBoolean found = new MutableBoolean(false);
      visitLayers((l, iterator) -> {
         if (l.id().equals(identifier)) {
            found.setTrue();
            return visitor.visit(l, iterator);
         } else {
            return false;
         }
      });
      return found.booleanValue();
   }

   @VisibleForTesting
   static boolean visitLayers(LayerVisitor visitor) {
      boolean modified = false;

      for(class_2960 id : VANILLA_ELEMENT_IDS) {
         RootLayer rootLayer = (RootLayer)ROOT_ELEMENTS.get(id);
         modified |= visitLayers(rootLayer.layers(), visitor);
      }

      return modified;
   }

   private static boolean visitLayers(List<HudLayer> layers, LayerVisitor visitor) {
      MutableBoolean modified = new MutableBoolean(false);
      ListIterator<HudLayer> iterator = layers.listIterator();

      while(iterator.hasNext()) {
         HudLayer layer = (HudLayer)iterator.next();
         if (visitor.visit(layer, iterator)) {
            modified.setTrue();
         }
      }

      return modified.booleanValue();
   }

   static {
      VANILLA_ELEMENT_IDS = List.of(VanillaHudElements.MISC_OVERLAYS, VanillaHudElements.CROSSHAIR, VanillaHudElements.SPECTATOR_MENU, VanillaHudElements.HOTBAR, VanillaHudElements.ARMOR_BAR, VanillaHudElements.HEALTH_BAR, VanillaHudElements.FOOD_BAR, VanillaHudElements.AIR_BAR, VanillaHudElements.MOUNT_HEALTH, VanillaHudElements.INFO_BAR, VanillaHudElements.EXPERIENCE_LEVEL, VanillaHudElements.HELD_ITEM_TOOLTIP, VanillaHudElements.SPECTATOR_TOOLTIP, VanillaHudElements.STATUS_EFFECTS, VanillaHudElements.BOSS_BAR, VanillaHudElements.SLEEP, VanillaHudElements.DEMO_TIMER, VanillaHudElements.SCOREBOARD, VanillaHudElements.OVERLAY_MESSAGE, VanillaHudElements.TITLE_AND_SUBTITLE, VanillaHudElements.CHAT, VanillaHudElements.PLAYER_LIST, VanillaHudElements.SUBTITLES);
      ROOT_ELEMENTS = (Map)VANILLA_ELEMENT_IDS.stream().map(RootLayer::new).collect(Collectors.toMap(RootLayer::id, Function.identity(), (a, b) -> a, IdentityHashMap::new));
      FIRST = (RootLayer)ROOT_ELEMENTS.get(VanillaHudElements.MISC_OVERLAYS);
      LAST = (RootLayer)ROOT_ELEMENTS.get(VanillaHudElements.SUBTITLES);
   }

   @Environment(EnvType.CLIENT)
   public static record RootLayer(class_2960 id, List<HudLayer> layers) {
      private RootLayer(class_2960 id) {
         this(id, new ArrayList());
         this.layers().add(HudLayer.ofVanilla(id));
      }

      public void render(class_332 context, class_9779 tickCounter, HudElement vanillaElement) {
         for(HudLayer layer : this.layers) {
            if (!layer.isRemoved()) {
               layer.element(vanillaElement).render(context, tickCounter);
            }
         }

      }
   }

   @Environment(EnvType.CLIENT)
   @VisibleForTesting
   interface LayerVisitor {
      boolean visit(HudLayer var1, ListIterator<HudLayer> var2);
   }
}