streamdeck-go/cmd/streamdeck/main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"image"
	"image/draw"
	"image/gif"
	_ "image/jpeg"
	_ "image/png"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/WoodardDigital/streamdeck-go/internal/config"
	"github.com/WoodardDigital/streamdeck-go/internal/device"
	"github.com/WoodardDigital/streamdeck-go/internal/modules"
	"github.com/fsnotify/fsnotify"
	"github.com/srwiley/oksvg"
	"github.com/srwiley/rasterx"
	"gopkg.in/yaml.v3"
)

// helperSocketPath returns the Unix socket path for the privileged helper.
// /run is the standard location on Linux; /var/run is used on macOS.
func helperSocketPath() string {
	if runtime.GOOS == "darwin" {
		return "/var/run/streamdeck-go/helper.sock"
	}
	return "/run/streamdeck-go/helper.sock"
}

func main() {
	cfgPath := flag.String("config", defaultConfigPath(), "path to config file")
	flag.Parse()

	if err := ensureConfigDir(*cfgPath); err != nil {
		log.Fatalf("setup: %v", err)
	}

	cfg, err := config.Load(*cfgPath)
	if err != nil {
		log.Fatalf("config: %v", err)
	}

	reg, err := modules.LoadRegistry(config.ModulesPath(*cfgPath))
	if err != nil {
		log.Fatalf("modules: %v", err)
	}

	// Open device — if not present at startup, wait for it.
	sd := mustConnect(cfg.Device.VendorID, cfg.Device.ProductID)

	// runDone receives true if run() exited because the device died,
	// false if it exited cleanly due to context cancellation.
	runDone := make(chan bool, 1)
	ctx, cancel := context.WithCancel(context.Background())

	startRun := func() {
		go func() {
			deviceDied := run(ctx, sd, cfg, reg)
			runDone <- deviceDied
		}()
	}
	startRun()

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatalf("watcher: %v", err)
	}
	defer watcher.Close()
	_ = watcher.Add(*cfgPath)
	_ = watcher.Add(filepath.Dir(*cfgPath))

	for {
		select {

		// Device died (disconnect, USB suspend, etc.) — reconnect and restart.
		case deviceDied := <-runDone:
			if !deviceDied {
				// Clean cancellation from a config reload below; already restarted.
				continue
			}
			log.Println("device disconnected — waiting for reconnect...")
			sd.Close()
			sd = mustConnect(cfg.Device.VendorID, cfg.Device.ProductID)
			log.Println("device reconnected")
			ctx, cancel = context.WithCancel(context.Background())
			startRun()

		// Config file or modules file changed — reload and restart run().
		case event, ok := <-watcher.Events:
			if !ok {
				cancel()
				return
			}
			cfgChanged := filepath.Clean(event.Name) == filepath.Clean(*cfgPath)
			modChanged := filepath.Clean(event.Name) == filepath.Clean(config.ModulesPath(*cfgPath))
			if !cfgChanged && !modChanged {
				continue
			}
			if !event.Has(fsnotify.Write) && !event.Has(fsnotify.Create) {
				continue
			}

			log.Println("config changed, reloading...")
			cancel()
			<-runDone // wait for run() to fully exit before restarting

			newCfg, err := config.Load(*cfgPath)
			if err != nil {
				log.Printf("reload: bad config: %v — keeping current config", err)
			} else {
				cfg = newCfg
			}
			newReg, err := modules.LoadRegistry(config.ModulesPath(*cfgPath))
			if err != nil {
				log.Printf("reload: bad modules: %v — keeping current modules", err)
			} else {
				reg = newReg
			}
			ctx, cancel = context.WithCancel(context.Background())
			startRun()

		case err, ok := <-watcher.Errors:
			if !ok {
				cancel()
				return
			}
			log.Printf("watcher error: %v", err)
		}
	}
}

