new-git/Assets/Better Git/Arbitrator/Editor/GUI/ArbitratorController.cs

// Copyright (c) 2025 TerraByte Inc.
//
// This script acts as the Controller for the ArbitratorWindow. It manages all
// state and business logic, separating it from the UI rendering code in the window.

using System;
using GitMerge;
using System.Linq;
using UnityEditor;
using UnityEngine.Scripting;
using Terra.Arbitrator.Settings;
using Terra.Arbitrator.Services;
using Terra.Arbitrator.Promises;
using System.Collections.Generic;
using UnityEditor.SceneManagement;

namespace Terra.Arbitrator.GUI
{
    public enum UserAction { Proceed, SaveAndProceed, Cancel }

    public enum SortColumn { Commit, Status, FilePath }
    
    [Preserve]
    public class ArbitratorController
    {
        private List<GitChange> _changes = new();
        public IReadOnlyList<GitChange> Changes => _changes;
        public string InfoMessage { get; private set; }
        public string ErrorMessage { get; private set; }
        public bool IsLoading { get; private set; }
        public string LoadingMessage { get; private set; } = "";
        public bool IsInConflictState { get; private set; }
        public int CommitsToPull { get; private set; }
        public SortColumn CurrentSortColumn { get; private set; } = SortColumn.FilePath;
        public float OperationProgress { get; private set; }
        public string OperationProgressMessage { get; private set; }
        public IReadOnlyList<string> RemoteBranchList { get; private set; } = new List<string>();
        public string CurrentBranchName { get; private set; } = "master";

        private readonly Action _requestRepaint;
        private readonly Func<string, string, string, string, bool> _displayDialog;
        private readonly Func<UserAction> _promptForUnsavedChanges;

        /// <summary>
        /// Initializes the controller.
        /// </summary>
        /// <param name="requestRepaint">A callback to the window's Repaint() method.</param>
        /// <param name="displayDialog">A callback to EditorUtility.DisplayDialog for user confirmations.</param>
        /// <param name="promptForUnsavedChanges">A callback to EditorUtility.DisplayDialog for user confirmations.</param>
        public ArbitratorController(Action requestRepaint, Func<string, string, string, string, bool> displayDialog, Func<UserAction> promptForUnsavedChanges)
        {
            _requestRepaint = requestRepaint;
            _displayDialog = displayDialog;
            _promptForUnsavedChanges = promptForUnsavedChanges;
        }
        
        public void OnEnable()
        {
            if (SessionState.GetBool(BetterGitStatePersistence.ResetQueueKey, false))
            {
                SessionState.EraseString(BetterGitStatePersistence.ResetQueueKey);
                InfoMessage = "Multi-file reset complete. Pulling again to confirm...";
                Pull();
            }
            else
            {
                Refresh();
            }
        }

        public void Refresh()
        {
            StartOperation("Refreshing status...");
            CommitsToPull = 0;
            
            UnstageStep()
                .Then(CompareStep)
                .Then(FetchUpstreamStep)
                .Then(FetchBranchDataStep)
                .Then(FinalizeRefresh)
                .Catch(HandleOperationError)
                .Finally(FinishOperation);
        }
        
        public void Pull()
        {
            if (IsLoading) return;
            if (CancelOperationIfUnsavedScenes()) return;
            if (CommitsToPull > 0)
            {
                if (!_displayDialog("Confirm Pull", $"There are {CommitsToPull} incoming changes. Are you sure you want to pull?", "Yes, Pull", "Cancel"))
                {
                    return;
                }
            }

            StartOperation("Analyzing for conflicts...");

            GitService.AnalyzePullConflicts()
                .Then(analysisResult =>
                {
                    FinishOperation();
                    if (analysisResult.HasConflicts)
                    {
                        var conflictingChanges = _changes.Where(c => analysisResult.ConflictingFiles.Contains(c.FilePath)).ToList();
                        ConflictResolutionWindow.ShowWindow(this, conflictingChanges);
                    }
                    else
                    {
                        PerformSafePullWithLock();
                    }
                })
                .Catch(ex => {
                    HandleOperationError(ex);
                    FinishOperation();
                })
                .Finally(FinishOperation);
        }

