i3-session-manager/i3/session.go

package i3

import (
	// Standard
	"encoding/json"
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var (
	outputIndices map[string]int
	workspaces    []I3Workspace
)

func init() {
	outputIndices = make(map[string]int)
}

func SetWorkspaces(wss []I3Workspace) {
	workspaces = wss
}

func (sess *I3Session) Open() (err error) {
	// Find path to i3 socket
	i3SocketPathCmd := exec.Command("i3", "--get-socketpath")
	i3SocketPathBytes, _ := i3SocketPathCmd.CombinedOutput()
	i3SocketPath := strings.TrimSpace(string(i3SocketPathBytes))
	sess.Socket, err = NewI3Socket(i3SocketPath)
	return
}

// Close closes the socket.
func (sess I3Session) Close() error {
	return sess.Socket.Close()
}

// Outputs query I3 for output data.
func (sess I3Session) Outputs() (outputs []I3Output, err error) {
	var outputJson []byte

	outputJson, err = sess.Socket.Request(GET_OUTPUTS, "")
	if err != nil {
		return
	}

	err = json.Unmarshal(outputJson, &outputs)
	return
}

// UpdateOutputIndices counts the outputs and assign incdices to them.
// The indices are used to calculate workspace IDs.
func (sess I3Session) UpdateOutputIndices() error {
	outputIndices = make(map[string]int)
	fmt.Printf("Update output indices:\n")
	defer fmt.Printf("\n")

	outputs, err := sess.Outputs()
	if err != nil {
		return err
	}
	idx := 0
	for _, output := range outputs {
		if output.Active {
			fmt.Printf(" %2d %s\n", idx, output.Name)
			outputIndices[output.Name] = idx
			idx++
		}
	}
	return nil
}

// Marks query I3 for mark data.
func (sess I3Session) Marks() (marks []string, err error) {
	var outputJson []byte

	outputJson, err = sess.Socket.Request(GET_MARKS, "")
	if err != nil {
		return
	}

	err = json.Unmarshal(outputJson, &marks)
	return
}

// Workspaces query I3 for workspace data.
func (sess I3Session) Workspaces() (workspaces []I3Workspace, err error) {
	var workspacesJson []byte
	workspacesJson, err = sess.Socket.Request(GET_WORKSPACES, "")
	if err != nil {
		return
	}

	err = json.Unmarshal(workspacesJson, &workspaces)
	return
}

// Subscribe subscribes to events.
func (sess I3Session) Subscribe(events []string) (err error) {
	eventsJSON, _ := json.Marshal(events)
	_, err = sess.Socket.Request(SUBSCRIBE, string(eventsJSON))
	return
}

// Config returns the currently running configuration data, as
// written in the configuration file.
func (sess I3Session) Config() (config string, err error) {
	var configJSON []byte
	var configStruct I3Config
	configJSON, err = sess.Socket.Request(GET_CONFIG, "")
	err = json.Unmarshal(configJSON, &configStruct)
	if err != nil {
		return
	}

	config = configStruct.Config
	return
}

// SwitchWorkspace switches to the given workspace index on the same output
// that has the workspace that is currently focused.
func (sess I3Session) SwitchWorkspace(workspaceIndex int) error {
	var output string
	i3Workspaces, err := sess.Workspaces()
	if err != nil {
		return err
	}

	for _, ws := range i3Workspaces {
		if ws.Focused {
			output = ws.Output
			break
		}
	}

	if output == "" {
		return fmt.Errorf("Focused workspace not found")
	}

	outputIdx := outputIndices[output]
	workspaceName := fmt.Sprintf(
		"%d:%s",
		outputIdx*10+workspaceIndex,
		workspaces[workspaceIndex].Name,
	)
	fmt.Printf("Switch to workspace: %s\n", workspaceName)
	sess.Socket.Request(
		RUN_COMMAND,
		fmt.Sprintf("workspace \"%s\"", workspaceName),
	)

	return nil
}

// SwitchOutput looks up the visible workspace on the given monitor,
// and switches to that.
func (sess I3Session) SwitchOutput(output string) error {
	workspaces, err := sess.Workspaces()
	if err != nil {
		return err
	}

	var nextWorkspace string
	for _, workspace := range workspaces {
		if workspace.Output == output && workspace.Visible {
			nextWorkspace = workspace.Name
			break
		}
	}

	if nextWorkspace == "" {
		return fmt.Errorf(
			"Switch to output (%s) failed, "+
				"no visible workspace found",
			output,
		)
	}

	fmt.Printf("Switching to output: %s\n", output)
	sess.Socket.Request(RUN_COMMAND, "workspace "+nextWorkspace)

	return nil
}

// ManageSession runs dmenu and asks for reload, restart or exit.
func (sess *I3Session) ManageSession() error {
	cmd := exec.Command(
		"sh", "-c",
		"echo \"reload\nrestart\nexit\" | dmenu -l 3 -i",
	)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return err
	}
	choice := strings.TrimSpace(string(out))

	switch choice {
	case "reload", "restart", "exit":
		_, err = sess.Socket.Request(RUN_COMMAND, choice)
		if err != nil {
			return err
		}

		if choice == "restart" {
			return fmt.Errorf("restarted")
		}
	default:
		return fmt.Errorf("Unknown management command")
	}

	return nil
}