// run initialises the deck and handles button presses until ctx is cancelled
// or the device dies. Returns true if the device died, false if ctx was cancelled.
func run(ctx context.Context, sd *device.StreamDeck, cfg *config.Config, reg *modules.Registry) (deviceDied bool) {
	// Recover from any panic inside this goroutine tree.
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic in run: %v — treating as device error", r)
			deviceDied = true
		}
	}()

	if err := sd.Reset(); err != nil {
		log.Printf("warn: reset: %v", err)
	}
	if err := sd.SetBrightness(cfg.Brightness); err != nil {
		log.Printf("warn: brightness: %v", err)
	}
	for i := 0; i < sd.KeyCount(); i++ {
		if err := sd.ClearKey(i); err != nil {
			log.Printf("warn: clear key %d: %v", i, err)
		}
	}

	// Track animation/poll goroutines so we can wait for them on exit.
	var wg sync.WaitGroup

	// triggers lets button presses immediately re-poll toggle keys.
	triggers := make(map[int]chan struct{})

	for keyIdx, keyCfg := range cfg.Keys {
		// Resolve module-based commands to shell strings before any other logic.
		if keyCfg.Module != "" {
			resolved, err := reg.Resolve(keyCfg.Module, keyCfg.Function, keyCfg.Params)
			if err != nil {
				log.Printf("key %d: module: %v", keyIdx, err)
				continue
			}
			keyCfg.Command = resolved
		}
		if keyCfg.Poll != nil && keyCfg.Poll.Module != "" {
			resolved, err := reg.Resolve(keyCfg.Poll.Module, keyCfg.Poll.Function, keyCfg.Poll.Params)
			if err != nil {
				log.Printf("key %d: poll module: %v", keyIdx, err)
				continue
			}
			keyCfg.Poll.Command = resolved
		}

		// Toggle/status key: managed by a polling goroutine.
		if keyCfg.Poll != nil {
			if keyCfg.IconTrue == "" || keyCfg.IconFalse == "" {
				log.Printf("key %d: toggle key requires icon_on and icon_off", keyIdx)
				continue
			}
			trigger := make(chan struct{}, 1)
			triggers[keyIdx] = trigger
			wg.Add(1)
			go func(idx int, kCfg config.KeyConfig, trig chan struct{}) {
				defer wg.Done()
				defer func() {
					if r := recover(); r != nil {
						log.Printf("panic in pollKey %d: %v", idx, r)
					}
				}()
				pollKey(ctx, sd, idx, kCfg, cfg.IconsDir, trig)
			}(keyIdx, keyCfg, trigger)
			continue
		}

		// Regular key: load icon once.
		if keyCfg.Icon == "" {
			continue
		}
		iconPath := filepath.Join(cfg.IconsDir, keyCfg.Icon)

		if strings.ToLower(filepath.Ext(keyCfg.Icon)) == ".gif" {
			frames, delays, err := loadGIF(sd, iconPath)
			if err != nil {
				log.Printf("key %d: load gif %q: %v", keyIdx, keyCfg.Icon, err)
				continue
			}
			wg.Add(1)
			go func(idx int, f [][]byte, d []time.Duration) {
				defer wg.Done()
				defer func() {
					if r := recover(); r != nil {
						log.Printf("panic in animateKey %d: %v", idx, r)
					}
				}()
				animateKey(ctx, sd, idx, f, d)
			}(keyIdx, frames, delays)
		} else {
			img, err := loadImage(iconPath)
			if err != nil {
				log.Printf("key %d: load icon %q: %v", keyIdx, keyCfg.Icon, err)
				continue
			}
			if err := sd.SetKeyImage(keyIdx, img); err != nil {
				log.Printf("key %d: set image: %v", keyIdx, err)
			}
		}
	}

	log.Printf("stream deck ready (%d keys)", sd.KeyCount())

	prev := make([]bool, sd.KeyCount())
	consecutiveErrors := 0

	for {
		select {
		case <-ctx.Done():
			wg.Wait()
			return false
		default:
		}

		buttons, err := sd.ReadButtons()
		if err != nil {
			consecutiveErrors++
			if consecutiveErrors >= 3 {
				log.Printf("device error (gave up after %d attempts): %v", consecutiveErrors, err)
				// Cancel animations, wait for them to stop, then signal device death.
				wg.Wait()
				return true
			}
			log.Printf("read error (%d/3): %v", consecutiveErrors, err)
			time.Sleep(200 * time.Millisecond)
			continue
		}
		consecutiveErrors = 0

		if buttons == nil {
			continue // 250 ms timeout, no data
		}

		for i, pressed := range buttons {
			if pressed && !prev[i] {
				keyCfg, ok := cfg.Keys[i]
				if !ok || keyCfg.Command == "" {
					continue
				}
				log.Printf("key %d pressed → %s", i, keyCfg.Command)
				go func(cmd string) {
					defer func() {
						if r := recover(); r != nil {
							log.Printf("panic running command %q: %v", cmd, r)
						}
					}()
					runCommand(cmd)
				}(keyCfg.Command)

				// For toggle keys, signal the poll goroutine to re-check state
				// shortly after the command runs (gives the system time to update).
				if ch, ok := triggers[i]; ok {
					select {
					case ch <- struct{}{}:
					default:
					}
				}
			}
		}
		prev = buttons
	}
}

