qmk-rgb/cmd/qmk-rgb-tool/effect.go

package main

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/spf13/cobra"
	intrgb "netdome.biz/paul/impact-80/internal/rgb"
	"netdome.biz/paul/impact-80/internal/via"
)

var listEffects bool

func resolveEffectTargets(name string, zones []intrgb.Zone) ([]intrgb.EffectTarget, []intrgb.Zone, error) {
	if name == "static" {
		name = "solid"
	}
	return intrgb.ResolveEffect(name, zones)
}

func NewEffectCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "effect [name]",
		Short: "Set or list RGB effects",
		Long:  "Set the RGB lighting effect on the connected keyboard. Without an argument, lists all effects.",
		Args:  cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if listEffects || len(args) == 0 {
				return listAllEffects(cmd)
			}
			return runEffectSet(cmd, args)
		},
	}
	cmd.Flags().BoolVar(&listEffects, "list", false, "List all available effects per zone")
	return cmd
}

func listAllEffects(cmd *cobra.Command) error {
	type ZoneEffectList struct {
		Zone   string `json:"zone"`
		Effect string `json:"effect"`
		ID     uint8  `json:"id"`
	}
	type EffectList struct {
		Zones []ZoneEffectList `json:"zones"`
	}

	var list EffectList
	for _, zone := range intrgb.AllZones() {
		names := intrgb.ZoneEffects(zone)
		for id, name := range names {
			list.Zones = append(list.Zones, ZoneEffectList{
				Zone:   string(zone),
				Effect: name,
				ID:     uint8(id),
			})
		}
	}

	data, _ := json.MarshalIndent(list.Zones, "", "  ")
	fmt.Fprintf(cmd.OutOrStdout(), "%s\n", string(data))
	return nil
}

func runEffectSet(cmd *cobra.Command, args []string) error {
	zones, err := selectedZones()
	if err != nil {
		return err
	}
	targets, skipped, err := resolveEffectTargets(args[0], zones)
	if err != nil {
		return err
	}
	if len(skipped) > 0 {
		skippedNames := make([]string, len(skipped))
		for i, zone := range skipped {
			skippedNames[i] = string(zone)
		}
		fmt.Fprintf(cmd.ErrOrStderr(), "Warning: effect not supported on zone(s): %s\n", strings.Join(skippedNames, ", "))
	}

	proto, err := openRGBProtocol()
	if err != nil {
		return err
	}
	defer proto.Close()

	for _, target := range targets {
		if err := proto.SetValue(via.LEDType(target.Zone.Channel()), uint8(intrgb.EffectID), target.ID); err != nil {
			return fmt.Errorf("set effect on %s: %w", target.Zone, err)
		}
	}

	fmt.Fprintf(cmd.OutOrStdout(), "Effect set to %q\n", args[0])
	return nil
}