        private void PerformSafePullWithLock()
        {
            StartOperation("Pulling changes...");
            EditorApplication.LockReloadAssemblies();
            GitService.PerformSafePull()
                .Then(successMessage =>
                {
                    InfoMessage = successMessage;
                    Refresh();
                })
                .Catch(ex => {
                    HandleOperationError(ex);
                    FinishOperation();
                })
                .Finally(EditorApplication.UnlockReloadAssemblies);
        }
        
        public void ResolveConflictsAndPull(List<GitChange> resolutions)
        {
            StartOperation("Resolving conflicts and pulling...");
            EditorApplication.LockReloadAssemblies();

            GitService.PullAndOverwrite(resolutions)
                .Then(successMessage =>
                {
                    InfoMessage = successMessage;
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    EditorApplication.UnlockReloadAssemblies();
                    AssetDatabase.Refresh();
                    Refresh();
                });
        }

        public void CommitAndPush(string commitMessage)
        {   
            if (CancelOperationIfUnsavedScenes()) return;
            
            var selectedFiles = _changes.Where(c => c.IsSelectedForCommit).ToList();
            
            var username = BetterGitSettings.Username;
            var email = BetterGitSettings.Email;
            
            StartOperation("Staging, committing, and pushing files...");

            GitService.CommitAndPush(selectedFiles, commitMessage, username, email, OnProgressModified)
                .Then(successMessage => {
                    InfoMessage = successMessage;
                    Refresh();
                })
                .Catch(ex => {
                    HandleOperationError(ex);
                    FinishOperation();
                });
            return;

            void OnProgressModified(float progress, string message)
            {
                OperationProgress = progress;
                OperationProgressMessage = message;
                _requestRepaint?.Invoke();
            }
        }
        
        public void SetSortColumn(SortColumn newColumn)
        {
            // If it's already the active column, do nothing.
            if (CurrentSortColumn == newColumn) return;

            CurrentSortColumn = newColumn;
            ApplyGrouping();
            _requestRepaint?.Invoke();
        }

        private void ApplyGrouping()
        {
            if (_changes == null || !_changes.Any()) return;

            _changes = CurrentSortColumn switch
            {
                SortColumn.Commit => _changes.OrderByDescending(c => c.IsSelectedForCommit).ThenBy(c => c.FilePath).ToList(),
                SortColumn.Status => _changes.OrderBy(c => GetStatusSortPriority(c.Status)).ThenBy(c => c.FilePath).ToList(),
                SortColumn.FilePath => _changes.OrderBy(c => c.FilePath).ToList(),
                _ => _changes
            };
        }
        
        private static int GetStatusSortPriority(LibGit2Sharp.ChangeKind status)
        {
            return status switch
            {
                LibGit2Sharp.ChangeKind.Conflicted => -1, // Always show conflicts on top
                LibGit2Sharp.ChangeKind.Modified => 0,
                LibGit2Sharp.ChangeKind.Added => 1,
                LibGit2Sharp.ChangeKind.Deleted => 2,
                LibGit2Sharp.ChangeKind.Renamed => 3,
                _ => 99
            };
        }

        public void ResetFile(GitChange change)
        {
            if (IsLoading) return;
            var userConfirmed = _displayDialog("Confirm Reset", $"Are you sure you want to revert all local changes to '{change.FilePath}'? This action cannot be undone.", "Yes, Revert", "Cancel");
            if (!userConfirmed) return;
            
            StartOperation($"Resetting {change.FilePath}...");
            
            GitService.ResetFileChanges(change)
                .Then(successMessage => {
                    InfoMessage = successMessage;
                    Refresh();
                })
                .Catch(ex => {
                    HandleOperationError(ex);
                    FinishOperation();
                });
        }
        
