From f9e815b4e315ccc9e8d13ca81b1fc3a1d78d3e60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Magnus=20=C3=85hall?= <magnus.ahall@gibon.se>
Date: Tue, 28 Dec 2021 10:57:18 +0100
Subject: [PATCH] Better handling of store/restore workspaces

---
 session.go | 87 +++++++++++++++++++++++++++++++++++++++++++++++++-----
 types.go   |  4 ++-
 2 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/session.go b/session.go
index 3299778..134fd99 100644
--- a/session.go
+++ b/session.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os/exec"
 	"regexp"
+	"sort"
 	"strconv"
 	"strings"
 )
@@ -299,14 +300,29 @@ func (sess I3Session) MarkMove() error {
 func (sess *I3Session) StoreWorkspaces() (err error) {
 	var workspaces []I3Workspace
 
-	sess.workspaces = make(map[int]string)
+	sess.workspaces = make(map[string]string)
+	sess.workspacesActive = make(map[string]string)
+
 	workspaces, err = sess.Workspaces()
 	if err != nil {
 		return
 	}
 
 	for _, ws := range workspaces {
-		sess.workspaces[ws.Id] = ws.Output
+		fmt.Printf("Store workspace: %s - %s\n", ws.Name, ws.Output)
+		sess.workspaces[ws.Name] = ws.Output
+
+		// Visible workspaces are stored to visit them at a restore,
+		// to make sure they still are visible afterwards.
+		if ws.Visible {
+			sess.workspacesActive[ws.Name] = ws.Output
+		}
+
+		// Focused workspace is switched to last, to ensure it is still
+		// the focused workspace.
+		if ws.Focused {
+			sess.activeWorkspace = ws.Name
+		}
 	}
 	return
 }
@@ -318,18 +334,71 @@ func (sess *I3Session) RestoreWorkspaces() (err error) {
 		return err
 	}
 
-	for wsId, output := range sess.workspaces {
+	var list [][]string
+	var wsNum int
+
+
+	for wsName, output := range sess.workspaces {
+		//fmt.Printf("Restore workspace: %s to %s\n", wsName, output)
+		components := strings.Split(wsName, ":")
+		wsNum, err = strconv.Atoi(components[0])
+		if err  != nil {
+			return err
+		}
+
 		cmd := fmt.Sprintf(
-			"[workspace=%d] move workspace to output %s",
-			wsId,
+			"[workspace=\"%s\"] move workspace to output \"%s\"",
+			wsName,
 			output,
 		)
-		res, err := sess.Socket.Request(RUN_COMMAND, cmd)
+		list = append(list, []string{
+			fmt.Sprintf("%03d", wsNum),
+			cmd,
+		})
+
+		if false {
+		}
+	}
+
+	sort.SliceStable(list, func(i, j int) bool {
+		return list[i][0] < list[j][0]
+	})
+
+	for _, cmd := range list {
+		fmt.Printf("%s: ", cmd)
+		res, err := sess.Socket.Request(RUN_COMMAND, cmd[1])
 		if err != nil {
 			return err
 		}
 		fmt.Printf("%s\n", res)
 	}
+
+	// Visible workspaces are switched to to make sure they still are
+	// visible afterwards.
+	for wsName := range sess.workspacesActive {
+		cmd := fmt.Sprintf(
+			"workspace \"%s\"",
+			wsName,
+		)
+		_, err = sess.Socket.Request(RUN_COMMAND, cmd)
+		if err != nil {
+			return err
+		}
+	}
+
+	// i3 will often leave user on workspace "1", going to last workspace
+	// when stored makes for a better user experience.
+	if sess.activeWorkspace != "" {
+		cmd := fmt.Sprintf(
+			"workspace \"%s\"",
+			sess.activeWorkspace,
+		)
+		_, err = sess.Socket.Request(RUN_COMMAND, cmd)
+		if err != nil {
+			return err
+		}
+	}
+
 	return
 }
 
@@ -361,6 +430,7 @@ func (sess I3Session) FixWorkspaces() error {
 		// Is workspace on the wrong output?
 		wsOutputIdx := ws.Num - (ws.Num % 10)
 		if wsOutputIdx != outputIdx {
+			fmt.Printf("%s is on wrong output\n", ws.Name)
 			if (ws.Num % 10) >= len(workspaces) {
 				return fmt.Errorf(
 					"Workspace index doesn't exist",
@@ -368,9 +438,10 @@ func (sess I3Session) FixWorkspaces() error {
 			}
 
 			cmd := fmt.Sprintf(
-				"[workspace=\"%s\"] move window to workspace \"%d:%s\"",
+				"[workspace=\"%s\"] move window to workspace "+
+				"\"%d:%s\"",
 				ws.Name,
-				(outputIdx*10)+(ws.Num % 10),
+				outputIdx+(ws.Num % 10),
 				workspaces[(ws.Num % 10)].Name,
 			)
 			_, err := sess.Socket.Request(RUN_COMMAND, cmd)
diff --git a/types.go b/types.go
index 147bd46..8db8ece 100644
--- a/types.go
+++ b/types.go
@@ -10,7 +10,9 @@ type I3Session struct {
 	filename string
 
 	// workspace ID → output name
-	workspaces map[int]string
+	workspaces       map[string]string
+	workspacesActive map[string]string
+	activeWorkspace string
 }
 
 type I3Socket struct {