qmk-rgb/cmd/wobkey/rgb/rgb.go

package rgb

import (
	"fmt"
	"strconv"

	"github.com/spf13/cobra"
	intdevice "github.com/wobkey/rgb/internal/device"
	intrgb "github.com/wobkey/rgb/internal/rgb"
	"github.com/wobkey/rgb/internal/via"
)

var (
	targetDevice string
	targetZone   string
)

type zoneProtocol interface {
	SetValue(via.LEDType, uint8, uint8) error
	SetColor(via.LEDType, uint8, uint8) error
}

func forEachSelectedZone(zones []intrgb.Zone, fn func(intrgb.Zone, via.LEDType) error) error {
	channels := zoneChannels(zones)
	for i, zone := range zones {
		if err := fn(zone, channels[i]); err != nil {
			return err
		}
	}
	return nil
}

func setValueOnZones(proto zoneProtocol, zones []intrgb.Zone, param, value uint8) error {
	return forEachSelectedZone(zones, func(_ intrgb.Zone, channel via.LEDType) error {
		return proto.SetValue(channel, param, value)
	})
}

func setBrightnessOnZones(proto zoneProtocol, zones []intrgb.Zone, value uint8) error {
	return setValueOnZones(proto, zones, uint8(intrgb.Brightness), value)
}

func setSpeedOnZones(proto zoneProtocol, zones []intrgb.Zone, value uint8) error {
	return setValueOnZones(proto, zones, uint8(intrgb.Speed), value)
}

func setColorOnZones(proto zoneProtocol, zones []intrgb.Zone, hue, saturation uint8) error {
	return forEachSelectedZone(zones, func(_ intrgb.Zone, channel via.LEDType) error {
		return proto.SetColor(channel, hue, saturation)
	})
}

func disableLightingOnZones(proto zoneProtocol, zones []intrgb.Zone) error {
	return forEachSelectedZone(zones, func(_ intrgb.Zone, channel via.LEDType) error {
		if err := proto.SetValue(channel, uint8(intrgb.EffectID), 0); err != nil {
			return err
		}
		return proto.SetValue(channel, uint8(intrgb.Brightness), 0)
	})
}

func enableLightingOnZones(proto zoneProtocol, zones []intrgb.Zone) error {
	return forEachSelectedZone(zones, func(zone intrgb.Zone, channel via.LEDType) error {
		if err := proto.SetValue(channel, uint8(intrgb.EffectID), intrgb.DefaultEffect(zone)); err != nil {
			return err
		}
		return proto.SetValue(channel, uint8(intrgb.Brightness), 160)
	})
}

func setModeOnZones(proto zoneProtocol, zones []intrgb.Zone, value uint8) error {
	return setValueOnZones(proto, zones, uint8(intrgb.EffectID), value)
}

func OpenDevice() (*via.Protocol, intdevice.Keyboard, error) {
	if _, err := selectedZones(); err != nil {
		return nil, intdevice.Keyboard{}, err
	}

	keyboards, err := intdevice.LoadKeyboards()
	if err != nil {
		return nil, intdevice.Keyboard{}, fmt.Errorf("load keyboards: %w", err)
	}

	devices, err := intdevice.Discover(keyboards)
	if err != nil {
		return nil, intdevice.Keyboard{}, fmt.Errorf("discover: %w", err)
	}

	if len(devices) == 0 {
		return nil, intdevice.Keyboard{}, fmt.Errorf("no VIA-compatible keyboard found")
	}

	var dev intdevice.Device
	if targetDevice != "" {
		for _, d := range devices {
			if d.Path == targetDevice {
				dev = d
				break
			}
		}
		if dev.Path == "" {
			return nil, intdevice.Keyboard{}, fmt.Errorf("device %s not found", targetDevice)
		}
	} else if len(devices) == 1 {
		dev = devices[0]
	} else {
		return nil, intdevice.Keyboard{}, fmt.Errorf("multiple devices found, use --device to specify")
	}

	proto, err := via.New(dev)
	if err != nil {
		return nil, intdevice.Keyboard{}, fmt.Errorf("open protocol: %w", err)
	}

	return proto, dev.Keyboard, nil
}

func lightingEnabled(mode, brightness uint8) bool {
	return mode != 0 && brightness > 0
}

func ParseUint8(s string) (uint8, error) {
	v, err := strconv.ParseUint(s, 10, 8)
	if err != nil {
		return 0, fmt.Errorf("invalid value: %w", err)
	}
	return uint8(v), nil
}

func Init() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "rgb",
		Short: "RGB lighting commands",
		PersistentPreRunE: func(*cobra.Command, []string) error {
			_, err := selectedZones()
			return err
		},
	}
	cmd.PersistentFlags().StringVar(&targetDevice, "device", "", "HID device path to use")
	cmd.PersistentFlags().StringVar(&targetZone, "zone", "", "RGB lighting zone")
	return cmd
}