TiedUp-/src/main/java/com/tiedup/remake/rig/anim/AnimationManager.java

/*
 * Derived from Epic Fight (https://github.com/Epic-Fight/epicfight)
 * by the Epic Fight Team, licensed under GPLv3.
 * Modifications © 2026 TiedUp! Remake Contributors, distributed under GPLv3.
 */

package com.tiedup.remake.rig.anim;

import java.io.IOException;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.logging.log4j.Logger;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonParseException;

import net.minecraft.client.Minecraft;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.chat.Component;
import net.minecraft.resources.FileToIdConverter;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.network.ServerGamePacketListenerImpl;
import net.minecraft.server.packs.resources.Resource;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.server.packs.resources.SimplePreparableReloadListener;
import net.minecraft.util.GsonHelper;
import net.minecraft.util.profiling.ProfilerFiller;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import net.minecraftforge.eventbus.api.Event;
import net.minecraftforge.fml.event.IModBusEvent;
import com.tiedup.remake.rig.anim.property.AnimationProperty;
import com.tiedup.remake.rig.anim.types.DynamicAnimation;
import com.tiedup.remake.rig.anim.types.StaticAnimation;
import com.tiedup.remake.rig.asset.AssetAccessor;
import com.tiedup.remake.rig.asset.JsonAssetLoader;
import com.tiedup.remake.rig.anim.client.AnimationSubFileReader;
import yesman.epicfight.api.data.reloader.SkillManager;
import com.tiedup.remake.rig.exception.AssetLoadingException;
import com.tiedup.remake.rig.util.InstantiateInvoker;
import com.tiedup.remake.rig.util.MutableBoolean;
import com.tiedup.remake.rig.TiedUpRigRegistry;
import com.tiedup.remake.rig.TiedUpRigConstants;
import yesman.epicfight.network.EpicFightNetworkManager;
import yesman.epicfight.network.client.CPCheckAnimationRegistryMatches;
import yesman.epicfight.network.server.SPDatapackSync;

@SuppressWarnings("unchecked")
public class AnimationManager extends SimplePreparableReloadListener<List<ResourceLocation>> {
	private static final AnimationManager INSTANCE = new AnimationManager();
	private static ResourceManager serverResourceManager = null;
	private static final Gson GSON = new GsonBuilder().create();
    private static final String DIRECTORY = "animmodels/animations";
    
	public static AnimationManager getInstance() {
		return INSTANCE;
	}
	
	private final Map<Integer, AnimationAccessor<? extends StaticAnimation>> animationById = Maps.newHashMap();
	private final Map<ResourceLocation, AnimationAccessor<? extends StaticAnimation>> animationByName = Maps.newHashMap();
	private final Map<AnimationAccessor<? extends StaticAnimation>, StaticAnimation> animations = Maps.newHashMap();
	private final Map<AnimationAccessor<? extends StaticAnimation>, String> resourcepackAnimationCommands = Maps.newHashMap();
	
	public static boolean checkNull(AssetAccessor<? extends StaticAnimation> animation) {
		if (animation == null || animation.isEmpty()) {
			if (animation != null) {
				TiedUpRigConstants.stacktraceIfDevSide("Empty animation accessor: " + animation.registryName(), NoSuchElementException::new);
			} else {
				TiedUpRigConstants.stacktraceIfDevSide("Null animation accessor", NoSuchElementException::new);
			}
			
			return true;
		}
		
		return false;
	}
	
	public static <T extends StaticAnimation> AnimationAccessor<T> byKey(String registryName) {
		return byKey(ResourceLocation.parse(registryName));
	}
	
	public static <T extends StaticAnimation> AnimationAccessor<T> byKey(ResourceLocation registryName) {
		return (AnimationAccessor<T>)getInstance().animationByName.get(registryName);
	}
	
	public static <T extends StaticAnimation> AnimationAccessor<T> byId(int animationId) {
		return (AnimationAccessor<T>)getInstance().animationById.get(animationId);
	}
	
	public Map<ResourceLocation, AnimationAccessor<? extends StaticAnimation>> getAnimations(Predicate<AssetAccessor<? extends StaticAnimation>> filter) {
		Map<ResourceLocation, AnimationAccessor<? extends StaticAnimation>> filteredItems =
			this.animationByName.entrySet().stream()
				.filter(entry -> {
					return filter.test(entry.getValue());
				})
				.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
		
		return ImmutableMap.copyOf(filteredItems);
	}
	
