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

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

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

// profilesDir is a var so tests can redirect it at a temp directory.
var profilesDir = "profiles"

type Profile struct {
	Name    string                   `json:"name"`
	Version int                      `json:"version"`
	Board   *Board                   `json:"board,omitempty"`
	Zones   map[string]*ZoneSettings `json:"zones"`
}

// Board identifies the keyboard a profile was saved from. A profile stores
// effect names, and names belong to a board, so without this a profile written
// for one keyboard would be applied to another without a word. The pair is
// spelled as hex strings, the way the VIA definition files
// spell it, so the three can be read side by side.
type Board struct {
	VendorID  string `json:"vendorId"`
	ProductID string `json:"productId"`
}

// boardMatch is what comparing a profile's board to the connected one tells us.
type boardMatch int

const (
	// boardUnknown is a profile written before profiles carried a board.
	boardUnknown boardMatch = iota
	boardSame
	boardOther
)

func boardFor(vendorID, productID uint16) *Board {
	return &Board{
		VendorID:  fmt.Sprintf("0x%04X", vendorID),
		ProductID: fmt.Sprintf("0x%04X", productID),
	}
}

// MatchesBoard reports whether a profile was saved from the connected keyboard,
// was saved from another one, or says nothing about it.
func (p *Profile) MatchesBoard(vendorID, productID uint16) boardMatch {
	if p == nil || p.Board == nil {
		return boardUnknown
	}
	if strings.EqualFold(p.Board.VendorID, fmt.Sprintf("0x%04X", vendorID)) &&
		strings.EqualFold(p.Board.ProductID, fmt.Sprintf("0x%04X", productID)) {
		return boardSame
	}
	return boardOther
}

// boardMismatchWarning says which profile belongs to which keyboard, because
// the user has to be able to tell which of their profiles is the wrong one.
func (p *Profile) boardMismatchWarning(vendorID, productID uint16) string {
	return fmt.Sprintf("Warning: profile %q was saved for keyboard %s/%s, and this keyboard is %s/%s; "+
		"its effect names may not exist here\n",
		p.Name, p.Board.VendorID, p.Board.ProductID,
		fmt.Sprintf("0x%04X", vendorID), fmt.Sprintf("0x%04X", productID))
}

// listLine is one profile as the list command prints it.
func (p *Profile) listLine() string {
	if p.Board == nil {
		return p.Name
	}
	return fmt.Sprintf("%s  %s/%s", p.Name, p.Board.VendorID, p.Board.ProductID)
}

type ZoneSettings struct {
	Enabled    bool   `json:"enabled"`
	Effect     string `json:"effect"`
	Brightness uint8  `json:"brightness"`
	Speed      uint8  `json:"speed"`
	Color      string `json:"color"`
}

func ProfilesPath() string {
	return filepath.Join(profilesDir)
}

func (p *Profile) Save() error {
	if p.Name == "" {
		return fmt.Errorf("profile name is required")
	}

	if err := os.MkdirAll(profilesDir, 0755); err != nil {
		return fmt.Errorf("create profiles directory: %w", err)
	}

	data, err := json.MarshalIndent(p, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal profile: %w", err)
	}

	path := filepath.Join(profilesDir, sanitizeFilename(p.Name)+".json")
	if err := os.WriteFile(path, data, 0644); err != nil {
		return fmt.Errorf("write profile: %w", err)
	}
	return nil
}

func LoadProfile(name string) (*Profile, error) {
	path := filepath.Join(profilesDir, sanitizeFilename(name)+".json")
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read profile %s: %w", name, err)
	}

	var p Profile
	if err := json.Unmarshal(data, &p); err != nil {
		return nil, fmt.Errorf("parse profile %s: %w", name, err)
	}
	return &p, nil
}

func ListProfiles() ([]string, error) {
	entries, err := os.ReadDir(profilesDir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("list profiles: %w", err)
	}

	var names []string
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if strings.HasSuffix(name, ".json") {
			name = name[:len(name)-5]
			names = append(names, name)
		}
	}
	return names, nil
}

func DeleteProfile(name string) error {
	path := filepath.Join(profilesDir, sanitizeFilename(name)+".json")
	if err := os.Remove(path); err != nil {
		if os.IsNotExist(err) {
			return fmt.Errorf("profile %s not found", name)
		}
		return fmt.Errorf("delete profile %s: %w", name, err)
	}
	return nil
}

