vmux/i3bar.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

// I3BarBlock represents a single block in the i3bar protocol.
type I3BarBlock struct {
	FullText  string `json:"full_text"`
	ShortText string `json:"short_text,omitempty"`
	Color     string `json:"color"`
	Name      string `json:"name"`
	Markup    string `json:"markup,omitempty"`
}

// formatI3BarBlocks builds the vmux status block for i3bar from session data.
func formatI3BarBlocks(sessions []SessionInfo) []I3BarBlock {
	if len(sessions) == 0 {
		return []I3BarBlock{{FullText: "vmux: no sessions", Color: "#00ff00", Name: "vmux"}}
	}

	hasWaiting := false
	parts := make([]string, 0, len(sessions))

	for _, s := range sessions {
		name := shortName(s)
		switch s.State {
		case "Needs Input":
			parts = append(parts, name+"[!]")
			hasWaiting = true
		case "Working":
			parts = append(parts, name+"[W]")
		case "Idle":
			parts = append(parts, name+"[I]")
		}
	}

	var text string
	if !hasWaiting {
		text = fmt.Sprintf("vmux: all working (%d)", len(sessions))
	} else {
		text = "vmux: " + strings.Join(parts, " ")
	}

	color := "#00ff00"
	if hasWaiting {
		color = "#ff0000"
	}

	return []I3BarBlock{{FullText: text, Color: color, Name: "vmux"}}
}

// queryVmuxBlock fetches sessions from daemon and returns the vmux i3bar block.
// Returns a "daemon offline" block if the daemon is unreachable.
func queryVmuxBlock(sockPath string) I3BarBlock {
	client := NewClient(sockPath)
	resp, err := client.Send(Request{Action: "list"})
	if err != nil {
		return I3BarBlock{FullText: "vmux: daemon offline", Color: "#888888", Name: "vmux"}
	}
	if !resp.OK {
		return I3BarBlock{FullText: "vmux: error", Color: "#888888", Name: "vmux"}
	}
	blocks := formatI3BarBlocks(resp.Sessions)
	return blocks[0]
}

// writeI3BarLine writes a JSON array of blocks as an i3bar protocol line.
// Uses os.Stdout.Write + Sync to avoid buffering issues.
func writeI3BarLine(blocks []I3BarBlock, first bool) {
	data, err := json.Marshal(blocks)
	if err != nil {
		log.Printf("i3bar marshal error: %v", err)
		return
	}

	var line string
	if first {
		line = string(data) + "\n"
	} else {
		line = "," + string(data) + "\n"
	}
	os.Stdout.Write([]byte(line))
	os.Stdout.Sync()
}

// runI3Bar runs the i3bar protocol loop on stdout.
// If i3statusCmd is non-empty, it wraps i3status output and prepends vmux blocks.
// If empty, runs in standalone mode with periodic polling.
func runI3Bar(sockPath, i3statusCmd string) {
	if i3statusCmd != "" {
		runI3BarWrapped(sockPath, i3statusCmd)
		// If i3status exits, fall through to standalone
	}
	runI3BarStandalone(sockPath)
}

func runI3BarStandalone(sockPath string) {
	os.Stdout.Write([]byte("{\"version\":1}\n"))
	os.Stdout.Write([]byte("[\n"))
	os.Stdout.Sync()

	first := true
	for {
		block := queryVmuxBlock(sockPath)
		writeI3BarLine([]I3BarBlock{block}, first)
		first = false
		time.Sleep(2 * time.Second)
	}
}

func runI3BarWrapped(sockPath, i3statusCmd string) {
	cmd := exec.Command("sh", "-c", i3statusCmd)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("i3bar: cannot pipe i3status stdout: %v", err)
		return
	}
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		log.Printf("i3bar: cannot start i3status: %v", err)
		return
	}

	scanner := bufio.NewScanner(stdout)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

	// Forward header line (e.g. {"version":1})
	if !scanner.Scan() {
		log.Print("i3bar: i3status produced no output")
		return
	}
	os.Stdout.Write([]byte(scanner.Text() + "\n"))
	os.Stdout.Sync()

	// Forward opening bracket
	if !scanner.Scan() {
		log.Print("i3bar: i3status produced no opening bracket")
		return
	}
	os.Stdout.Write([]byte(scanner.Text() + "\n"))
	os.Stdout.Sync()

	first := true
	for scanner.Scan() {
		line := scanner.Text()
		trimmed := strings.TrimLeft(line, ",")
		if trimmed == "" {
			continue
		}

		var blocks []I3BarBlock
		if err := json.Unmarshal([]byte(trimmed), &blocks); err != nil {
			// Not a JSON array line, forward as-is
			os.Stdout.Write([]byte(line + "\n"))
			os.Stdout.Sync()
			continue
		}

		vmuxBlock := queryVmuxBlock(sockPath)
		blocks = append([]I3BarBlock{vmuxBlock}, blocks...)

		writeI3BarLine(blocks, first)
		first = false
	}

	cmd.Wait()
}