SharedVRCStuff/AAC/AACCore/Editor/AACCore.cs

using System.Linq;
using System.Diagnostics;
using AnimatorAsCode.V1;
using gay.lilyy.aaccore;
using nadena.dev.ndmf;
using AnimatorAsCode.V1.ModularAvatar;
using UnityEditor;
using UnityEngine;
using VRC.SDK3.Avatars.Components;
using VRC.SDK3.Dynamics.Constraint.Components;
using System.Collections.Generic;
using gay.lilyy.EditorNotes;

[assembly: ExportsPlugin(typeof(AACPlugin))]

namespace gay.lilyy.aaccore
{



    public class AACAssets
    {
        public AacFlBase aac;
        public AacFlController fx;
        public AacFlClip emptyClip;
        public BuildContext ctx;
        public AACRoot root;
        public MaAc modularAvatar;
        
        public bool experimentalEnabled;
        public bool isPC;
    }

    public static class AACUtils
    {
        public static GameObject FindChildRecursive(GameObject parent, string name)
        {
            return FindChildRecursive(parent.transform, name)?.gameObject;
        }
        public static Transform FindChildRecursive(Transform parent, string name)
        {
            Transform childTransform = parent.GetComponentsInChildren<Transform>()
                                            .FirstOrDefault(t => t.gameObject.name == name);
            return childTransform;
        }

        public static AacFlClip CreateConstraintWeightClip(AACAssets assets, VRCParentConstraint constraint, int index, int indexCount) {
            return assets.aac.NewClip()
                .Animating(action => {
                    action.Animates(constraint, $"Sources.source{index}.Weight").WithOneFrame(1f);
                    for (int i = 0; i < indexCount; i++) {
                        if (i == index) continue;
                        action.Animates(constraint, $"Sources.source{i}.Weight").WithOneFrame(0f);
                    }
                });
        }

        public static AacFlClip CreateEmptyClipWithFrames(AACAssets assets, int frames) {
            if (frames == 0) return assets.emptyClip;
            var emptyGo = new GameObject();
            emptyGo.name = "_EmptyClipInstant";
            emptyGo.transform.SetParent(assets.ctx.AvatarRootTransform);
            var clip = assets.aac.NewClip().Animating(action => {
                action.Animates(emptyGo.transform, "m_LocalScale.z").WithFrameCountUnit(unit => {
                    unit.Constant(0, 0);
                    unit.Constant(frames, 1);
                });
            });
            Object.DestroyImmediate(emptyGo);
            return clip;
        }

        public static AacFlClip CreateEmptyClipWithSeconds(AACAssets assets, float seconds) {
            if (seconds == 0) return assets.emptyClip;
            var emptyGo = new GameObject();
            emptyGo.name = "_EmptyClipInstant";
            emptyGo.transform.SetParent(assets.ctx.AvatarRootTransform);
            var clip = assets.aac.NewClip().Animating(action => {
                action.Animates(emptyGo.transform, "m_LocalScale.z").WithFrameCountUnit(unit => {
                    unit.Constant(0, 0);
                    unit.Constant(seconds, 1);
                });
            });
            Object.DestroyImmediate(emptyGo);
            return clip;
        }

        public static List<SkinnedMeshRenderer> FindWithBlendshape(AACAssets assets, string shapeName) {
            var list = new List<SkinnedMeshRenderer>();
            foreach (var smr in assets.ctx.AvatarRootObject.GetComponentsInChildren<SkinnedMeshRenderer>(true)) {
                var mesh = smr.sharedMesh;
                if (!mesh) continue;

                int count = mesh.blendShapeCount;
                for (int i = 0; i < count; i++) {
                    if (mesh.GetBlendShapeName(i) == shapeName) {
                        list.Add(smr);
                        break;
                    }
                }
            }
            return list;
        }

        public static List<T> FindNamedComponents<T>(Transform parent, string name) where T : Component {
            var list = new List<T>();
            foreach (var component in parent.GetComponentsInChildren<T>(true)) {
                if (component.gameObject.name == name) {
                    list.Add(component);
                }
            }
            UnityEngine.Debug.Log($"Found {list.Count} {name} components");
            foreach (var component in list) {
                UnityEngine.Debug.Log($"Component: {component.gameObject.name}");
            }
            return list;
        }

