// 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 _changes = new(); public IReadOnlyList 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 RemoteBranchList { get; private set; } = new List(); public string CurrentBranchName { get; private set; } = "master"; private readonly Action _requestRepaint; private readonly Func _displayDialog; private readonly Func _promptForUnsavedChanges; /// /// Initializes the controller. /// /// A callback to the window's Repaint() method. /// A callback to EditorUtility.DisplayDialog for user confirmations. /// A callback to EditorUtility.DisplayDialog for user confirmations. public ArbitratorController(Action requestRepaint, Func displayDialog, Func 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(); // Stop loading before showing dialog if (analysisResult.HasConflicts) { ConflictResolutionWindow.ShowWindow(this, analysisResult.ConflictingFiles); } else { PerformSafePullWithLock(); } }) .Catch(ex => { HandleOperationError(ex); FinishOperation(); }) .Finally(FinishOperation); } public void PerformSafePullWithLock() { StartOperation("Pulling changes..."); EditorApplication.LockReloadAssemblies(); GitService.PerformSafePull() .Then(successMessage => { InfoMessage = successMessage; Refresh(); }) .Catch(ex => { HandleOperationError(ex); FinishOperation(); }) .Finally(EditorApplication.UnlockReloadAssemblies); } public void ResetSingleConflictingFile(string filePath, Action onResetComplete) { var change = GitService.GetChangeForFile(filePath); if (change == null) { ErrorMessage = $"Could not find file '{filePath}' to reset."; onResetComplete?.Invoke(); return; } GitService.ResetFileChanges(change) .Then(successMessage => { InfoMessage = successMessage; }) .Catch(HandleOperationError) .Finally(onResetComplete); } public void CommitAndPush(string commitMessage) { 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 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 UnstageStep() { return GitService.UnstageAllFilesIfSafe(); } private IPromise> CompareStep(bool wasUnstaged) { if (wasUnstaged) { InfoMessage = "Found and unstaged files for review."; } return GitService.CompareLocalToRemote(); } private IPromise FetchUpstreamStep(List changes) { _changes = changes; IsInConflictState = _changes.Any(c => c.Status == LibGit2Sharp.ChangeKind.Conflicted); return IsInConflictState ? new Promise((resolve, _) => resolve(0)) : GitService.GetUpstreamAheadBy(OnProgressModified); void OnProgressModified(float progress, string message) { OperationProgress = progress; OperationProgressMessage = message; _requestRepaint?.Invoke(); } } private IPromise 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 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 promiseChain = new Promise((resolve, _) => resolve("")); foreach (var path in filePaths) { promiseChain = promiseChain.Then(_ => { var change = GitService.GetChangeForFile(path); return change != null ? GitService.ResetFileChanges(change) : new Promise((res, _) => res("")); }); } promiseChain .Then(successMsg => { InfoMessage = $"Successfully reset {filePaths.Count} file(s)."; Refresh(); }) .Catch(ex => { HandleOperationError(ex); FinishOperation(); }); } } private 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; } } } }