terra-llm/Assets/LLM/Editor/GUI/MCPWindow.cs

using System.Linq;
using UnityEditor;
using UnityEngine;
using LLM.Editor.Core;
using System.Threading;
using LLM.Editor.Client;
using LLM.Editor.Helper;
using LLM.Editor.Commands;
using LLM.Editor.Analysis;
using System.Threading.Tasks;
using System.Collections.Generic;

namespace LLM.Editor.GUI
{
    public class MCPWindow : EditorWindow
    {
        private string _promptText = "The grenade launcher should hit the target.";
        private Vector2 _scrollPos;
        private bool _isRequestInProgress;
        
        private readonly List<Object> _stagedContextObjects = new();
        private List<string> _requestedRoles = new();
        private bool _isAnalysisContextRequested;
        private ILlmApiClient _apiClient;
        private CancellationTokenSource _cancellationTokenSource;
        private List<string> _currentPlan = new();

        // --- Fields for Dummy API Client UI ---
        private List<string> _dummyTriggerPhrases = new();
        private int _selectedDummyTriggerIndex;
        private bool _dummyResponsesLoaded;

        [MenuItem("LLM/MCP Assistant")]
        public static void ShowWindow()
        {
            GetWindow<MCPWindow>("MCP Assistant");
        }

        private void OnEnable()
        {
            CommandExecutor.OnQueueUpdated += Repaint;
            CommandExecutor.OnPlanUpdated += HandlePlanUpdated;
            RequestAnalysisContextCommand.OnAnalysisContextRequested += HandleAnalysisContextRequested;
            CommandExecutor.OnContextReadyForNextTurn += HandleContextReadyForNextTurn;
            _apiClient = ApiClientFactory.GetClient();
            
            LoadDummyTestCases();
        }

        private void OnDisable()
        {
            CommandExecutor.OnQueueUpdated -= Repaint;
            CommandExecutor.OnPlanUpdated -= HandlePlanUpdated;
            RequestAnalysisContextCommand.OnAnalysisContextRequested -= HandleAnalysisContextRequested;
            CommandExecutor.OnContextReadyForNextTurn -= HandleContextReadyForNextTurn;
            
            _cancellationTokenSource?.Cancel();
            _cancellationTokenSource?.Dispose();
        }

        private void HandlePlanUpdated(List<string> plan)
        {
            _currentPlan = plan;
            Repaint();
        }

        private void HandleAnalysisContextRequested(RequestAnalysisContextParams contextParams)
        {
            _isAnalysisContextRequested = true;
            _requestedRoles = new List<string>(contextParams.subjectRoles);
            while (_stagedContextObjects.Count < _requestedRoles.Count)
            {
                _stagedContextObjects.Add(null);
            }
            Repaint();
        }
        
        private void HandleContextReadyForNextTurn(string detailedContext)
        {
            Debug.Log("[MCPWindow] Received detailed context. Sending follow-up to LLM.");
            _ = SendFollowUpAsync(detailedContext);
        }

        private void OnGUI()
        {
            EditorGUILayout.LabelField("LLM Co-Pilot", EditorStyles.boldLabel);
            
            DrawSessionManagement();
            EditorGUILayout.Space();

            EditorGUI.BeginDisabledGroup(!SessionManager.HasActiveSession());

            DrawContextStagingArea();
            EditorGUILayout.Space();
            if (!(_apiClient is DummyApiClient))
            {
                DrawSelectionHint();
            }
            EditorGUILayout.Space();

            EditorGUILayout.LabelField("Enter your request:");
            
            // --- Conditional UI for Prompt Input ---
            if (_apiClient is DummyApiClient)
            {
                if (!_dummyResponsesLoaded)
                {
                    EditorGUILayout.HelpBox("Loading dummy responses...", MessageType.Info);
                }
                else if (_dummyTriggerPhrases.Any())
                {
                    _selectedDummyTriggerIndex = EditorGUILayout.Popup("Select Test Case", _selectedDummyTriggerIndex, _dummyTriggerPhrases.ToArray());
                    // Update prompt text based on selection to be used by the send function
                    if(_selectedDummyTriggerIndex < _dummyTriggerPhrases.Count)
                    {
                        _promptText = _dummyTriggerPhrases[_selectedDummyTriggerIndex];
                    }
                }
                else
                {
                    EditorGUILayout.HelpBox("No dummy test cases found in DummyResponses.json.", MessageType.Warning);
                }
            }
            else
            {
                // Default text area for other clients like Gemini
                _promptText = EditorGUILayout.TextArea(_promptText, GUILayout.Height(60));
            }
            // --- End of Conditional UI ---

            DrawActionButtons();
            EditorGUILayout.Space();
            
            DrawPlanArea();
            EditorGUILayout.Space();
            
            DrawCommandQueue();
            
            EditorGUI.EndDisabledGroup();
        }
        