// pollKey watches the state of a toggle key and keeps its icon up to date.
// It polls on an interval and also re-polls when triggered (e.g. after a button press).
func pollKey(ctx context.Context, sd *device.StreamDeck, keyIdx int, keyCfg config.KeyConfig, iconsDir string, trigger <-chan struct{}) {
	interval := 2 * time.Second
	if keyCfg.Poll.Interval != "" {
		if d, err := time.ParseDuration(keyCfg.Poll.Interval); err == nil {
			interval = d
		} else {
			log.Printf("key %d: invalid poll interval %q, using 2s", keyIdx, keyCfg.Poll.Interval)
		}
	}

	imgTrue, err := loadImage(filepath.Join(iconsDir, keyCfg.IconTrue))
	if err != nil {
		log.Printf("key %d: load icon_true %q: %v", keyIdx, keyCfg.IconTrue, err)
		return
	}
	imgFalse, err := loadImage(filepath.Join(iconsDir, keyCfg.IconFalse))
	if err != nil {
		log.Printf("key %d: load icon_false %q: %v", keyIdx, keyCfg.IconFalse, err)
		return
	}

	lastState := -1 // unknown, forces initial icon set

	applyState := func() {
		state := queryPollState(keyCfg.Poll)
		if state == lastState {
			return
		}
		lastState = state
		img := imgFalse
		if state == 1 {
			img = imgTrue
		}
		if err := sd.SetKeyImage(keyIdx, img); err != nil {
			log.Printf("key %d: set poll icon: %v", keyIdx, err)
		}
	}

	applyState() // set icon immediately on startup

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			applyState()
		case <-trigger:
			// Wait briefly for the toggle command to take effect, then re-poll.
			select {
			case <-ctx.Done():
				return
			case <-time.After(400 * time.Millisecond):
			}
			applyState()
		}
	}
}

// queryPollState runs the poll command and returns 1 (on) or 0 (off).
// If Match is set, checks for that substring in stdout.
// If Match is empty, uses exit code: 0 → on, non-zero → off.
func queryPollState(poll *config.PollConfig) int {
	cmd := exec.Command("sh", "-c", poll.Command)
	output, err := cmd.Output()
	if poll.Match == "" {
		if err == nil {
			return 1
		}
		return 0
	}
	if strings.Contains(string(output), poll.Match) {
		return 1
	}
	return 0
}

// mustConnect blocks until the device opens successfully.
//
// Strategy: try quickly at first (device may just be enumerating), then settle
// into a short fixed interval. We deliberately avoid long exponential backoff
// because in KVM setups the device is guaranteed to come back — we just don't
// know when. A 2s poll means the deck is live within ~2s of switching back.
func mustConnect(vendorID, productID uint16) *device.StreamDeck {
	const (
		quickRetries  = 5
		quickInterval = 500 * time.Millisecond
		pollInterval  = 2 * time.Second
	)
	for i := 0; ; i++ {
		sd, err := device.Open(vendorID, productID)
		if err == nil {
			return sd
		}
		wait := pollInterval
		if i < quickRetries {
			wait = quickInterval
		}
		log.Printf("could not open device: %v — retrying in %s", err, wait)
		time.Sleep(wait)
	}
}

// animateKey cycles pre-encoded GIF frames on a key until ctx is cancelled.
func animateKey(ctx context.Context, sd *device.StreamDeck, keyIdx int, frames [][]byte, delays []time.Duration) {
	for {
		for i, frame := range frames {
			select {
			case <-ctx.Done():
				return
			default:
			}
			if err := sd.SetKeyFrame(keyIdx, frame); err != nil {
				log.Printf("key %d: animate frame %d: %v", keyIdx, i, err)
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(delays[i]):
			}
		}
	}
}

// loadGIF decodes a GIF and pre-encodes every frame as JPEG bytes.
func loadGIF(sd *device.StreamDeck, path string) (frames [][]byte, delays []time.Duration, err error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, nil, err
	}
	defer f.Close()

	g, err := gif.DecodeAll(f)
	if err != nil {
		return nil, nil, err
	}

	canvas := image.NewRGBA(image.Rect(0, 0, g.Config.Width, g.Config.Height))
	for i, frame := range g.Image {
		draw.Draw(canvas, frame.Bounds(), frame, frame.Bounds().Min, draw.Over)
		encoded, err := sd.EncodeFrame(canvas)
		if err != nil {
			return nil, nil, err
		}
		frames = append(frames, encoded)
		d := g.Delay[i]
		if d <= 0 {
			d = 10
		}
		delays = append(delays, time.Duration(d)*10*time.Millisecond)
	}
	return frames, delays, nil
}

func loadImage(path string) (image.Image, error) {
	if strings.ToLower(filepath.Ext(path)) == ".svg" {
		return loadSVG(path)
	}
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	img, _, err := image.Decode(f)
	return img, err
}

func loadSVG(path string) (image.Image, error) {
	icon, err := oksvg.ReadIcon(path, oksvg.StrictErrorMode)
	if err != nil {
		return nil, err
	}
	const size = 96
	icon.SetTarget(0, 0, size, size)
	rgba := image.NewRGBA(image.Rect(0, 0, size, size))
	scanner := rasterx.NewScannerGV(size, size, rgba, rgba.Bounds())
	raster := rasterx.NewDasher(size, size, scanner)
	icon.Draw(raster, 1.0)
	return rgba, nil
}

