// GitAutoSync.cs - part of the Git Auto Sync package (net.lenrui.git-auto-sync).
//
// When Unity opens:  fetches, and if the branch's upstream has new commits, updates to them.
//                    Unpushed local commits are rebased on top; uncommitted files are stashed and restored.
// When Unity closes: saves scenes/assets, commits everything under the project folder,
//                    brings in anything pushed meanwhile, then pushes.
//                    If anything fails you can stay open or quit anyway; your work is always committed locally first.
//
// Menu: Tools > Git Auto Sync  (on/off, ask-before-commit, and manual Pull / Commit and Push)

using System;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace Lenrui.GitSync
{
    [InitializeOnLoad]
    public static class GitAutoSync
    {
        const string Title = "Git Auto Sync";
        const string MenuRoot = "Tools/Git Auto Sync/";
        const string StartupDoneKey = "GitAutoSync.StartupDone";

        static readonly string ProjectRoot = Path.GetFullPath(Path.Combine(Application.dataPath, ".."));
        static bool busy;

        // ---------------------------------------------------------------- settings (per project, per machine)

        static string Key(string name) { return "GitAutoSync." + name + "." + ProjectRoot; }

        static bool Enabled
        {
            get { return EditorPrefs.GetBool(Key("Enabled"), true); }
            set { EditorPrefs.SetBool(Key("Enabled"), value); }
        }

        static bool AskBeforeCommit
        {
            get { return EditorPrefs.GetBool(Key("AskBeforeCommit"), false); }
            set { EditorPrefs.SetBool(Key("AskBeforeCommit"), value); }
        }

        // ---------------------------------------------------------------- hooks

        static GitAutoSync()
        {
            if (Application.isBatchMode) return;

            EditorApplication.wantsToQuit -= OnWantsToQuit;
            EditorApplication.wantsToQuit += OnWantsToQuit;

            // Static constructors re-run on every script reload; only pull once per editor launch.
            if (SessionState.GetBool(StartupDoneKey, false)) return;
            SessionState.SetBool(StartupDoneKey, true);
            EditorApplication.delayCall += () => { if (Enabled) Guarded(() => PullLatest(false)); };
        }

        static bool OnWantsToQuit()
        {
            if (!Enabled || busy) return true;
            bool quit = true;
            Guarded(() => quit = CommitAndPush(true), onError: () =>
                quit = EditorUtility.DisplayDialog(Title, "Unexpected error while syncing; see the Console.", "Quit Anyway", "Stay Open"));
            return quit;
        }

        static void Guarded(Action action, Action onError = null)
        {
            busy = true;
            try { action(); }
            catch (Exception e) { Debug.LogException(e); if (onError != null) onError(); }
            finally { busy = false; EditorUtility.ClearProgressBar(); }
        }

        // ---------------------------------------------------------------- menu

        [MenuItem(MenuRoot + "Enabled", false, 0)]
        static void ToggleEnabled() { Enabled = !Enabled; }

        [MenuItem(MenuRoot + "Enabled", true)]
        static bool ToggleEnabledCheck() { Menu.SetChecked(MenuRoot + "Enabled", Enabled); return true; }

        [MenuItem(MenuRoot + "Ask Before Committing on Close", false, 1)]
        static void ToggleAsk() { AskBeforeCommit = !AskBeforeCommit; }

        [MenuItem(MenuRoot + "Ask Before Committing on Close", true)]
        static bool ToggleAskCheck() { Menu.SetChecked(MenuRoot + "Ask Before Committing on Close", AskBeforeCommit); return true; }

        [MenuItem(MenuRoot + "Pull Latest Now", false, 20)]
        static void PullNow() { Guarded(() => PullLatest(true)); }

        [MenuItem(MenuRoot + "Commit and Push Now", false, 21)]
        static void PushNow() { Guarded(() => CommitAndPush(false)); }

        // ---------------------------------------------------------------- open: fetch + update

        static void PullLatest(bool manual)
        {
            string why;
            if (!CheckRepo(out why)) { Info(why, manual); return; }

            string branch = CurrentBranch();
            if (branch == null) { Info("HEAD is detached (not on a branch), so nothing was pulled.", manual); return; }

            if (manual) AssetDatabase.SaveAssets();

            var fetch = Git("fetch --prune", "Checking for remote changes...", 120);
            if (!fetch.Ok) { Warn("Couldn't fetch (offline?). Working with the local copy.\n\n" + fetch.Tail, manual); return; }

            if (!HasUpstream(branch)) { Info("'" + branch + "' has no upstream yet, so there's nothing to pull. It will be pushed when you close Unity.", manual); return; }

            int ahead, behind;
            AheadBehind(out ahead, out behind);
            if (behind == 0) { Info("'" + branch + "' is up to date.", manual); return; }

            // Uncommitted edits to files that also changed remotely would end up full of conflict markers,
            // which breaks Unity assets. Leave them alone; they get committed on close and any clash is reported then.
            var remoteChanged = Lines(Git("diff --name-only HEAD...@{u}").Out);
            var localDirty = Lines(Git("diff --name-only HEAD").Out)
                .Concat(Lines(Git("ls-files --others --exclude-standard --full-name").Out));
            string[] overlap = localDirty.Intersect(remoteChanged).ToArray();
            if (overlap.Length > 0)
            {
                Alert("Didn't pull " + behind + " new commit(s) on '" + branch + "': you have uncommitted edits to files that were also changed remotely:\n\n" +
                      Clip(string.Join("\n", overlap), 10) + "\n\nThey'll be committed when you close Unity, and you'll be told if they conflict.");
                return;
            }

            string incoming = Git("log --oneline --no-decorate HEAD..@{u}").Out.Trim();
            int stashesBefore = StashCount();

            var rebase = Git("rebase --autostash @{u}", "Pulling " + behind + " new commit(s)...", 300);
            if (!rebase.Ok)
            {
                Git("rebase --abort");
                Alert("Couldn't update '" + branch + "': your " + ahead + " unpushed commit(s) conflict with the " + behind +
                      " new remote commit(s). Your files were left as they were; resolve it in your git client.\n\n" + rebase.Tail);
                return;
            }

            AssetDatabase.Refresh();
            Debug.Log("[GitAutoSync] Pulled " + behind + " commit(s) into '" + branch + "':\n" + incoming);

            if (StashCount() > stashesBefore)
                Alert("Pulled the latest changes, but your uncommitted edits clashed with them. Check for conflict markers; " +
                      "a copy of your edits is saved in `git stash`.");
            else if (manual)
                EditorUtility.DisplayDialog(Title, "Pulled " + behind + " commit(s):\n\n" + Clip(incoming, 15), "OK");
        }

        // ---------------------------------------------------------------- close: commit + push

        // Returns true if it's OK to quit (when quitting).
        static bool CommitAndPush(bool quitting)
        {
            string why;
            if (!CheckRepo(out why)) { Info(why, !quitting); return true; }

            // Make what's on disk match the editor.
            if (!EditorSceneManager.SaveCurrentModifiedScenesIfUserWantsTo()) return false; // user hit Cancel
            AssetDatabase.SaveAssets();

            // Never commit Unity's cache.
            if (Git("check-ignore -q Library/").Code != 0)
                return Fail("Library/ isn't in .gitignore, so auto-commit is off to avoid committing gigabytes of cache. " +
                            "Add a Unity .gitignore to the repo.", quitting);

            string branch = CurrentBranch();
            if (branch == null) return Fail("HEAD is detached (not on a branch), so nothing was committed or pushed.", quitting);

            // ---- commit
            string[] pending = Lines(Git("status --porcelain -- .").Out);
            if (pending.Length > 0)
            {
                if (quitting && AskBeforeCommit)
                {
                    int choice = EditorUtility.DisplayDialogComplex(Title,
                        "Commit and push " + pending.Length + " change(s) on '" + branch + "'?\n\n" + Clip(string.Join("\n", pending), 15),
                        "Commit and Push", "Stay Open", "Quit Without Committing");
                    if (choice == 1) return false;
                    if (choice == 2) return true;
                }

                var add = Git("add -A -- .", "Staging changes...", 300);
                if (!add.Ok) return Fail("git add failed:\n\n" + add.Tail, quitting);

                string[] staged = Lines(Git("diff --cached --name-status").Out);
                string message = "Auto-commit from Unity (" + Environment.MachineName + ", " + DateTime.Now.ToString("yyyy-MM-dd HH:mm") + ")\n\n" +
                                 Clip(string.Join("\n", staged), 100);

                string msgFile = Path.GetTempFileName();
                File.WriteAllText(msgFile, message, new UTF8Encoding(false));
                var commit = Git("commit -q -F \"" + msgFile + "\"", "Committing...", 300);
                File.Delete(msgFile);
                if (!commit.Ok) return Fail("Commit failed:\n\n" + commit.Tail, quitting);

                Debug.Log("[GitAutoSync] Committed " + staged.Length + " file(s) on '" + branch + "'.");
            }

            // ---- push
            bool hasUpstream = HasUpstream(branch);
            string remote = hasUpstream ? GitValue("config --get branch." + branch + ".remote") : DefaultRemote();
            if (remote == null) { Info("No git remote is set up, so changes were only committed locally.", !quitting); return true; }

            var fetch = Git("fetch --prune", "Checking for remote changes...", 120);
            if (!fetch.Ok)
                return Fail("Committed locally, but couldn't reach the remote, so nothing was pushed. " +
                            "It will be pushed the next time you close Unity.\n\n" + fetch.Tail, quitting);

            GitResult push;
            if (hasUpstream)
            {
                int ahead, behind;
                AheadBehind(out ahead, out behind);
                if (behind > 0)
                {
                    var rebase = Git("rebase --autostash @{u}", "Bringing in " + behind + " remote commit(s)...", 300);
                    if (!rebase.Ok)
                    {
                        Git("rebase --abort");
                        return Fail("New commits on the remote conflict with yours. Your work is committed locally but not pushed; " +
                                    "resolve it in your git client, then push.\n\n" + rebase.Tail, quitting);
                    }
                    AheadBehind(out ahead, out behind);
                }

                if (ahead == 0) { Info("'" + branch + "' is already in sync.", !quitting); return true; }

                string mergeRef = GitValue("config --get branch." + branch + ".merge");
                push = Git("push " + remote + " HEAD:" + mergeRef, "Pushing " + ahead + " commit(s)...", 900);
            }
            else
            {
                push = Git("push -u " + remote + " HEAD", "Pushing new branch '" + branch + "'...", 900);
            }

            if (!push.Ok)
                return Fail("Committed locally, but the push failed. It will be retried the next time you close Unity.\n\n" + push.Tail, quitting);

            Info("Pushed '" + branch + "' to " + remote + ".", !quitting);
            if (!quitting) AssetDatabase.Refresh();
            return true;
        }

        // ---------------------------------------------------------------- git helpers

        static bool CheckRepo(out string why)
        {
            var r = Git("rev-parse --is-inside-work-tree");
            why = null;
            if (r.Code == -1) why = r.Err;
            else if (!r.Ok || r.Out.Trim() != "true") why = "This project isn't inside a git repository, so Git Auto Sync is idle.";
            return why == null;
        }

        static string CurrentBranch() { return GitValue("symbolic-ref --short -q HEAD"); }

        static bool HasUpstream(string branch)
        {
            return GitValue("config --get branch." + branch + ".merge") != null && Git("rev-parse -q --verify @{u}").Ok;
        }

        static string DefaultRemote()
        {
            string[] remotes = Lines(Git("remote").Out);
            if (remotes.Contains("origin")) return "origin";
            return remotes.Length > 0 ? remotes[0] : null;
        }

        static void AheadBehind(out int ahead, out int behind)
        {
            string[] parts = Git("rev-list --left-right --count HEAD...@{u}").Out.Split(new[] { '\t', ' ', '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries);
            ahead = parts.Length > 0 ? int.Parse(parts[0]) : 0;
            behind = parts.Length > 1 ? int.Parse(parts[1]) : 0;
        }

        static int StashCount() { return Lines(Git("stash list").Out).Length; }

        static string GitValue(string args)
        {
            var r = Git(args);
            string v = r.Out.Trim();
            return r.Ok && v.Length > 0 ? v : null;
        }

        struct GitResult
        {
            public int Code;
            public string Out, Err;
            public bool Ok { get { return Code == 0; } }
            public string Tail
            {
                get
                {
                    var all = Lines((Err ?? "") + "\n" + (Out ?? ""))
                        .Where(l => !l.StartsWith("hint:") && !l.StartsWith("Rebasing (")).ToArray();
                    return string.Join("\n", all.Skip(Math.Max(0, all.Length - 12)).ToArray());
                }
            }
        }

        // Runs git in the project folder. With a progress title it shows a cancelable progress bar.
        static GitResult Git(string args, string progress = null, int timeoutSeconds = 30)
        {
            var psi = new ProcessStartInfo("git", args)
            {
                WorkingDirectory = ProjectRoot,
                UseShellExecute = false,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                CreateNoWindow = true,
                StandardOutputEncoding = Encoding.UTF8,
                StandardErrorEncoding = Encoding.UTF8,
            };
            psi.EnvironmentVariables["GIT_TERMINAL_PROMPT"] = "0"; // fail instead of waiting on an invisible password prompt

            var output = new StringBuilder();
            var errors = new StringBuilder();
            using (var p = new Process { StartInfo = psi })
            {
                p.OutputDataReceived += (s, e) => { if (e.Data != null) lock (output) output.AppendLine(e.Data); };
                p.ErrorDataReceived += (s, e) => { if (e.Data != null) lock (errors) errors.AppendLine(e.Data); };

                try { p.Start(); }
                catch (Exception e)
                {
                    return new GitResult { Code = -1, Out = "", Err = "Couldn't run git (is it installed and on PATH?): " + e.Message };
                }
                p.BeginOutputReadLine();
                p.BeginErrorReadLine();

                var timer = Stopwatch.StartNew();
                string abortReason = null;
                while (!p.WaitForExit(100))
                {
                    if (progress != null &&
                        EditorUtility.DisplayCancelableProgressBar(Title, progress + "  (" + (int)timer.Elapsed.TotalSeconds + "s)",
                                                                   (float)(timer.Elapsed.TotalSeconds % 4.0 / 4.0)))
                        abortReason = "Cancelled.";
                    else if (timer.Elapsed.TotalSeconds > timeoutSeconds)
                        abortReason = "Timed out after " + timeoutSeconds + "s.";

                    if (abortReason != null)
                    {
                        try { p.Kill(); } catch { }
                        break;
                    }
                }
                p.WaitForExit(); // flush async output
                if (progress != null) EditorUtility.ClearProgressBar();

                string err;
                lock (errors) err = errors.ToString();
                if (abortReason != null) err += "\n" + abortReason;
                lock (output) return new GitResult { Code = abortReason != null ? -2 : p.ExitCode, Out = output.ToString(), Err = err };
            }
        }

        // ---------------------------------------------------------------- reporting

        static void Info(string msg, bool showDialog)
        {
            Debug.Log("[GitAutoSync] " + msg);
            if (showDialog) EditorUtility.DisplayDialog(Title, msg, "OK");
        }

        static void Warn(string msg, bool showDialog)
        {
            Debug.LogWarning("[GitAutoSync] " + msg);
            if (showDialog) EditorUtility.DisplayDialog(Title, msg, "OK");
        }

        static void Alert(string msg) { Warn(msg, true); }

        // When quitting: ask whether to quit anyway. Otherwise: just report. Returns true to quit.
        static bool Fail(string msg, bool quitting)
        {
            Debug.LogWarning("[GitAutoSync] " + msg);
            if (!quitting)
            {
                EditorUtility.DisplayDialog(Title, msg, "OK");
                AssetDatabase.Refresh();
                return false;
            }
            bool quit = EditorUtility.DisplayDialog(Title, msg, "Quit Anyway", "Stay Open");
            if (!quit) AssetDatabase.Refresh(); // a partial sync may have changed files on disk
            return quit;
        }

        static string[] Lines(string s)
        {
            return (s ?? "").Split(new[] { '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries);
        }

        static string Clip(string text, int maxLines)
        {
            string[] lines = Lines(text);
            if (lines.Length <= maxLines) return string.Join("\n", lines);
            return string.Join("\n", lines.Take(maxLines).ToArray()) + "\n...and " + (lines.Length - maxLines) + " more";
        }
    }
}