chultun/internal/budget/units.go

package budget

import (
	"fmt"
	"slices"
	"strings"
)

const serviceText = "[Unit]\nDescription=Terminal workload resource profiles\n\n[Service]\nExecStart=%h/.local/bin/budget watch\n\n[Install]\nWantedBy=default.target\n"

func escape(name string) string {
	var result strings.Builder
	for i := range len(name) {
		b := name[i]
		switch {
		case b == '/':
			result.WriteByte('-')
		case b >= 'a' && b <= 'z', b >= 'A' && b <= 'Z', b >= '0' && b <= '9', b == ':', b == '_', b == '.' && i != 0:
			result.WriteByte(b)
		default:
			fmt.Fprintf(&result, `\x%02x`, b)
		}
	}
	return result.String()
}

func appDirs(app string) []string {
	app = escape(app)
	return []string{"app-" + app + ".service.d", "app-" + app + "@.service.d", "app-" + app + "-.scope.d", "app-flatpak-" + app + "-.scope.d"}
}

func appMatches(app, unit string) bool {
	app = escape(app)
	return unit == "app-"+app+".service" || strings.HasPrefix(unit, "app-"+app+"@") || strings.HasPrefix(unit, "app-"+app+"-") || strings.HasPrefix(unit, "app-flatpak-"+app+"-")
}

func (c config) sharedSlice() string {
	apps := keys(c.Shared)
	if len(apps) == 0 {
		return ""
	}
	return "app-" + escape(apps[0]) + ".slice"
}

func (c config) unitProfile(unit string) string {
	for _, name := range keys(c.Units) {
		if slices.Contains(c.Units[name], unit) {
			return name
		}
	}
	for _, app := range keys(c.Shared) {
		if appMatches(app, unit) {
			return c.Shared[app]
		}
	}
	for _, name := range keys(c.Apps) {
		for _, app := range c.Apps[name] {
			if appMatches(app, unit) {
				return name
			}
		}
	}
	for _, name := range keys(c.Prefixes) {
		for _, prefix := range c.Prefixes[name] {
			if prefixMatches(prefix, unit) {
				return name
			}
		}
	}
	if unit == "search.slice" {
		return "search"
	}
	if slice := c.sharedSlice(); slice != "" && unit == slice {
		return "default"
	}
	if strings.HasPrefix(unit, "app-") && (strings.HasSuffix(unit, ".scope") || strings.HasSuffix(unit, ".service")) {
		return "default"
	}
	return ""
}

func (c config) activeProfile(u unit) string {
	if name := c.unitProfile(u.Name); name != "" {
		return name
	}
	if strings.HasPrefix(u.Name, "budget-") && strings.HasSuffix(u.Name, ".scope") {
		if name, ok := strings.CutSuffix(u.Description, " workload"); ok && name != "search" {
			return name
		}
	}
	return ""
}

func prefixMatches(directory, unit string) bool {
	for _, kind := range []string{".service", ".scope", ".slice"} {
		if stem, ok := strings.CutSuffix(directory, kind+".d"); ok {
			return strings.HasPrefix(unit, stem) && strings.HasSuffix(unit, kind)
		}
	}
	return false
}

func dropSection(directory string) (string, error) {
	for _, kind := range []struct{ suffix, section string }{
		{".service.d", "Service"}, {".scope.d", "Scope"}, {".slice.d", "Slice"},
	} {
		if stem, ok := strings.CutSuffix(directory, kind.suffix); ok && stem != "" {
			return kind.section, nil
		}
	}
	return "", fmt.Errorf("invalid drop-in directory %q", directory)
}

func (c config) render() (map[string]string, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}
	files := map[string]string{}
	add := func(directory, name, extra string) error {
		values, err := c.assignments(name)
		if err != nil {
			return err
		}
		section, err := dropSection(directory)
		if err != nil {
			return err
		}
		files[directory+"/50-budget.conf"] = "[" + section + "]\n" + strings.Join(values, "\n") + "\n" + extra
		return nil
	}
	for _, directory := range []string{"app-.service.d", "app-.scope.d"} {
		if err := add(directory, "default", ""); err != nil {
			return nil, err
		}
	}
	for _, name := range keys(c.Apps) {
		for _, app := range c.Apps[name] {
			for _, directory := range appDirs(app) {
				if err := add(directory, name, ""); err != nil {
					return nil, err
				}
			}
		}
	}
	for _, name := range keys(c.Units) {
		for _, unit := range c.Units[name] {
			if err := add(unit+".d", name, ""); err != nil {
				return nil, err
			}
		}
	}
	for _, name := range keys(c.Prefixes) {
		for _, directory := range c.Prefixes[name] {
			if err := add(directory, name, ""); err != nil {
				return nil, err
			}
		}
	}
	for _, app := range keys(c.Shared) {
		for _, directory := range appDirs(app) {
			if err := add(directory, c.Shared[app], "Slice="+c.sharedSlice()+"\n"); err != nil {
				return nil, err
			}
		}
	}
	groups := map[string]string{"search.slice": "search"}
	if len(c.Shared) > 0 {
		groups[c.sharedSlice()] = "default"
	}
	for unit, name := range groups {
		values, err := c.assignments(name)
		if err != nil {
			return nil, err
		}
		files[unit] = "[Slice]\n" + strings.Join(values, "\n") + "\n"
	}
	files["budget.service"] = serviceText
	return files, nil
}