terra-llm/Assets/AssetBank/Editor/Tools/JsonDataTrimmer.cs

using System.Collections.Generic;
using System.Linq;
using AssetBank.Editor.SchemaConverter.Data.Input;
using AssetBank.Settings;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;
using UnityEngine;

namespace AssetBank.Editor.Tools
{
    public static class JsonDataTrimmer
    {
        public static string Process(string rawJsonText, ProjectExporterSettings settings)
        {
            var processedJson = ShortenBrokenPrefabGameObjects(rawJsonText);
            
            var allComponents = JArray.Parse(processedJson);
            
            // Pass 1: Build the definitive ruleset
            var rules = settings.SafeComponents.ToDictionary(r => r.Component, r => r.ExposedFields.ToHashSet());
            if (!settings.OverrideHardcodedDefaults)
            {
                foreach (var hardcodedRule in settings.HardcodedComponents)
                {
                    if (rules.TryGetValue(hardcodedRule.Component, out var userRules))
                    {
                        userRules.UnionWith(hardcodedRule.ExposedFields);
                    }
                    else
                    {
                        rules[hardcodedRule.Component] = hardcodedRule.ExposedFields.ToHashSet();
                    }
                }
            }

            // Pass 2: Create a fully trimmed copy of every component
            var trimmedComponentMap = new Dictionary<string, JObject>();
            foreach (var component in allComponents.Children<JObject>())
            {
                var anchorId = component["anchor_id"]?.ToString();
                if (string.IsNullOrEmpty(anchorId)) continue;

                var trimmedComponent = component.DeepClone() as JObject;
                var dataBlock = trimmedComponent["data"] as JObject;
                if (dataBlock == null || !dataBlock.HasValues)
                {
                    trimmedComponentMap[anchorId] = trimmedComponent;
                    continue;
                }

                var componentProperty = dataBlock.Properties().FirstOrDefault();
                if (componentProperty == null)
                {
                    trimmedComponentMap[anchorId] = trimmedComponent;
                    continue;
                }

                var componentName = componentProperty.Name;
                var componentJObject = componentProperty.Value as JObject;
                if (componentJObject == null)
                {
                    trimmedComponentMap[anchorId] = trimmedComponent;
                    continue;
                }

                if (rules.TryGetValue(componentName, out var exposedFields))
                {
                    if (exposedFields.Contains("*"))
                    {
                        trimmedComponentMap[anchorId] = trimmedComponent;
                        continue;
                    }

                    var propertiesToKeep = new HashSet<string>(exposedFields);
                    var currentKeys = componentJObject.Properties().Select(p => p.Name).ToList();
                    foreach (var key in currentKeys)
                    {
                        if (!propertiesToKeep.Contains(key))
                        {
                            componentJObject.Remove(key);
                        }
                    }
                }
                else
                {
                    componentJObject.RemoveAll();
                }
                trimmedComponentMap[anchorId] = trimmedComponent;
            }
            
            // Pass 3: Final Assembly
            var finalOutput = new JArray();
            var embeddedIds = new HashSet<string>();

            // First, identify all components that will be embedded into GameObjects
            foreach (var originalComponent in allComponents.Children<JObject>())
            {
                if (originalComponent["data"]?["GameObject"]?["m_Component"] is JArray componentList)
                {
                    foreach (var compRef in componentList.Children<JObject>())
                    {
                        if (compRef["component"]?["fileID"]?.ToString() is var fileID && !string.IsNullOrEmpty(fileID))
                        {
                            embeddedIds.Add(fileID);
                        }
                    }
                }
            }

            // Now, build the final output, preserving order
            foreach (var originalComponent in allComponents.Children<JObject>())
            {
                var anchorId = originalComponent["anchor_id"]?.ToString();
                if (string.IsNullOrEmpty(anchorId) || embeddedIds.Contains(anchorId))
                {
                    // If a component has no ID, or it's a child of a GameObject, skip it.
                    // Children will be added via their parent GameObject.
                    continue;
                }

                // Get the trimmed version of the current component
                if (!trimmedComponentMap.TryGetValue(anchorId, out var trimmedComponent))
                {
                    // Should not happen, but as a safeguard, add the original.
                    finalOutput.Add(originalComponent.DeepClone());
                    continue;
                }

                // If it's a GameObject, embed its already-trimmed children
                if (originalComponent["data"]?["GameObject"] is JObject originalGameObjectData)
                {
                    var trimmedGameObjectData = trimmedComponent["data"]["GameObject"] as JObject;
                    var componentList = originalGameObjectData["m_Component"] as JArray;
                    if (trimmedGameObjectData != null && componentList != null)
                    {
                        var newComponentList = new JArray();
                        var gameObjectName = originalGameObjectData["m_Name"]?.ToString() ?? "Unnamed";
                        foreach (var compRef in componentList.Children<JObject>())
                        {
                            var fileID = compRef["component"]?["fileID"]?.ToString();
                            if (string.IsNullOrEmpty(fileID)) continue;

                            if (trimmedComponentMap.TryGetValue(fileID, out var linkedComponent))
                            {
                                var clonedData = linkedComponent["data"].DeepClone() as JObject;
                                (clonedData?.Properties().FirstOrDefault()?.Value as JObject)?.Remove("m_GameObject");
                                
                                // Re-inject the anchor_id to preserve the link information after trimming.
                                if (clonedData != null)
                                {
                                    clonedData.AddFirst(new JProperty("anchor_id", fileID));
                                }
                                
                                newComponentList.Add(clonedData);
                            }
                            else
                            {
                                Debug.LogWarning($"JsonDataTrimmer: Could not find component with anchor_id '{fileID}' referenced by GameObject '{gameObjectName}'.");
                            }
                        }
                        trimmedGameObjectData["m_Component"] = newComponentList;
                    }
                }
                
                finalOutput.Add(trimmedComponent);
            }

            return finalOutput.ToString(Newtonsoft.Json.Formatting.None);
        }

