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/extract_high_level.py";
        private const string MidLevelScript = "Assets/LLM/source/extract_mid_level.py";
        private const string LowLevelScript = "Assets/LLM/source/extract_low_level.py";

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

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

        [MenuItem("Tools/DataExtractor/Low Level Export")]
        private static void ExportLowLevel()
        {
            RunExtractor("Low Level", LowLevelScript);
        }

        private static void RunExtractor(string level, string scriptPath)
        {
            var 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 = Path.Combine(projectRoot, "venv", "bin", "python3");
            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} extractor was not found.", "OK");
                return;
            }

            var arguments = $"\"{fullScriptPath}\" --input \"{projectRoot}\" --output \"{outputPath}\"";
        
            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} export...", 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("Export Complete", $"{level} data export finished successfully.", "OK");
            }
            else
            {
                EditorUtility.DisplayDialog("Export Failed", $"The {level} data export failed. Check the Unity Console for error messages.", "OK");
            }
        }
    }
}