terra-llm/Assets/LLM/Editor/DataExtractorMenu.cs

using System.IO;
using UnityEditor;
using UnityEngine;
using System.Diagnostics;

namespace LLM.Editor
{
    public static class DataExtractorMenu
    {
        private const string HighLevelScript = "Assets/LLM/source/orchestrators/extract_high_level.py";
        private const string MidLevelScript = "Assets/LLM/source/orchestrators/extract_mid_level.py";
        private const string LowLevelScript = "Assets/LLM/source/orchestrators/extract_low_level.py";
        private const string CreateConfigScript = "Assets/LLM/source/orchestrators/create_config.py";
        private const string TRSExtractorScript = "Assets/LLM/source/orchestrators/extract_scene_trs.py";

        [MenuItem("Tools/DataExtractor/High Level Export")]
        private static void ExportHighLevel()
        {
            RunExtractor("High Level", HighLevelScript, true);
        }

        [MenuItem("Tools/DataExtractor/Mid Level Export")]
        private static void ExportMidLevel()
        {
            RunExtractor("Mid Level", MidLevelScript, true);
        }

        [MenuItem("Tools/DataExtractor/Low Level Export")]
        private static void ExportLowLevel()
        {
            RunExtractor("Low Level", LowLevelScript, true);
        }
        
        [MenuItem("Tools/DataExtractor/TRS Extractor")]
        private static void ExportTRS()
        {
            RunExtractor("TRS", TRSExtractorScript, true);
        }
        
        [MenuItem("Tools/DataExtractor/Create or Update Config File")]
        private static void CreateOrUpdateConfigFile()
        {
            RunExtractor("Create Config", CreateConfigScript, false);
        }

        private static void RunExtractor(string level, string scriptPath, bool requireOutputPath)
        {
            string outputPath = "";
            if (requireOutputPath)
            {
                outputPath = EditorUtility.OpenFolderPanel($"Select Output Folder for {level} Export", "", "");
                if (string.IsNullOrEmpty(outputPath))
                {
                    UnityEngine.Debug.Log($"{level} export cancelled by user.");
                    return;
                }
            }

            var projectRoot = Path.GetFullPath(Path.Combine(Application.dataPath, ".."));

            var pythonExecutable = Application.platform == RuntimePlatform.WindowsEditor ? 
                Path.Combine(projectRoot, "venv", "Scripts", "python.exe") : 
                Path.Combine(projectRoot, "venv", "bin", "python");
            
            var fullScriptPath = Path.Combine(projectRoot, scriptPath);

            if (!File.Exists(pythonExecutable))
            {
                UnityEngine.Debug.LogError($"Python executable not found at: {pythonExecutable}");
                EditorUtility.DisplayDialog("Export Error", "The Python virtual environment executable was not found. Please ensure the 'venv' directory is set up correctly.", "OK");
                return;
            }

            if (!File.Exists(fullScriptPath))
            {
                UnityEngine.Debug.LogError($"Extractor script not found at: {fullScriptPath}");
                EditorUtility.DisplayDialog("Export Error", $"The Python script for the {level} task was not found.", "OK");
                return;
            }

            var arguments = requireOutputPath ? 
                $"\"{fullScriptPath}\" --input \"{projectRoot}\" --output \"{outputPath}\"" : 
                $"\"{fullScriptPath}\"    ";
        
            UnityEngine.Debug.Log($"Running command: \"{pythonExecutable}\" {arguments}");

            var process = new Process
            {
                StartInfo = new ProcessStartInfo
                {
                    FileName = pythonExecutable,
                    Arguments = arguments,
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true,
                    WorkingDirectory = projectRoot
                }
            };

            EditorUtility.DisplayProgressBar("Data Extractor", $"Running {level} task...", 0.5f);
        
            process.OutputDataReceived += (_, e) => { if (e.Data != null) UnityEngine.Debug.Log($"[Extractor] {e.Data}"); };
            process.ErrorDataReceived += (_, e) => { if (e.Data != null) UnityEngine.Debug.LogError($"[Extractor ERROR] {e.Data}"); };

            process.Start();
            process.BeginOutputReadLine();
            process.BeginErrorReadLine();

            process.WaitForExit();
            EditorUtility.ClearProgressBar();

            if (process.ExitCode == 0)
            {
                EditorUtility.DisplayDialog("Task Complete", $"{level} task finished successfully.", "OK");
            }
            else
            {
                EditorUtility.DisplayDialog("Task Failed", $"The {level} task failed. Check the Unity Console for error messages.", "OK");
            }
        }
    }
}