	public AnimationClip loadAnimationClip(StaticAnimation animation, BiFunction<JsonAssetLoader, StaticAnimation, AnimationClip> clipLoader) {
		try {
			if (getAnimationResourceManager() == null) {
				return null;
			}
			
			JsonAssetLoader modelLoader = new JsonAssetLoader(getAnimationResourceManager(), animation.getLocation());
			AnimationClip loadedClip = clipLoader.apply(modelLoader, animation);
			
			return loadedClip;
		} catch (AssetLoadingException e) {
			throw new AssetLoadingException("Failed to load animation clip from: " + animation, e);
		}
	}
	
	public static void readAnimationProperties(StaticAnimation animation) {
		ResourceLocation dataLocation = getSubAnimationFileLocation(animation.getLocation(), AnimationSubFileReader.SUBFILE_CLIENT_PROPERTY);
		ResourceLocation povLocation = getSubAnimationFileLocation(animation.getLocation(), AnimationSubFileReader.SUBFILE_POV_ANIMATION);
		
		getAnimationResourceManager().getResource(dataLocation).ifPresent((rs) -> {
			AnimationSubFileReader.readAndApply(animation, rs, AnimationSubFileReader.SUBFILE_CLIENT_PROPERTY);
		});
		
		getAnimationResourceManager().getResource(povLocation).ifPresent((rs) -> {
			AnimationSubFileReader.readAndApply(animation, rs, AnimationSubFileReader.SUBFILE_POV_ANIMATION);
		});
	}
	
	@Override
	protected List<ResourceLocation> prepare(ResourceManager resourceManager, ProfilerFiller profilerIn) {
		if (!TiedUpRigConstants.isPhysicalClient() && serverResourceManager == null) {
			serverResourceManager = resourceManager;
		}
		
		this.animations.clear();
		this.animationById.entrySet().removeIf(entry -> !entry.getValue().inRegistry());
		this.animationByName.entrySet().removeIf(entry -> !entry.getValue().inRegistry());
		this.resourcepackAnimationCommands.clear();
		
        List<ResourceLocation> directories = new ArrayList<> ();
        scanDirectoryNames(resourceManager, directories);
        
		return directories;
	}
	
    private static void scanDirectoryNames(ResourceManager resourceManager, List<ResourceLocation> output) {
        FileToIdConverter filetoidconverter = FileToIdConverter.json(DIRECTORY);
        filetoidconverter.listMatchingResources(resourceManager).keySet().stream().map(AnimationManager::pathToId).forEach(output::add);
    }
	