        private void LoadDummyTestCases()
        {
            var testCases = DummyApiClient.GetTestCases();
            if (testCases != null && testCases.Any())
            {
                _dummyTriggerPhrases = testCases.Select(c => c.TriggerPhrase).ToList();
            }
            else
            {
                _dummyTriggerPhrases = new List<string> { "No C# test cases found" };
            }
            _dummyResponsesLoaded = true;
            Repaint();
        }
        
        private void DrawSessionManagement()
        {
            EditorGUILayout.BeginHorizontal();
            if (SessionManager.HasActiveSession())
            {
                EditorGUILayout.LabelField($"Session: {SessionManager.GetCurrentSessionId()}");
                if (GUILayout.Button("End Session"))
                {
                    SessionManager.EndSession();
                    ClearStagingArea();
                    _currentPlan.Clear();
                }
            }
            else
            {
                if (GUILayout.Button("Start New Session"))
                {
                    SessionManager.StartNewSession();
                }
            }
            EditorGUILayout.EndHorizontal();
        }
        
        private void DrawContextStagingArea()
        {
            EditorGUILayout.LabelField("Manual Context", EditorStyles.boldLabel);
            EditorGUILayout.HelpBox("You can manually provide context by dragging and dropping any asset (prefab, script, etc.) or GameObject into the slots below.", MessageType.Info);

            if (GUILayout.Button("Add Manual Context Slot", GUILayout.Width(200)))
            {
                _stagedContextObjects.Add(null);
                if (_isAnalysisContextRequested) _requestedRoles.Add("Custom Role");
            }

            var indexToRemove = -1;

            for (var i = 0; i < _stagedContextObjects.Count; i++)
            {
                EditorGUILayout.BeginHorizontal();
                var label = (_isAnalysisContextRequested && i < _requestedRoles.Count) ? _requestedRoles[i] : $"Slot {i + 1}";
                _stagedContextObjects[i] = EditorGUILayout.ObjectField(label, _stagedContextObjects[i], typeof(Object), true);

                if (GUILayout.Button("X", GUILayout.Width(20)))
                {
                    indexToRemove = i;
                }
                EditorGUILayout.EndHorizontal();
            }

            if (indexToRemove != -1)
            {
                _stagedContextObjects.RemoveAt(indexToRemove);
                if (_isAnalysisContextRequested && indexToRemove < _requestedRoles.Count)
                {
                    _requestedRoles.RemoveAt(indexToRemove);
                }
                Repaint();
            }
        }

        private void DrawSelectionHint()
        {
            EditorGUILayout.LabelField("Automatic Context (From Selection)", EditorStyles.boldLabel);
            var selectedObjects = Selection.objects;

            if (selectedObjects.Length > 0)
            {
                EditorGUILayout.HelpBox("The following selected objects will automatically be included as context.", MessageType.None);
                
                EditorGUI.BeginDisabledGroup(true);
                foreach (var obj in selectedObjects)
                {
                    EditorGUILayout.ObjectField(obj.name, obj, typeof(Object), true);
                }
                EditorGUI.EndDisabledGroup();
            }
            else
            {
                EditorGUILayout.HelpBox("No objects are currently selected in the editor.", MessageType.None);
            }
        }
        