func runCommand(cmd string) {
	if strings.HasPrefix(cmd, "priv:") {
		name := strings.TrimPrefix(cmd, "priv:")
		if err := runPrivileged(name); err != nil {
			log.Printf("privileged command %q: %v", name, err)
		}
		return
	}
	c := exec.Command("sh", "-c", cmd)
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	if err := c.Run(); err != nil {
		log.Printf("command %q: %v", cmd, err)
	}
}

// runPrivileged dispatches a named privileged command.
// On macOS: reads the user's local whitelist and runs via osascript (admin auth dialog).
// On Linux: sends to the root helper daemon over its Unix socket.
func runPrivileged(name string) error {
	if runtime.GOOS == "darwin" {
		return runPrivilegedDarwin(name)
	}
	return runPrivilegedHelper(name)
}

// runPrivilegedDarwin looks up name in ~/.config/streamdeck-go/privileged.yaml
// and executes it via osascript, which shows the standard macOS admin auth dialog.
func runPrivilegedDarwin(name string) error {
	wlPath := darwinWhitelistPath()
	commands, err := loadPrivilegedCommands(wlPath)
	if err != nil {
		return fmt.Errorf("load whitelist %q: %w", wlPath, err)
	}
	shell, ok := commands[name]
	if !ok {
		return fmt.Errorf("unknown command %q — add it to %s", name, wlPath)
	}
	script := fmt.Sprintf(`do shell script %q with administrator privileges`, shell)
	out, err := exec.Command("osascript", "-e", script).CombinedOutput()
	if err != nil {
		return fmt.Errorf("%w: %s", err, out)
	}
	return nil
}

// runPrivilegedHelper sends a named command to the root helper daemon over its Unix socket.
func runPrivilegedHelper(name string) error {
	conn, err := net.Dial("unix", helperSocketPath())
	if err != nil {
		return fmt.Errorf("helper unavailable (is streamdeck-go-helper.service running?): %w", err)
	}
	defer conn.Close()

	req, _ := json.Marshal(map[string]string{"command": name})
	if _, err := fmt.Fprintf(conn, "%s\n", req); err != nil {
		return fmt.Errorf("send: %w", err)
	}

	scanner := bufio.NewScanner(conn)
	if !scanner.Scan() {
		return fmt.Errorf("no response from helper")
	}

	var resp struct {
		OK    bool   `json:"ok"`
		Error string `json:"error"`
	}
	if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
		return fmt.Errorf("bad response: %w", err)
	}
	if !resp.OK {
		return fmt.Errorf("%s", resp.Error)
	}
	return nil
}

func darwinWhitelistPath() string {
	base := os.Getenv("XDG_CONFIG_HOME")
	if base == "" {
		home, _ := os.UserHomeDir()
		base = filepath.Join(home, ".config")
	}
	return filepath.Join(base, "streamdeck-go", "privileged.yaml")
}

func loadPrivilegedCommands(path string) (map[string]string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var wl struct {
		Commands map[string]string `yaml:"commands"`
	}
	if err := yaml.Unmarshal(data, &wl); err != nil {
		return nil, err
	}
	if wl.Commands == nil {
		return map[string]string{}, nil
	}
	return wl.Commands, nil
}

func defaultConfigPath() string {
	return config.DefaultConfigPath()
}

func ensureConfigDir(cfgPath string) error {
	dir := filepath.Dir(cfgPath)
	iconsDir := filepath.Join(dir, "icons")
	if err := os.MkdirAll(iconsDir, 0755); err != nil {
		return err
	}
	if _, err := os.Stat(cfgPath); os.IsNotExist(err) {
		return os.WriteFile(cfgPath, []byte(defaultConfig(iconsDir)), 0644)
	}
	return nil
}

func defaultConfig(iconsDir string) string {
	return `# streamdeck-go configuration
# https://github.com/WoodardDigital/streamdeck-go

icons_dir: ` + iconsDir + `
brightness: 70

# USB IDs — defaults match Stream Deck XL v2.
# Run: lsusb | grep Elgato
device:
  vendor_id: 0x0fd9
  product_id: 0x00ba

# Keys are 0-indexed, left-to-right, top-to-bottom.
# Stream Deck XL layout (8 columns x 4 rows):
#
#  0  1  2  3  4  5  6  7
#  8  9 10 11 12 13 14 15
# 16 17 18 19 20 21 22 23
# 24 25 26 27 28 29 30 31
#
# icon: filename inside icons_dir (PNG, JPEG, or GIF)
# command: shell command to run on press

keys: {}
`
}