    @Override
	protected void apply(List<ResourceLocation> objects, ResourceManager resourceManager, ProfilerFiller profilerIn) {
		// RIG : Armatures.reload() (EF gameasset registry) retiré.
		// TiedUpArmatures.reload() sera appelé ici en Phase 2 quand le registry
		// sera créé. En Phase 0, no-op.
		
		Set<ResourceLocation> registeredAnimation =
            this.animationById.values().stream()
                .reduce(
                    new HashSet<> (),
                    (set, accessor) -> {
                        set.add(accessor.registryName());

                        for (AssetAccessor<? extends StaticAnimation> subAnimAccessor : accessor.get().getSubAnimations()) {
                            set.add(subAnimAccessor.registryName());
                        }

                        return set;
                    },
                    (set1, set2) -> {
                        set1.addAll(set2);
                        return set1;
                    }
                );
		
		// Load animations that are not registered by AnimationRegistryEvent
        // Reads from /assets folder in physical client, /datapack in physical server.
        objects.stream()
            .filter(animId -> !registeredAnimation.contains(animId) && !animId.getPath().contains("/data/") && !animId.getPath().contains("/pov/"))
			.sorted(Comparator.comparing(ResourceLocation::toString))
			.forEach(animId -> {
                Optional<Resource> resource = resourceManager.getResource(idToPath(animId));
                
                try (Reader reader = resource.orElseThrow().openAsReader()) {
                    JsonElement jsonelement = GsonHelper.fromJson(GSON, reader, JsonElement.class);
                    this.readResourcepackAnimation(animId, jsonelement.getAsJsonObject());
                } catch (IOException | JsonParseException | IllegalArgumentException resourceReadException) {
                    TiedUpRigConstants.LOGGER.error("Couldn't parse animation data from {}", animId, resourceReadException);
                } catch (Exception e) {
                    TiedUpRigConstants.LOGGER.error("Failed at constructing {}", animId, e);
                }
            });
        
        SkillManager.reloadAllSkillsAnimations();
        
		this.animations.entrySet().stream()
            .reduce(
                new ArrayList<AssetAccessor<? extends StaticAnimation>>(),
                (list, entry) -> {
                    MutableBoolean init = new MutableBoolean(true);
                    
                    if (entry.getValue() == null || entry.getValue().getAccessor() == null) {
                        TiedUpRigConstants.logAndStacktraceIfDevSide(Logger::error, "Invalid animation implementation: " + entry.getKey(), AssetLoadingException::new);
                        init.set(false);
                    }
                    
                    entry.getValue().getSubAnimations().forEach((subAnimation) -> {
                        if (subAnimation == null || subAnimation.get() == null) {
                            TiedUpRigConstants.logAndStacktraceIfDevSide(Logger::error, "Invalid sub animation implementation: " + entry.getKey(), AssetLoadingException::new);
                            init.set(false);
                        }
                    });
                    
                    if (init.value()) {
                        list.add(entry.getValue().getAccessor());
                        list.addAll(entry.getValue().getSubAnimations());
                    }
                    
                    return list;
                },
                (list1, list2) -> {
                    list1.addAll(list2);
                    return list1;
                }
            )
            .forEach(accessor -> {
                accessor.doOrThrow(StaticAnimation::postInit);

                if (TiedUpRigConstants.isPhysicalClient()) {
                    AnimationManager.readAnimationProperties(accessor.get());
                }
            });
	}
	
	public static ResourceLocation getSubAnimationFileLocation(ResourceLocation location, AnimationSubFileReader.SubFileType<?> subFileType) {
		int splitIdx = location.getPath().lastIndexOf('/');
		
		if (splitIdx < 0) {
			splitIdx = 0;
		}
		
		return ResourceLocation.fromNamespaceAndPath(location.getNamespace(), String.format("%s/" + subFileType.getDirectory() + "%s", location.getPath().substring(0, splitIdx), location.getPath().substring(splitIdx)));
	}
	
	/// Converts animation id, acquired by [StaticAnimation#getRegistryName], to animation resource path acquired by [StaticAnimation#getLocation]
    public static ResourceLocation idToPath(ResourceLocation rl) {
        return rl.getPath().matches(DIRECTORY + "/.*\\.json") ? rl : ResourceLocation.fromNamespaceAndPath(rl.getNamespace(), DIRECTORY + "/" + rl.getPath() + ".json");
    }

    /// Converts animation resource path, acquired by [StaticAnimation#getLocation], to animation id acquired by [StaticAnimation#getRegistryName]
    public static ResourceLocation pathToId(ResourceLocation rl) {
        return ResourceLocation.fromNamespaceAndPath(rl.getNamespace(), rl.getPath().replace(DIRECTORY + "/", "").replace(".json", ""));
    }
	
	public static void setServerResourceManager(ResourceManager pResourceManager) {
		serverResourceManager = pResourceManager;
	}
	
	public static ResourceManager getAnimationResourceManager() {
		return TiedUpRigConstants.isPhysicalClient() ? Minecraft.getInstance().getResourceManager() : serverResourceManager;
	}
	
	public int getResourcepackAnimationCount() {
		return this.resourcepackAnimationCommands.size();
	}
	
	public Stream<CompoundTag> getResourcepackAnimationStream() {
		return this.resourcepackAnimationCommands.entrySet().stream().map((entry) -> {
			CompoundTag compTag = new CompoundTag();
			compTag.putString("registry_name", entry.getKey().registryName().toString());
			compTag.putInt("id", entry.getKey().id());
			compTag.putString("invoke_command", entry.getValue());
			
			return compTag;
		});
	}
	