        private void DrawActionButtons()
        {
            EditorGUILayout.BeginHorizontal();
            
            EditorGUI.BeginDisabledGroup(_isRequestInProgress);

            if (_isAnalysisContextRequested)
            {
                if (GUILayout.Button("Provide Staged Context to LLM"))
                {
                    _stagedContextObjects.AddRange(Selection.objects);
                    var contextSummary = ContextBuilder.BuildTier1Summary(_stagedContextObjects);
                    var followUpMessage = "Here is the context you requested:\n" + contextSummary;
                    _ = SendFollowUpAsync(followUpMessage);
                }
            }
            else
            {
                if (GUILayout.Button("Send Prompt"))
                {
                    _ = SendInitialPromptAsync();
                }
            }
            
            EditorGUI.EndDisabledGroup();
            
            EditorGUI.BeginDisabledGroup(!_isRequestInProgress);
            if (GUILayout.Button("Stop", GUILayout.Width(80)))
            {
                _cancellationTokenSource?.Cancel();
            }
            EditorGUI.EndDisabledGroup();
            
            EditorGUILayout.EndHorizontal();
        }

        private void DrawPlanArea()
        {
            if (_currentPlan == null || !_currentPlan.Any()) return;
            
            EditorGUILayout.LabelField("Current Plan", EditorStyles.boldLabel);
            
            EditorGUILayout.BeginVertical(EditorStyles.helpBox);
            for(var i = 0; i < _currentPlan.Count; i++)
            {
                var step = _currentPlan[i];
                var label = $"{i + 1}. {step}";
                EditorGUILayout.LabelField(label, EditorStyles.wordWrappedLabel);
            }
            EditorGUILayout.EndVertical();
        }
        
        private void DrawCommandQueue()
        {
            EditorGUILayout.LabelField("Pending Commands", EditorStyles.boldLabel);
            _scrollPos = EditorGUILayout.BeginScrollView(_scrollPos, EditorStyles.helpBox);
            
            if (CommandExecutor.HasPendingCommands())
            {
                var nextCommand = CommandExecutor.GetNextCommand();
                EditorGUILayout.LabelField("Next Up: " + nextCommand.commandName);

                EditorGUI.BeginDisabledGroup(_isRequestInProgress);
                if (GUILayout.Button("Execute Next Command"))
                {
                    _ = CommandExecutor.ExecuteNextCommand();
                }
                EditorGUI.EndDisabledGroup();
            }
            else
            {
                EditorGUILayout.LabelField("No pending commands.");
            }

            EditorGUILayout.EndScrollView();
        }
        
        private async Task SendInitialPromptAsync()
        {
            _isRequestInProgress = true;
            _cancellationTokenSource = new CancellationTokenSource();
            _currentPlan.Clear();
            Repaint();
            
            try
            {
                var combinedContext = new List<Object>();
                combinedContext.AddRange(_stagedContextObjects);
                if (!(_apiClient is DummyApiClient))
                {
                    combinedContext.AddRange(Selection.objects);
                }
            
                await _apiClient.SendPrompt(_promptText, combinedContext, _cancellationTokenSource.Token);
            }
            catch (TaskCanceledException)
            {
                Debug.Log("[MCPWindow] Task was cancelled.");
            }
            finally
            {
                _isRequestInProgress = false;
                _cancellationTokenSource.Dispose();
                _cancellationTokenSource = null;
                Repaint();
            }
        }
        
        private async Task SendFollowUpAsync(string followUpMessage)
        {
            _isRequestInProgress = true;
            _cancellationTokenSource = new CancellationTokenSource();
            Repaint();

            try
            {
                await _apiClient.SendFollowUp(followUpMessage, _cancellationTokenSource.Token);
            }
            catch (TaskCanceledException)
            {
                Debug.Log("[MCPWindow] Task was cancelled.");
            }
            finally
            {
                _isRequestInProgress = false;
                _cancellationTokenSource.Dispose();
                _cancellationTokenSource = null;
                
                if (_isAnalysisContextRequested)
                {
                    ClearStagingArea();
                }
                Repaint();
            }
        }
        
        private void ClearStagingArea()
        {
            _isAnalysisContextRequested = false;
            _stagedContextObjects.Clear();
            _requestedRoles.Clear();
        }
    }
}