func sanitizeFilename(name string) string {
	name = strings.ToLower(name)
	name = strings.Map(func(r rune) rune {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' || r == '_' {
			return r
		}
		return '-'
	}, name)
	if len(name) > 0 && name[0] == '-' {
		name = "unnamed-" + name
	}
	return name
}

// applyProfile reads RGB state from the device and stores it in a Profile.
func applyProfileToProfile(proto rgbProtocol, channels []via.Channel, display map[uint16]string, catalog *intrgb.Catalog, p *Profile) error {
	out, err := readInfo(proto, channels, display, catalog)
	if err != nil {
		return fmt.Errorf("read device state: %w", err)
	}

	p.Zones = make(map[string]*ZoneSettings)
	for _, zi := range out.Zones {
		if zi.Error != "" {
			continue
		}
		p.Zones[zi.Zone] = &ZoneSettings{
			Enabled:    zi.Enabled,
			Effect:     zi.Effect,
			Brightness: zi.Brightness,
			Speed:      zi.Speed,
			Color:      fmt.Sprintf("%02x%02x", zi.Color.Hue, zi.Color.Saturation),
		}
	}
	return nil
}

// loadProfileFromDevice reads RGB state from device and saves it. Warnings go
// to warn, which is the command's stderr.
func loadProfileFromDevice(name string, warn io.Writer) error {
	proto, target, channels, err := openTarget()
	if err != nil {
		return err
	}
	defer proto.Close()

	catalog, _, err := resolveCatalog(target)
	if err != nil {
		return err
	}
	if catalog == nil {
		// The profile stores effect names and a board without a catalog has none
		// to store, so every channel is written as "unknown" and cannot be
		// restored. Say so here, where the user can still act on it.
		fmt.Fprintf(warn,
			"Warning: this keyboard has no effect names, so the profile records effect %q and cannot restore it; "+
				"run `definition fetch` for its VIA definition, or set an effect with `effect <index>`\n", "unknown")
	}

	p := &Profile{
		Name:    name,
		Version: 1,
		Board:   boardFor(target.Device.VendorID, target.Device.ProductID),
	}
	if err := applyProfileToProfile(proto, channels, target.Display, catalog, p); err != nil {
		return err
	}
	return p.Save()
}

func NewProfileSaveCmd() *cobra.Command {
	var name string
	cmd := &cobra.Command{
		Use:   "save [name]",
		Short: "Save current RGB state to a profile",
		Long:  "Read the current RGB settings from the keyboard and save them as a JSON profile in the profiles/ directory.",
		Args:  cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 0 {
				name = "default"
			} else {
				name = args[0]
			}
			return loadProfileFromDevice(name, cmd.ErrOrStderr())
		},
	}
	return cmd
}