        public void DiffFile(GitChange change)
        {
            if (IsLoading) return;
            StartOperation($"Launching diff for {change.FilePath}...");
            
            GitService.LaunchExternalDiff(change)
                .Catch(HandleOperationError)
                .Finally(Refresh);
        }
        
        public void ResolveConflict(GitChange change)
        {
            if (IsLoading) return;
            
            StartOperation($"Opening merge tool for {change.FilePath}...");
            
            var fileExtension = System.IO.Path.GetExtension(change.FilePath).ToLower();
            if (fileExtension is ".prefab" or ".unity")
            {
                try
                {
                    GitMergeWindow.ResolveConflict(change.FilePath);
                }
                catch (Exception e)
                {
                    ErrorMessage = e.Message;
                }
                
                Refresh();
                return;
            }
            
            GitService.LaunchMergeTool(change)
                .Then(successMessage => { InfoMessage = successMessage; })
                .Catch(HandleOperationError)
                .Finally(Refresh);
        }
        
        public void ResetSelected()
        {
            var selectedFiles = _changes.Where(c => c.IsSelectedForCommit).ToList();
            if (!selectedFiles.Any()) return;

            var fileList = string.Join("\n - ", selectedFiles.Select(f => f.FilePath));
            if (!_displayDialog("Confirm Reset Selected", $"Are you sure you want to revert changes for the following {selectedFiles.Count} file(s)?\n\n - {fileList}", "Yes, Revert Selected", "Cancel")) return;
            
            var pathsToReset = selectedFiles.Select(c => c.FilePath).ToList();
            ResetMultipleFiles(pathsToReset);
        }

        public void SetAllSelection(bool selected)
        {
            if (_changes == null) return;
            foreach (var change in _changes.Where(change => change.Status != LibGit2Sharp.ChangeKind.Conflicted))
            {
                change.IsSelectedForCommit = selected;
            }
        }
        
        public void SwitchToBranch(string targetBranch)
        {
            if (IsLoading || targetBranch == CurrentBranchName) return;

            if (Changes.Any())
            {
                if (!_displayDialog("Discard Local Changes?", $"You have local changes. To switch branches, these changes must be discarded.\n\nDiscard changes and switch to '{targetBranch}'?", "Yes, Discard and Switch", "Cancel"))
                {
                    return;
                }
                
                StartOperation($"Discarding changes and switching to {targetBranch}...");
                GitService.ResetAndSwitchBranch(targetBranch)
                    .Then(successMsg => { InfoMessage = successMsg; Refresh(); })
                    .Catch(ex => { HandleOperationError(ex); FinishOperation(); })
                    .Finally(() => { EditorApplication.delayCall += AssetDatabase.Refresh; });
            }
            else
            {
                if (!_displayDialog("Confirm Branch Switch", $"Are you sure you want to switch to branch '{targetBranch}'?", "Yes, Switch", "Cancel"))
                {
                    return;
                }

                StartOperation($"Switching to {targetBranch}...");
                GitService.SwitchBranch(targetBranch)
                    .Then(successMsg => { InfoMessage = successMsg; Refresh(); })
                    .Catch(ex => { HandleOperationError(ex); FinishOperation(); })
                    .Finally(() => { EditorApplication.delayCall += AssetDatabase.Refresh; });
            }
        }

        // --- Private Methods ---
        
        private static IPromise<bool> UnstageStep()
        {
            return GitService.UnstageAllFilesIfSafe();
        }

        private IPromise<List<GitChange>> CompareStep(bool wasUnstaged)
        {
            if (wasUnstaged)
            {
                InfoMessage = "Found and unstaged files for review.";
            }
            return GitService.CompareLocalToRemote();
        }
        