	/**
	 * @param mandatoryPack : creates dummy animations for animations from the server without animation clips when the server has mandatory resource pack.
	 *                        custom weapon types & mob capabilities won't be created because they won't be able to find the animations from the server
	 *                        dummy animations will be automatically removed right after reloading resourced as the server forces using resource pack
	 */
	@OnlyIn(Dist.CLIENT)
	public void processServerPacket(SPDatapackSync packet, boolean mandatoryPack) {
		if (mandatoryPack) {
			for (CompoundTag tag : packet.getTags()) {
				String invocationCommand = tag.getString("invoke_command");
				ResourceLocation registryName = ResourceLocation.parse(tag.getString("registry_name"));
				int id = tag.getInt("id");
						
				if (this.animationByName.containsKey(registryName)) {
					continue;
				}
				
				AnimationAccessor<? extends StaticAnimation> accessor = AnimationAccessorImpl.create(registryName, getResourcepackAnimationCount(), false, (accessor$2) -> {
					try {
						return InstantiateInvoker.invoke(invocationCommand, StaticAnimation.class).getResult();
					} catch (Exception e) {
						TiedUpRigConstants.LOGGER.warn("Failed at creating animation from server resource pack");
						e.printStackTrace();
						return TiedUpRigRegistry.EMPTY_ANIMATION;
					}
				});
				
				this.animationById.put(id, accessor);
				this.animationByName.put(registryName, accessor);
			}
		}
		
		int animationCount = this.animations.size();
		String[] registryNames = new String[animationCount];
		
		for (int i = 0; i < animationCount; i++) {
			String registryName = this.animationById.get(i + 1).registryName().toString();
			registryNames[i] = registryName;
		}
		
		CPCheckAnimationRegistryMatches registrySyncPacket = new CPCheckAnimationRegistryMatches(animationCount, registryNames);
		EpicFightNetworkManager.sendToServer(registrySyncPacket);
	}
	
	public void validateClientAnimationRegistry(CPCheckAnimationRegistryMatches msg, ServerGamePacketListenerImpl connection) {
		StringBuilder messageBuilder = new StringBuilder();
		int count = 0;
		
		Set<String> clientAnimationRegistry = new HashSet<> (Set.of(msg.registryNames));
		
		for (String registryName : this.animations.keySet().stream().map((rl) -> rl.toString()).toList()) {
			if (!clientAnimationRegistry.contains(registryName)) {
				// Animations that don't exist in client
				if (count < 10) {
					messageBuilder.append(registryName);
					messageBuilder.append("\n");
				}
				
				count++;
			} else {
				clientAnimationRegistry.remove(registryName);
			}
		}
		
		// Animations that don't exist in server
		for (String registryName : clientAnimationRegistry) {
			if (registryName.equals("empty")) {
				continue;
			}
			
			if (count < 10) {
				messageBuilder.append(registryName);
				messageBuilder.append("\n");
			}
			
			count++;
		}
		
		if (count >= 10) {
			messageBuilder.append(Component.translatable("gui.epicfight.warn.animation_unsync.etc", (count - 9)).getString());
			messageBuilder.append("\n");
		}
		
		if (!messageBuilder.isEmpty()) {
			connection.disconnect(Component.translatable("gui.epicfight.warn.animation_unsync", messageBuilder.toString()));
		}
	}
	
	private static final Set<String> NO_WARNING_MODID = Sets.newHashSet();
	
	public static void addNoWarningModId(String modid) {
		NO_WARNING_MODID.add(modid);
	}
	
