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;
using UnityEngine;

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 bool HasStash { 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()
        {
            Refresh();
        }

        public void Refresh()
        {
            StartOperation("Refreshing status...");
            CommitsToPull = 0;
            
            UnstageStep()
                .Then(CompareStep)
                .Then(FetchUpstreamStep)
                .Then(CheckForStashStep)
                .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, new PullConflictSource());
                    }
                    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 StashSelectedFiles()
        {
            var selectedFiles = _changes.Where(c => c.IsSelectedForCommit).ToList();
            if (!selectedFiles.Any())
            {
                InfoMessage = "No files selected to stash.";
                return;
            }
            
            EditorApplication.LockReloadAssemblies();
            
            StartOperation("Checking for existing stash...");
            GitService.HasStash()
                .Then(hasStash =>
                {
                    FinishOperation();
                    bool confirmed;
                    if (hasStash)
                    {
                        confirmed = _displayDialog("Overwrite Stash?", 
                            "A 'Better Git Stash' already exists. Do you want to overwrite it with your currently selected files?", 
                            "Yes, Overwrite", "Cancel");
                    }
                    else
                    {
                        confirmed = _displayDialog("Create Stash?", 
                            $"Are you sure you want to stash the selected {selectedFiles.Count} file(s)?", 
                            "Yes, Stash", "Cancel");
                    }

                    if (!confirmed) return;
                    
                    StartOperation("Stashing selected files...");
                    GitService.CreateOrOverwriteStash(selectedFiles)
                        .Then(successMsg => { InfoMessage = successMsg; })
                        .Catch(HandleOperationError)
                        .Finally(Refresh);
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    FinishOperation();
                    Refresh();
                    EditorApplication.UnlockReloadAssemblies();
                    AssetDatabase.Refresh();
                });
        }

        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<bool> CheckForStashStep(int? pullCount)
        {
            CommitsToPull = pullCount ?? 0;
            return GitService.HasStash();
        }

        private IPromise<BranchData> FetchBranchDataStep(bool hasStash)
        {
            HasStash = hasStash;
            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}";
            Debug.LogException(ex);
        }
        
        private void FinishOperation()
        {
            IsLoading = false;
            _requestRepaint?.Invoke();
        }
        
        private void ClearMessages()
        {
            ErrorMessage = null;
            InfoMessage = null;
        }

        private void ResetMultipleFiles(List<string> filePaths)
        {
            EditorApplication.LockReloadAssemblies();
            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(_ =>
                {
                    InfoMessage = $"Successfully reset {filePaths.Count} file(s).";
                    Refresh();
                })
                .Catch(ex =>
                {
                    HandleOperationError(ex);
                    FinishOperation();
                })
                .Finally(() =>
                {
                    EditorApplication.UnlockReloadAssemblies();
                    AssetDatabase.Refresh();
                });
        }
        
        public void ResolveConflicts(List<GitChange> resolutions, IConflictSource source)
        {
            StartOperation("Resolving conflicts...");
            EditorApplication.LockReloadAssemblies();

            source.Resolve(resolutions)
                .Then(successMessage =>
                {
                    InfoMessage = successMessage;
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    EditorApplication.UnlockReloadAssemblies();
                    AssetDatabase.Refresh();
                    Refresh();
                });
        }
        
        public void ShowStashedChangesWindow()
        {
            var tempChanges = new List<GitChange>();
            tempChanges.AddRange(_changes);
            StartOperation("Loading stashed files...");
            GitService.GetStashedFiles()
                .Then(stashedFiles =>
                {
                    if (stashedFiles.Any())
                    {
                        StashedChangesWindow.ShowWindow(this, stashedFiles);
                    }
                    else
                    {
                        InfoMessage = "No files found in the stash.";
                    }
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    _changes = tempChanges;
                    FinishOperation();
                });
        }

        public void DiscardStash()
        {
            StartOperation("Discarding stash...");
            GitService.DropStash()
                .Then(successMessage =>
                {
                    InfoMessage = successMessage;
                })
                .Catch(HandleOperationError)
                .Finally(Refresh);
        }
        
        public void DiffStashedFile(GitChange change)
        {
            GitService.DiffStashedFile(change)
                .Catch(ex =>
                {
                    var res = _displayDialog("Stash Diff Error", $"Error diffing '{change.FilePath}': {ex.Message}", "Uh-oh", "Ok");
                    if (!res) return;
                    Debug.Log($"Error diffing '{change.FilePath}'");
                    Debug.LogException(ex);
                });
        }
        
        public void ApplyStash()
        {
            StartOperation("Analyzing stash for conflicts...");
            GitService.AnalyzeStashConflicts()
                .Then(analysisResult =>
                {
                    Debug.Log($"Any conflicts? {analysisResult.HasConflicts}");
                    FinishOperation();
                    if (analysisResult.HasConflicts)
                    {
                        
                        var conflictSource = new StashConflictSource();
                        GitService.GetStashedFiles().Then(stashedFiles =>
                        {
                            var conflictingChanges = stashedFiles
                                .Where(sf => analysisResult.ConflictingFiles.Contains(sf.FilePath))
                                .ToList();
                            Debug.Log($"Found {conflictingChanges.Count} conflicting files. Showing window.");
                            ConflictResolutionWindow.ShowWindow(this, conflictingChanges, conflictSource);
                        });
                    }
                    else
                    {
                        StartOperation("Applying stash...");
                        var resolutions = new List<GitChange>();
                        var source = new StashConflictSource();
                        source.Resolve(resolutions)
                            .Then(successMsg => { InfoMessage = successMsg; })
                            .Catch(HandleOperationError)
                            .Finally(Refresh);
                    }
                })
                .Catch(HandleOperationError)
                .Finally(() =>
                {
                    FinishOperation();
                    Refresh();
                    AssetDatabase.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;
            }
        }
    }
}