        public static List<T> FindNamedComponents<T>(AACAssets assets, string name) where T : Component {
            return FindNamedComponents<T>(assets.ctx.AvatarRootTransform, name);
        }

        public static List<Transform> FindTransformWithNote(Transform parent, string tag) {
            var notes = parent.GetComponentsInChildren<EditorNote>(true);
            var matched = new List<Transform>();
            foreach (var note in notes) {
                if (note.valueType != EditorNote.ValueType.String) continue;
                if (note.stringValue != "aactag") continue;
                if (note.note == tag) {
                    matched.Add(note.transform);
                }
            }
            return matched;
        }
    }


    public class AACPlugin : Plugin<AACPlugin>
    {
        public override string QualifiedName => "gay.lilyy.aaccore.plugin";
        public override string DisplayName => "AACCore";

        private const string SystemName = "AACCore";
        private const bool UseWriteDefaults = true;
        protected override void Configure()
        {
            InPhase(BuildPhase.PlatformInit).BeforePlugin("nadena.dev.modular-avatar").Run($"Platform Init {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.PlatformInit));
            InPhase(BuildPhase.Resolving).BeforePlugin("nadena.dev.modular-avatar").Run($"Resolving {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.Resolving));
            InPhase(BuildPhase.Generating).BeforePlugin("nadena.dev.modular-avatar").Run($"Generating {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.Generating));
            InPhase(BuildPhase.Transforming).BeforePlugin("nadena.dev.modular-avatar").Run($"Transforming {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.Transforming));
            InPhase(BuildPhase.Optimizing).BeforePlugin("nadena.dev.modular-avatar").Run($"Optimizing {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.Optimizing));
            InPhase(BuildPhase.PlatformFinish).BeforePlugin("nadena.dev.modular-avatar").Run($"Platform Finish {DisplayName}", (ctx) => RunPhase(ctx, BuildPhase.PlatformFinish));
        }

        private void RunPhase(BuildContext ctx, BuildPhase phase)
        {

            var allLayerGroups = LayerGroup.Instances.ToList();
            var layerGroups = allLayerGroups
                .Where(lg => lg.SystemName != "template" && lg.buildPhase == phase)
                .ToList();
            
            // Start overall timer
            var overallStopwatch = Stopwatch.StartNew();

            V5AACLogger.LogInfo("Starting Lillith V5 AAC generation...");
            var root = ctx.AvatarRootObject.GetComponent<AACRoot>();
            if (root == null)
            {
                V5AACLogger.LogInfo("No LillithV5AACRoot component found. Skipping.");
                return;
            }

            // Skip if the root GameObject or the component itself is disabled
            if (!root.enabled || !root.gameObject.activeInHierarchy)
            {
                V5AACLogger.LogInfo($"LillithV5AACRoot on GameObject '{root.name}' is disabled or inactive. Skipping.");
                return;
            }

            AACAssets assets = new();
            assets.ctx = ctx;
            assets.root = root;
            assets.modularAvatar = MaAc.Create(new GameObject(SystemName)
            {
                transform = { parent = ctx.AvatarRootTransform }
            });

            assets.isPC = EditorUserBuildSettings.activeBuildTarget == BuildTarget.StandaloneWindows
                        || EditorUserBuildSettings.activeBuildTarget == BuildTarget.StandaloneWindows64;


            // Time AAC initialization
            var initStopwatch = Stopwatch.StartNew();
            V5AACLogger.LogInfo("Initializing Animator As Code...");
            assets.aac = AacV1.Create(new AacConfiguration
            {
                SystemName = SystemName,
                AnimatorRoot = ctx.AvatarRootTransform,
                DefaultValueRoot = ctx.AvatarRootTransform,
                AssetKey = GUID.Generate().ToString(),
                AssetContainer = ctx.AssetContainer,
                ContainerMode = AacConfiguration.Container.OnlyWhenPersistenceRequired,
                AssetContainerProvider = new NDMFContainerProvider(ctx),
                DefaultsProvider = new AacDefaultsProvider(UseWriteDefaults)
            });
            assets.emptyClip = assets.aac.NewClip();
            
            assets.fx = assets.aac.NewAnimatorController();
            initStopwatch.Stop();
            V5AACLogger.LogInfo($"AAC initialization completed in {initStopwatch.ElapsedMilliseconds}ms");

            // Process layer groups with individual timing
            // Filter out the template LayerGroup (by SystemName)
            V5AACLogger.LogInfo($"Processing {layerGroups.Count} layer groups...");

            var totalLayerGroupTime = 0L;
            var processedCount = 0;
            var skippedCount = 0;

            if (root.experimentalPlayMode && Application.isPlaying)
            {
                assets.experimentalEnabled = true;
            }
            if (root.experimentalUpload && !Application.isPlaying)
            {
                assets.experimentalEnabled = true;
            }

            foreach (var layerGroup in layerGroups)
            {
                var layerStopwatch = Stopwatch.StartNew();

                bool skipped = false;

                if (!layerGroup.enabled)
                {
                    V5AACLogger.LogWarning($"Skipping layer group: {layerGroup.DisplayName} (disabled)");
                    skipped = true;
                }
                else if (layerGroup.experimental && !assets.experimentalEnabled)
                {
                    V5AACLogger.LogWarning($"Skipping layer group: {layerGroup.DisplayName} (Experimental)");
                    skipped = true;
                }
                else if (!layerGroup.IsApplicable(assets))
                {
                    if (layerGroup.shouldWarnIfNotApplicable)
                        V5AACLogger.LogWarning($"Skipping layer group: {layerGroup.DisplayName} (Not applicable)");
                    
                    skipped = true;
                }

                if (!skipped)
                {
                    V5AACLogger.LogInfo($"Running layer group: {layerGroup.DisplayName}");
                    layerGroup.Run(assets);
                    processedCount++;
                }
                else
                {
                    skippedCount++;
                }

                layerStopwatch.Stop();
                totalLayerGroupTime += layerStopwatch.ElapsedMilliseconds;
                V5AACLogger.LogInfo($"Layer group '{layerGroup.DisplayName}' completed in {layerStopwatch.ElapsedMilliseconds}ms");
            }

            // Time Modular Avatar setup
            var maStopwatch = Stopwatch.StartNew();
            V5AACLogger.LogInfo("Creating Modular Avatar merge animator...");

            assets.modularAvatar.NewMergeAnimator(assets.fx.AnimatorController, VRCAvatarDescriptor.AnimLayerType.FX);
            maStopwatch.Stop();
            V5AACLogger.LogInfo($"Modular Avatar setup completed in {maStopwatch.ElapsedMilliseconds}ms");

            // Final timing summary
            overallStopwatch.Stop();
            V5AACLogger.LogSuccess($"Lillith V5 AAC generation completed successfully!");
            V5AACLogger.LogInfo($"=== TIMING SUMMARY ===");
            V5AACLogger.LogInfo($"Total time: {overallStopwatch.ElapsedMilliseconds}ms ({overallStopwatch.Elapsed.TotalSeconds:F2}s)");
            V5AACLogger.LogInfo($"AAC initialization: {initStopwatch.ElapsedMilliseconds}ms");
            V5AACLogger.LogInfo($"Layer groups: {totalLayerGroupTime}ms (processed: {processedCount}, skipped: {skippedCount})");
            V5AACLogger.LogInfo($"Modular Avatar: {maStopwatch.ElapsedMilliseconds}ms");
            V5AACLogger.LogInfo($"Average per layer group: {(processedCount > 0 ? totalLayerGroupTime / processedCount : 0)}ms");
        }
    }



    internal class NDMFContainerProvider : IAacAssetContainerProvider
    {
        private readonly BuildContext _ctx;
        public NDMFContainerProvider(BuildContext ctx) => _ctx = ctx;
        public void SaveAsPersistenceRequired(Object objectToAdd) => _ctx.AssetSaver.SaveAsset(objectToAdd);
        public void SaveAsRegular(Object objectToAdd) { } // Let NDMF crawl our assets when it finishes
        public void ClearPreviousAssets() { } // ClearPreviousAssets is never used in non-destructive contexts
    }
}