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

using AssetBank.Settings;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using UnityEditor;
using UnityEngine;
using Newtonsoft.Json.Linq;

namespace AssetBank.Editor.Tools
{
    /// <summary>
    /// Handles the core logic for scanning the project and managing the component whitelist.
    /// </summary>
    public class ComponentWhitelistController
    {
        private static readonly string ProjectRoot = Path.GetDirectoryName(Application.dataPath);
        
        public enum ScanSource
        {
            Scenes,
            Prefabs,
            Both
        }

        /// <summary>
        /// Scans the project for components based on the selected source.
        /// </summary>
        /// <param name="scanSource">The type of assets to scan.</param>
        /// <returns>A HashSet of unique component type names found.</returns>
        public HashSet<string> ScanProject(ScanSource scanSource)
        {
            var settings = ProjectExporterSettings.GetOrCreateSettings();
            var ignoredFolders = settings.FoldersToIgnore;
            var assetsToScan = new List<string>();

            if (scanSource == ScanSource.Scenes || scanSource == ScanSource.Both)
            {
                assetsToScan.AddRange(AssetDatabase.FindAssets("t:Scene", new[] { "Assets" }));
            }
            if (scanSource == ScanSource.Prefabs || scanSource == ScanSource.Both)
            {
                assetsToScan.AddRange(AssetDatabase.FindAssets("t:Prefab", new[] { "Assets" }));
            }

            var assetPaths = assetsToScan.Distinct().Select(AssetDatabase.GUIDToAssetPath).ToList();
            var filteredPaths = assetPaths.Where(path => !ignoredFolders.Any(folder => path.StartsWith(folder, StringComparison.OrdinalIgnoreCase))).ToList();

            var tempDir = Path.Combine("Library", "ComponentScanTemp");
            if (Directory.Exists(tempDir))
            {
                Directory.Delete(tempDir, true);
            }
            Directory.CreateDirectory(tempDir);

            var uniqueComponents = new HashSet<string>();

            try
            {
                EditorUtility.DisplayProgressBar("Scanning Components", "Preparing to scan...", 0f);
                for (int i = 0; i < filteredPaths.Count; i++)
                {
                    var assetPath = filteredPaths[i];
                    var progress = (float)i / filteredPaths.Count;
                    EditorUtility.DisplayProgressBar("Scanning Components", $"Processing: {Path.GetFileName(assetPath)}", progress);

                    var jsonPath = Path.Combine(tempDir, $"{Path.GetFileNameWithoutExtension(assetPath)}_{i}.json");
                    if (ExecutePythonConversion(assetPath, jsonPath))
                    {
                        var components = ParseComponentsFromJson(jsonPath);
                        uniqueComponents.UnionWith(components);
                    }
                }
            }
            finally
            {
                if (Directory.Exists(tempDir))
                {
                    Directory.Delete(tempDir, true);
                }
                EditorUtility.ClearProgressBar();
            }

            return uniqueComponents;
        }
        
        private string FindPythonExecutable()
        {
            var venvPath = Path.Combine(ProjectRoot, "venv", "bin", "python3");
            if (File.Exists(venvPath))
            {
                return venvPath;
            }
            venvPath = Path.Combine(ProjectRoot, "venv", "bin", "python");
            if (File.Exists(venvPath))
            {
                return venvPath;
            }
            return "python3";
        }

        /// <summary>
        /// Executes the Python conversion script for a given asset.
        /// </summary>
        private bool ExecutePythonConversion(string assetPath, string jsonOutputPath)
        {
            var pythonExecutable = FindPythonExecutable();
            var scriptGuid = AssetDatabase.FindAssets("convert_scene").FirstOrDefault();
            if (string.IsNullOrEmpty(scriptGuid))
            {
                UnityEngine.Debug.LogError("Conversion script 'convert_scene.py' not found.");
                return false;
            }
            var scriptPath = Path.Combine(ProjectRoot, AssetDatabase.GUIDToAssetPath(scriptGuid));
            var absoluteAssetPath = Path.Combine(ProjectRoot, assetPath);

            var process = new Process
            {
                StartInfo = new ProcessStartInfo
                {
                    FileName = pythonExecutable,
                    Arguments = $"\"{scriptPath}\" \"{absoluteAssetPath}\" \"{jsonOutputPath}\"",
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true
                }
            };

            process.Start();
            string error = process.StandardError.ReadToEnd();
            process.WaitForExit();

            if (process.ExitCode == 0) return true;

            UnityEngine.Debug.LogError($"Failed to convert {assetPath}. Error: {error}");
            return false;
        }

        /// <summary>
        /// Parses a temporary JSON file to extract unique component names.
        /// </summary>
        private IEnumerable<string> ParseComponentsFromJson(string jsonPath)
        {
            var components = new HashSet<string>();
            try
            {
                var jsonContent = File.ReadAllText(jsonPath);
                var jsonArray = JArray.Parse(jsonContent);

                foreach (var item in jsonArray)
                {
                    var data = item["data"];
                    if (data is JObject dataObject)
                    {
                        var componentName = dataObject.Properties().FirstOrDefault()?.Name;
                        if (!string.IsNullOrEmpty(componentName))
                        {
                            components.Add(componentName);
                        }
                    }
                }
            }
            catch (Exception e)
            {
                UnityEngine.Debug.LogError($"Failed to parse JSON file {jsonPath}. Error: {e.Message}");
            }
            return components;
        }

        /// <summary>
        /// Saves the selected components to the ComponentWhitelistSettings asset.
        /// </summary>
        public void SaveWhitelist(Dictionary<string, bool> selection)
        {
            var settings = ComponentWhitelistSettings.GetOrCreateSettings();
            settings.WhitelistedComponents.Clear();
            settings.WhitelistedComponents.AddRange(selection.Where(kvp => kvp.Value).Select(kvp => kvp.Key));

            EditorUtility.SetDirty(settings);
            AssetDatabase.SaveAssets();

            EditorUtility.DisplayDialog("Success", $"Whitelist saved with {settings.WhitelistedComponents.Count} components.", "OK");
        }
    }
}