        private IPromise<int?> FetchUpstreamStep(List<GitChange> changes)
        {
            _changes = changes;
            IsInConflictState = _changes.Any(c => c.Status == LibGit2Sharp.ChangeKind.Conflicted);
            return IsInConflictState ? new Promise<int?>((resolve, _) => resolve(0)) : GitService.GetUpstreamAheadBy(OnProgressModified);

            void OnProgressModified(float progress, string message)
            {
                OperationProgress = progress;
                OperationProgressMessage = message;
                _requestRepaint?.Invoke();
            }
        }

        private IPromise<BranchData> FetchBranchDataStep(int? pullCount)
        {
            CommitsToPull = pullCount ?? 0;
            OperationProgress = 0f;
            OperationProgressMessage = "";
            return GitService.GetBranchData();
        }

        private void FinalizeRefresh(BranchData branchData)
        {
            CurrentBranchName = branchData.CurrentBranch;
            RemoteBranchList = branchData.AllBranches;
            ApplyGrouping();
        }
        
        // --- Shared Helper Methods ---

        private void StartOperation(string loadingMessage)
        {
            IsLoading = true;
            LoadingMessage = loadingMessage;
            OperationProgress = 0f;
            OperationProgressMessage = "";
            ClearMessages();
            _changes = null;
            _requestRepaint?.Invoke();
        }

        private void HandleOperationError(Exception ex)
        {
            ErrorMessage = $"Operation Failed: {ex.Message}";
        }
        
        private void FinishOperation()
        {
            IsLoading = false;
            _requestRepaint?.Invoke();
        }
        
        private void ClearMessages()
        {
            ErrorMessage = null;
            InfoMessage = null;
        }

        private void ResetMultipleFiles(List<string> filePaths)
        {
            var hasScripts = filePaths.Any(p => p.EndsWith(".cs", StringComparison.OrdinalIgnoreCase));

            if (hasScripts)
            {
                InfoMessage = $"Starting reset for {filePaths.Count} file(s)... This may trigger script compilation.";
                _requestRepaint?.Invoke();
                SessionState.SetString("BetterGit.ResetQueue", string.Join(";", filePaths));
                EditorApplication.delayCall += BetterGitStatePersistence.ContinueInterruptedReset;
            }
            else
            {
                StartOperation($"Resetting {filePaths.Count} file(s)...");

                IPromise<string> promiseChain = new Promise<string>((resolve, _) => resolve(""));

                foreach (var path in filePaths)
                {
                    promiseChain = promiseChain.Then(_ => {
                        var change = GitService.GetChangeForFile(path);
                        return change != null ? GitService.ResetFileChanges(change) : new Promise<string>((res, _) => res(""));
                    });
                }

                promiseChain
                    .Then(successMsg => {
                        InfoMessage = $"Successfully reset {filePaths.Count} file(s).";
                        Refresh();
                    })
                    .Catch(ex => {
                        HandleOperationError(ex);
                        FinishOperation();
                    });
            }
        }
        
        public void ForcePull()
        {
            StartOperation("Attempting to pull and create conflicts...");
            
            EditorApplication.LockReloadAssemblies();
            GitService.ForcePull()
                .Then(_ =>
                {
                    InfoMessage = "Pull resulted in conflicts. Please resolve them below.";
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    EditorApplication.UnlockReloadAssemblies();
                    AssetDatabase.Refresh();
                    Refresh();
                });
        }
        
        private bool CancelOperationIfUnsavedScenes()
        {
            var isAnySceneDirty = false;
            for (var i = 0; i < EditorSceneManager.sceneCount; i++)
            {
                var scene = EditorSceneManager.GetSceneAt(i);
                if (!scene.isDirty) continue;
                isAnySceneDirty = true;
                break;
            }

            if (!isAnySceneDirty)
            {
                return false;
            }
            
            var userChoice = _promptForUnsavedChanges();

            switch (userChoice)
            {
                case UserAction.SaveAndProceed:
                    EditorSceneManager.SaveOpenScenes();
                    return false;
                case UserAction.Proceed:
                    return false;
                case UserAction.Cancel:
                default:
                    return true;
            }
        }
    }
}