	/**************************************************
	 * User-animation loader
	 **************************************************/
	@SuppressWarnings({ "deprecation" })
	private void readResourcepackAnimation(ResourceLocation rl, JsonObject json) throws Exception {
		JsonElement constructorElement = json.get("constructor");
		
		if (constructorElement == null) {
			if (NO_WARNING_MODID.contains(rl.getNamespace())) {
				return;
			} else {
				TiedUpRigConstants.logAndStacktraceIfDevSide(
					  Logger::error
					, "Datapack animation reading failed: No constructor information has provided: " + rl
					, IllegalStateException::new
					, "No constructor information has provided in User animation, " + rl + "\nPlease remove this resource if it's unnecessary to optimize your project."
				);
				return;
			}
		}
		
		JsonObject constructorObject = constructorElement.getAsJsonObject();
		String invocationCommand = constructorObject.get("invocation_command").getAsString();
		StaticAnimation animation = InstantiateInvoker.invoke(invocationCommand, StaticAnimation.class).getResult();
		JsonElement propertiesElement = json.getAsJsonObject().get("properties");
		
		if (propertiesElement != null) {
			JsonObject propertiesObject = propertiesElement.getAsJsonObject();
			
			for (Map.Entry<String, JsonElement> entry : propertiesObject.entrySet()) {
				AnimationProperty<?> propertyKey = AnimationProperty.getSerializableProperty(entry.getKey());
				Object value = propertyKey.parseFrom(entry.getValue());
				animation.addPropertyUnsafe(propertyKey, value);
			}
		}
		
		AnimationAccessor<StaticAnimation> accessor = AnimationAccessorImpl.create(rl, this.animations.size() + 1, false, null);
		animation.setAccessor(accessor);
		
		this.resourcepackAnimationCommands.put(accessor, invocationCommand);
		this.animationById.put(accessor.id(), accessor);
		this.animationByName.put(accessor.registryName(), accessor);
		this.animations.put(accessor, animation);
	}



	public interface AnimationAccessor<A extends DynamicAnimation> extends AssetAccessor<A> {
		int id();
		
		default boolean idBetween(AnimationAccessor<? extends StaticAnimation> a1, AnimationAccessor<? extends StaticAnimation> a2) {
			return a1.id() <= this.id() && a2.id() >= this.id();
		}
	}
	
	public static record AnimationAccessorImpl<A extends StaticAnimation> (ResourceLocation registryName, int id, boolean inRegistry, Function<AnimationAccessor<A>, A> onLoad) implements AnimationAccessor<A> {
		private static <A extends StaticAnimation> AnimationAccessor<A> create(ResourceLocation registryName, int id, boolean inRegistry, Function<AnimationAccessor<A>, A> onLoad) {
			return new AnimationAccessorImpl<A> (registryName, id, inRegistry, onLoad);
		}
		
		@Override
		public A get() {
			if (!INSTANCE.animations.containsKey(this)) {
				INSTANCE.animations.put(this, this.onLoad.apply(this));
			}
			
			return (A)INSTANCE.animations.get(this);
		}
		
		public String toString() {
			return this.registryName.toString();
		}
		
		public int hashCode() {
			return this.registryName.hashCode();
		}
		
		public boolean equals(Object obj) {
			if (this == obj) {
				return true;
			} else if (obj instanceof AnimationAccessor armatureAccessor) {
				return this.registryName.equals(armatureAccessor.registryName());
			} else if (obj instanceof ResourceLocation rl) {
				return this.registryName.equals(rl);
			} else if (obj instanceof String name) {
				return this.registryName.toString().equals(name);
			} else {
				return false;
			}
		}
	}
	
	public static class AnimationRegistryEvent extends Event implements IModBusEvent {
		private List<AnimationBuilder> builders = Lists.newArrayList();
		private Set<String> namespaces = Sets.newHashSet();
		
		public void newBuilder(String namespace, Consumer<AnimationBuilder> build) {
			if (this.namespaces.contains(namespace)) {
				throw new IllegalArgumentException("Animation builder namespace '" + namespace + "' already exists!");
			}
			
			this.namespaces.add(namespace);
			this.builders.add(new AnimationBuilder(namespace, build));
		}
		
		public List<AnimationBuilder> getBuilders() {
			return this.builders;
		}
	}
	
	public static record AnimationBuilder(String namespace, Consumer<AnimationBuilder> task) {
		public <T extends StaticAnimation> AnimationManager.AnimationAccessor<T> nextAccessor(String id, Function<AnimationManager.AnimationAccessor<T>, T> onLoad) {
			AnimationAccessor<T> accessor = AnimationAccessorImpl.create(ResourceLocation.fromNamespaceAndPath(this.namespace, id), INSTANCE.animations.size() + 1, true, onLoad);
			
			INSTANCE.animationById.put(accessor.id(), accessor);
			INSTANCE.animationByName.put(accessor.registryName(), accessor);
			INSTANCE.animations.put(accessor, null);
			
			return accessor; 
		}
	}
}