// MarkTag marks a window with a tag for moving
func (sess I3Session) MarkTag() error {
	var max int
	var cur int
	var match []string
	r := regexp.MustCompile(`^move (\d+)$`)

	marks, err := sess.Marks()
	if err != nil {
		return err
	}

	for _, mark := range marks {
		match = r.FindStringSubmatch(mark)
		if len(match) == 2 {
			cur, _ = strconv.Atoi(match[1])
			if cur > max {
				max = cur
			}
		}
	}

	res, err := sess.Socket.Request(
		RUN_COMMAND,
		fmt.Sprintf("mark \"move %d\"", max+1),
	)
	if err != nil {
		return err
	}
	fmt.Printf("Mark tag: %s\n", res)

	return nil
}

// MarkTag marks a window with a tag for moving
func (sess I3Session) MarkClear() error {
	marks, err := sess.Marks()
	if err != nil {
		return err
	}

	r := regexp.MustCompile(`^move \d+$`)
	for _, mark := range marks {
		if r.Match([]byte(mark)) {
			_, err := sess.Socket.Request(
				RUN_COMMAND,
				fmt.Sprintf("unmark \"%s\"", mark),
			)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// MarkMove moves tagged windows to current workspace
func (sess I3Session) MarkMove() error {
	i3Workspaces, err := sess.Workspaces()
	if err != nil {
		return err
	}

	var focusedWorkspace *I3Workspace
	for _, ws := range i3Workspaces {
		if ws.Focused {
			focusedWorkspace = &ws
			break
		}
	}

	if focusedWorkspace == nil {
		return fmt.Errorf("No focused workspace")
	}

	cmd := fmt.Sprintf(
		"[con_mark=\"^move [0-9]\"] move workspace \"%s\"",
		focusedWorkspace.Name,
	)
	res, err := sess.Socket.Request(RUN_COMMAND, cmd)
	if err != nil {
		return err
	}

	fmt.Printf("Mark move: %s\n", res)

	err = sess.MarkClear()
	if err != nil {
		return err
	}

	return nil
}

// StoreWorkspaces stores the output of each workspace.
func (sess *I3Session) StoreWorkspaces() (err error) {
	var workspaces []I3Workspace

	sess.workspaces = make(map[string]string)
	sess.workspacesActive = make(map[string]string)

	workspaces, err = sess.Workspaces()
	if err != nil {
		return
	}

	for _, ws := range workspaces {
		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
}

// RestoreWorkspaces moves the workspace back to its stored place, if possible.
func (sess *I3Session) RestoreWorkspaces() (err error) {
	err = sess.UpdateOutputIndices()
	if err != nil {
		return err
	}

	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=\"%s\"] move workspace to output \"%s\"",
			wsName,
			output,
		)
		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
}

// FixWorkspaces moves windows from workspaces on the wrong monitor to
// to workspaces on the correct monitor. For fixing when going to less
// outputs.
func (sess I3Session) FixWorkspaces() error {
	// Outputs could be completely different, and could need to be updated.
	err := sess.UpdateOutputIndices()
	if err != nil {
		return err
	}

	// Identify workspaces on the wrong output
	wss, err := sess.Workspaces()
	if err != nil {
		return err
	}

	var outputIdx int
	var found bool
	for _, ws := range wss {
		outputIdx, found = outputIndices[ws.Output]
		outputIdx *= 10
		if !found {
			return fmt.Errorf("Output %s not found", ws.Output)
		}

		// 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",
				)
			}

			cmd := fmt.Sprintf(
				"[workspace=\"%s\"] move window to workspace "+
					"\"%d:%s\"",
				ws.Name,
				outputIdx+(ws.Num%10),
				workspaces[(ws.Num%10)].Name,
			)
			fmt.Printf("Move windows: %s\n", cmd)
			_, err := sess.Socket.Request(RUN_COMMAND, cmd)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// Loop reads the subscription socket and reacts to the implemented messages.
func (sess *I3Session) Loop() error {
	var msg []byte
	var err error
	var cmdBinding I3EventBinding

	cmdOutput := regexp.MustCompile(`(?i)^\s*nop\s*output\s*(.*?)\s*$`)
	cmdWS := regexp.MustCompile(`(?i)^\s*nop\s*workspace\s*(\d+)\s*$`)
	cmdSession := regexp.MustCompile(`(?i)^\s*nop\s*manage\s+session`)
	cmdMarkTag := regexp.MustCompile(`(?i)^\s*nop\s*mark\s+tag`)
	cmdMarkMove := regexp.MustCompile(`(?i)^\s*nop\s*mark\s+move`)
	cmdMarkClear := regexp.MustCompile(`(?i)^\s*nop\s*mark\s+clear`)

	for {
		msg, err = I3ReadMessage(sess.Socket)
		if err != nil {
			return err
		}

		cmdBinding = I3EventBinding{}
		err = json.Unmarshal(msg, &cmdBinding)
		if err != nil {
			fmt.Printf(
				"\nLoop JSON parse: %s\nMessage:\n(%s)\n\n",
				err,
				msg,
			)
		}
		binding := cmdBinding.Binding

		// Process is restarted
		if cmdBinding.Change == "restart" {
			return fmt.Errorf("restarted")
		}

		// Switch output
		m := cmdOutput.FindAllStringSubmatch(binding.Command, 1)
		if len(m) == 1 && len(m[0]) == 2 {
			err = sess.SwitchOutput(m[0][1])
			if err != nil {
				fmt.Printf("%s\n", err)
			}
			continue
		}

		// Switch output
		m = cmdWS.FindAllStringSubmatch(binding.Command, 1)
		if len(m) == 1 && len(m[0]) == 2 {
			i, _ := strconv.Atoi(m[0][1])
			err = sess.SwitchWorkspace(i)
			if err != nil {
				fmt.Printf("%s\n", err)
			}
			continue
		}

		// Manage session
		if cmdSession.Match([]byte(binding.Command)) {
			err = sess.ManageSession()
			if err != nil {
				return err
			}
			continue
		}

		// Mark current window with move tag
		if cmdMarkTag.Match([]byte(binding.Command)) {
			err = sess.MarkTag()
			if err != nil {
				fmt.Printf("Mark tag: %s\n", err)
			}
			continue
		}

		// Remove all move tags
		if cmdMarkClear.Match([]byte(binding.Command)) {
			err = sess.MarkClear()
			if err != nil {
				fmt.Printf("Mark clear: %s\n", err)
			}
			continue
		}

		// Move all windows with move tag
		if cmdMarkMove.Match([]byte(binding.Command)) {
			err = sess.MarkMove()
			if err != nil {
				fmt.Printf("Mark move: %s\n", err)
			}
			continue
		}

		//fmt.Printf("%s\n", msg)
	}
}

// vim: foldmethod=syntax foldnestmax=1