        private static string ShortenBrokenPrefabGameObjects(string data)
        {
            var nodes = JsonConvert.DeserializeObject<List<OriginalNode>>(data);
            
            if (nodes == null)
            {
                return data;
            }
            
            const string gameObjectTypeId = "1";
            const string prefabTypeId = "1001";
            const string transformTypeId = "4";
            var allGameObjects = nodes.Where(n => n.type_id == gameObjectTypeId).ToList();
            var allTransforms = nodes.Where(n => n.type_id == transformTypeId).ToList();
            var allPrefabs = nodes.Where(n => n.type_id == prefabTypeId).ToList();
            var restOfComponents = nodes.Where(n => n.type_id != gameObjectTypeId && n.type_id != prefabTypeId && n.type_id != transformTypeId).ToList();
            
            foreach (var gameObject in allGameObjects)
            {
                var nodeIndex = nodes.IndexOf(gameObject);
                var prefabInstanceId = gameObject.data?["GameObject"]?["m_PrefabInstance"]?["fileID"]?.ToString();
                if (string.IsNullOrEmpty(prefabInstanceId) || prefabInstanceId == "0") continue;
                var prefab = allPrefabs.FirstOrDefault(p => p.anchor_id == prefabInstanceId);
                if (prefab == null)
                {
                    Debug.LogWarning($"JsonDataTrimmer: Could not find prefab with anchor_id '{prefabInstanceId}' referenced by GameObject '{gameObject.data?["GameObject"]?["m_Name"]}'.");
                    continue;
                }
                var prefabGuid = prefab.data?["PrefabInstance"]?["m_SourcePrefab"]?["guid"]?.ToString();
                if (string.IsNullOrEmpty(prefabGuid))
                {
                    Debug.LogWarning($"JsonDataTrimmer: Could not find GUID for prefab with anchor_id '{prefabInstanceId}' referenced by GameObject '{gameObject.data?["GameObject"]?["m_Name"]}'.");
                    continue;
                }
                var relatedTransform = allTransforms.FirstOrDefault(t => t.data?["Transform"]?["m_PrefabInstance"]?["fileID"]?.ToString() == prefabInstanceId);
                if (relatedTransform == null)
                {
                    Debug.LogWarning($"JsonDataTrimmer: Could not find Transform with anchor_id '{prefabInstanceId}' referenced by GameObject '{gameObject.data?["GameObject"]?["m_Name"]}'.");
                    continue;
                }
                var relatedComponents = restOfComponents.Where(c => c.data?.SelectToken("*.m_GameObject.fileID")?.ToString() == gameObject.anchor_id).ToList();
                
                var newNode = new OriginalNode
                {
                    type_id = gameObject.type_id,
                    anchor_id = gameObject.anchor_id,
                    data = gameObject.data
                };
                var jArray = new JArray();
                foreach (var component in relatedComponents)
                {
                    jArray.Add(JObject.FromObject(new
                    {
                        component = new
                        {
                            fileID = component.anchor_id
                        }
                    }));
                }
                jArray.Add(JObject.FromObject(new
                {
                    component = new
                    {
                        fileID = relatedTransform.anchor_id
                    }
                }));
                newNode.data["GameObject"]!["m_Component"] = jArray;
                nodes[nodeIndex] = newNode;
            }
            return JsonConvert.SerializeObject(nodes, Formatting.None);
        }
    }
}