vmux/main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	i3 "go.i3wm.org/i3/v4"
)

func main() {
	home := os.Getenv("HOME")
	sockPath := filepath.Join(home, ".vmux", "vmux.sock")

	noColor := os.Getenv("NO_COLOR") != ""

	args := os.Args[1:]

	// Filter --no-color from args
	var filteredArgs []string
	for _, arg := range args {
		if arg == "--no-color" {
			noColor = true
		} else {
			filteredArgs = append(filteredArgs, arg)
		}
	}

	if len(filteredArgs) == 0 {
		printUsage()
		os.Exit(1)
	}

	switch filteredArgs[0] {
	case "list":
		if err := EnsureDaemon(sockPath); err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		client := NewClient(sockPath)
		resp, err := client.Send(Request{Action: "list"})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		if !resp.OK {
			fmt.Fprintf(os.Stderr, "Error: %s\n", resp.Error)
			os.Exit(1)
		}
		DisplaySessionInfos(os.Stdout, resp.Sessions, noColor, time.Now())

	case "switch":
		if len(filteredArgs) < 2 {
			fmt.Fprintf(os.Stderr, "Usage: vmux switch <query>\n")
			os.Exit(1)
		}
		query := strings.Join(filteredArgs[1:], " ")
		if err := EnsureDaemon(sockPath); err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		client := NewClient(sockPath)
		switchArgs, _ := json.Marshal(SwitchArgs{Query: query})
		resp, err := client.Send(Request{Action: "switch", Args: switchArgs})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		if !resp.OK {
			fmt.Fprintf(os.Stderr, "Error: %s\n", resp.Error)
			os.Exit(1)
		}
		fmt.Println("Switched.")

	case "label":
		if len(filteredArgs) < 3 {
			fmt.Fprintf(os.Stderr, "Usage: vmux label <session-id> <text...>\n")
			os.Exit(1)
		}
		sessionID := filteredArgs[1]
		label := strings.Join(filteredArgs[2:], " ")
		if err := EnsureDaemon(sockPath); err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		client := NewClient(sockPath)
		labelArgs, _ := json.Marshal(LabelArgs{SessionID: sessionID, Label: label})
		resp, err := client.Send(Request{Action: "label", Args: labelArgs})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		if !resp.OK {
			fmt.Fprintf(os.Stderr, "Error: %s\n", resp.Error)
			os.Exit(1)
		}
		fmt.Println("Label set.")

	case "focus":
		if len(filteredArgs) < 2 {
			fmt.Fprintf(os.Stderr, "Usage: vmux focus <minutes>\n")
			os.Exit(1)
		}
		minutes, err := strconv.Atoi(filteredArgs[1])
		if err != nil || minutes < 0 {
			fmt.Fprintf(os.Stderr, "Error: invalid minutes %q\n", filteredArgs[1])
			os.Exit(1)
		}
		if err := EnsureDaemon(sockPath); err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		client := NewClient(sockPath)
		focusArgs, _ := json.Marshal(FocusArgs{Minutes: minutes})
		resp, err := client.Send(Request{Action: "focus", Args: focusArgs})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		if !resp.OK {
			fmt.Fprintf(os.Stderr, "Error: %s\n", resp.Error)
			os.Exit(1)
		}
		fmt.Printf("Focus mode: notifications suppressed for %d minutes\n", minutes)

	case "stop":
		client := NewClient(sockPath)
		resp, err := client.Send(Request{Action: "stop"})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		if !resp.OK {
			fmt.Fprintf(os.Stderr, "Error: %s\n", resp.Error)
			os.Exit(1)
		}
		fmt.Println("Daemon stopped.")

	case "i3bar":
		if err := EnsureDaemon(sockPath); err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}
		runI3Bar(sockPath)

	case "daemon":
		runDaemon(sockPath)

	default:
		printUsage()
		os.Exit(1)
	}
}

func runDaemon(sockPath string) {
	home := os.Getenv("HOME")
	procDir := "/proc"
	claudeDir := filepath.Join(home, ".claude", "projects")
	labelsPath := filepath.Join(home, ".vmux", "labels.json")

	if err := os.MkdirAll(filepath.Dir(sockPath), 0o755); err != nil {
		log.Fatalf("create vmux dir: %v", err)
	}

	labels, err := NewLabelStore(labelsPath)
	if err != nil {
		log.Fatalf("load labels: %v", err)
	}

	d := NewDaemon(sockPath, procDir, claudeDir, labels)

	// Wire i3 + X11 workspace resolution (graceful degradation)
	x11, err := NewRealX11Resolver()
	if err != nil {
		log.Printf("Warning: X11 unavailable, no workspace resolution: %v", err)
	} else {
		d.i3commander = RealI3Commander{}
		d.InitWorkspaceResolver(&realI3TreeProvider{}, x11)
	}

	if err := d.Start(); err != nil {
		log.Fatalf("start daemon: %v", err)
	}

	log.Printf("vmux daemon listening on %s", sockPath)
	d.Wait()
}

// realI3TreeProvider wraps the real i3.GetTree call.
type realI3TreeProvider struct{}

func (p *realI3TreeProvider) GetTree() (i3.Tree, error) {
	return i3.GetTree()
}


func printUsage() {
	fmt.Fprintf(os.Stderr, `Usage: vmux <command> [args...]

Commands:
  list              List active Claude Code sessions
  switch <query>    Switch to the workspace of the matching session
  label <id> <text> Assign a label to a session
  focus <minutes>   Suppress notifications for N minutes
  i3bar             Output i3bar JSON (use as status_command in i3 config)
  stop              Stop the vmux daemon
  daemon            Run the daemon in foreground (used internally)

Flags:
  --no-color        Disable colored output
`)
}