func NewProfileLoadCmd() *cobra.Command {
	var name string
	cmd := &cobra.Command{
		ValidArgsFunction: completeProfileNames,
		Use:               "load [name]",
		Short:             "Load a profile and apply it to the keyboard",
		Long:              "Read a JSON profile from the profiles/ directory and apply the saved RGB settings to the keyboard.",
		Args:              cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 0 {
				name = "default"
			} else {
				name = args[0]
			}

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

			p, err := LoadProfile(name)
			if err != nil {
				return err
			}

			// A profile belongs to the keyboard it was saved from: its effect
			// names are that board's. Applying it elsewhere is allowed, because
			// the names that do not exist are reported per key below, but it is
			// said out loud, since a mismatch is the likeliest reason.
			if p.MatchesBoard(target.Device.VendorID, target.Device.ProductID) == boardOther {
				fmt.Fprintln(cmd.ErrOrStderr(), p.boardMismatchWarning(target.Device.VendorID, target.Device.ProductID))
			}

			catalog, _, err := resolveCatalog(target)
			if err != nil {
				return err
			}

			keys := make([]string, 0, len(p.Zones))
			for key := range p.Zones {
				keys = append(keys, key)
			}
			sort.Strings(keys)

			// A key is a channel name as the file wrote it. Resolving it here
			// means a profile written before a board was renamed reports the
			// key it cannot place instead of silently applying nothing.
			selectedSet := make(map[via.Channel]bool, len(selected))
			for _, ch := range selected {
				selectedSet[ch] = true
			}

			// Presence and selection are two different questions. A key the
			// selection leaves out is the user's own choice and stays quiet; a
			// key naming a channel the keyboard does not have has nowhere to go
			// and is reported, whatever the selection says.
			present, err := proto.DetectChannels()
			if err != nil {
				return err
			}
			presentSet := make(map[via.Channel]bool, len(present))
			for _, ch := range present {
				presentSet[ch] = true
			}

			// Without a catalog there are no effect names to look the
			// profile's value up in, so the whole load is impossible. Say that
			// once instead of reporting every key's name as not found.
			if catalog == nil {
				fmt.Fprintf(cmd.ErrOrStderr(),
					"Warning: profile %q has no effect names for this keyboard, so nothing applied; "+
						"run `definition fetch` for its VIA definition\n", name)
				return nil
			}

			applied := 0
			for _, key := range keys {
				settings := p.Zones[key]
				keyChannels, err := resolveZoneName(key, target.Display, target.Alternatives)
				if err != nil {
					fmt.Fprintf(cmd.ErrOrStderr(),
						"Warning: profile %q names zone %q, which this keyboard does not have; skipping\n", name, key)
					continue
				}

				placed := false
				for _, ch := range keyChannels {
					if presentSet[ch] {
						placed = true
						break
					}
				}
				if !placed {
					fmt.Fprintf(cmd.ErrOrStderr(),
						"Warning: profile %q names zone %q, which this keyboard does not have; skipping\n", name, key)
					continue
				}

				for _, ch := range keyChannels {
					// A key is applied only where the selection allows it, so
					// `--zone logo load <name>` leaves the other channels alone.
					if !selectedSet[ch] {
						continue
					}
					targets, _, err := intrgb.ResolveEffect(catalog, settings.Effect, []via.Channel{ch}, targetZone != "")
					if err != nil {
						fmt.Fprintf(cmd.ErrOrStderr(),
							"Warning: effect %q not found on %s, skipping\n", settings.Effect, channelName(ch, target.Display))
						continue
					}

					// ResolveEffect can answer with nothing to do and no
					// error: the name exists on the board but not on this
					// channel, and no zone was named, so it is skipped. Indexing
					// its result would crash here instead of saying so.
					if len(targets) == 0 {
						fmt.Fprintf(cmd.ErrOrStderr(),
							"Warning: effect %q not found on %s, skipping\n", settings.Effect, channelName(ch, target.Display))
						continue
					}

					if err := proto.SetValue(ch, uint8(intrgb.EffectID), targets[0].ID); err != nil {
						return err
					}
					if err := proto.SetValue(ch, uint8(intrgb.Brightness), settings.Brightness); err != nil {
						return err
					}
					if err := proto.SetValue(ch, uint8(intrgb.Speed), settings.Speed); err != nil {
						return err
					}

					if settings.Enabled && settings.Color != "" {
						hue, sat, err := hexToHSV(settings.Color)
						if err != nil {
							fmt.Fprintf(cmd.ErrOrStderr(),
								"Warning: invalid color %q on zone %s, skipping\n", settings.Color, channelName(ch, target.Display))
						} else if err := proto.SetColor(ch, hue, sat); err != nil {
							return err
						}
					}
					applied++
				}
			}

			if applied == 0 {
				fmt.Fprintf(cmd.ErrOrStderr(),
					"Warning: profile %q has no settings for the selected zone(s); nothing applied\n", name)
			}
			return nil
		},
	}
	return cmd
}

func NewProfileListCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "List saved profiles",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			names, err := ListProfiles()
			if err != nil {
				return err
			}
			if names == nil {
				names = []string{}
			}
			lines := make([]string, 0, len(names))
			for _, name := range names {
				p, err := LoadProfile(name)
				if err != nil {
					// A file that cannot be read is still a name in the
					// directory; the name is what the user can act on.
					lines = append(lines, name)
					continue
				}
				lines = append(lines, p.listLine())
			}
			if jsonOutput {
				return encodeJSON(cmd.OutOrStdout(), struct {
					Profiles []string `json:"profiles"`
				}{Profiles: names})
			}
			for _, line := range lines {
				fmt.Fprintln(cmd.OutOrStdout(), line)
			}
			return nil
		},
	}
}

func NewProfileDeleteCmd() *cobra.Command {
	var name string
	cmd := &cobra.Command{
		ValidArgsFunction: completeProfileNames,
		Use:               "delete [name]",
		Short:             "Delete a saved profile",
		Args:              cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 0 {
				name = "default"
			} else {
				name = args[0]
			}
			return DeleteProfile(name)
		},
	}
	return cmd
}

func hexToHSV(s string) (uint8, uint8, error) {
	if len(s) != 4 {
		return 0, 0, fmt.Errorf("invalid HSV hex (expected 4 hex digits): %s", s)
	}
	var hue, sat uint8
	_, err := fmt.Sscanf(s, "%02x%02x", &hue, &sat)
	return hue, sat, err
}