yay/pkg/upgrade/service.go

package upgrade

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"

	"github.com/Jguer/aur"
	alpm "github.com/Jguer/dyalpm"
	mapset "github.com/deckarep/golang-set/v2"
	"github.com/leonelquinteros/gotext"

	"github.com/Jguer/yay/v13/pkg/db"
	"github.com/Jguer/yay/v13/pkg/dep"
	"github.com/Jguer/yay/v13/pkg/dep/topo"
	"github.com/Jguer/yay/v13/pkg/intrange"
	"github.com/Jguer/yay/v13/pkg/query"
	"github.com/Jguer/yay/v13/pkg/settings"
	settingslua "github.com/Jguer/yay/v13/pkg/settings/lua"
	"github.com/Jguer/yay/v13/pkg/text"
	"github.com/Jguer/yay/v13/pkg/vcs"
)

const cutOffExtra = 2

type UpgradeService struct {
	grapher    *dep.Grapher
	aurCache   aur.QueryClient
	dbExecutor db.Executor
	vcsStore   vcs.Store
	cfg        *settings.Configuration
	log        *text.Logger
	lua        *settingslua.Engine
	noConfirm  bool

	AURWarnings *query.AURWarnings
}

func NewUpgradeService(grapher *dep.Grapher, aurCache aur.QueryClient,
	dbExecutor db.Executor, vcsStore vcs.Store,
	cfg *settings.Configuration, noConfirm bool, logger *text.Logger,
) *UpgradeService {
	return &UpgradeService{
		grapher:     grapher,
		aurCache:    aurCache,
		dbExecutor:  dbExecutor,
		vcsStore:    vcsStore,
		cfg:         cfg,
		noConfirm:   noConfirm,
		log:         logger,
		AURWarnings: query.NewWarnings(logger.Child("warnings")),
	}
}

func (u *UpgradeService) SetLua(engine *settingslua.Engine) {
	u.lua = engine
}

// upGraph adds packages to upgrade to the graph.
func (u *UpgradeService) upGraph(ctx context.Context, graph *topo.Graph[string, *dep.InstallInfo],
	enableDowngrade bool,
	filter Filter,
) (err error) {
	var (
		develUp      UpSlice
		errs         []error
		aurdata      = make(map[string]*aur.Pkg)
		aurUp        UpSlice
		syncUpgrades map[string]db.SyncUpgrade
	)

	remote := u.dbExecutor.InstalledRemotePackages()
	remoteNames := u.dbExecutor.InstalledRemotePackageNames()

	if u.cfg.Mode.AtLeastRepo() {
		syncUpgrades, err = u.dbExecutor.SyncUpgrades(enableDowngrade)
		errs = append(errs, err)

		replaced := u.syncReplacedPackageNames(syncUpgrades)
		if replaced.Cardinality() > 0 {
			filteredRemote := make(map[string]db.IPackage, len(remote))
			for name, pkg := range remote {
				if !replaced.Contains(name) {
					filteredRemote[name] = pkg
				}
			}
			remote = filteredRemote
			remoteNames = slices.DeleteFunc(slices.Clone(remoteNames), func(name string) bool {
				return replaced.Contains(name)
			})
		}
	}

	if u.cfg.Mode.AtLeastAUR() {
		u.log.OperationInfoln(gotext.Get("Searching AUR for updates..."))

		_aurdata, err := u.aurCache.Get(ctx, &aur.Query{Needles: remoteNames, By: aur.Name})

		errs = append(errs, err)

		if err == nil {
			for i := range _aurdata {
				pkg := &_aurdata[i]
				aurdata[pkg.Name] = pkg
				u.AURWarnings.AddToWarnings(remote, pkg)
			}

			u.AURWarnings.CalculateMissing(remoteNames, remote, aurdata)

			aurUp = UpAUR(u.log, remote, aurdata, enableDowngrade)

			if u.cfg.Devel {
				u.log.OperationInfoln(gotext.Get("Checking development packages..."))

				develUp = UpDevel(ctx, u.log, remote, aurdata, u.vcsStore)

				u.vcsStore.CleanOrphans(remote)
			}
		}
	}

	aurPkgsAdded := []*aur.Pkg{}

	names := mapset.NewThreadUnsafeSet[string]()
	for i := range develUp.Up {
		up := &develUp.Up[i]
		// check if deps are satisfied for aur packages
		reason := dep.Explicit
		if up.Reason == alpm.PkgReasonDepend {
			reason = dep.Dep
		}

		if filter != nil && !filter(up) {
			continue
		}

		aurPkg := aurdata[up.Name]
		graph = u.grapher.GraphAURTarget(ctx, graph, aurPkg, &dep.InstallInfo{
			Reason:       reason,
			Source:       dep.AUR,
			AURBase:      aurPkg.PackageBase,
			Upgrade:      true,
			Devel:        true,
			LocalVersion: up.LocalVersion,
			Version:      up.RemoteVersion,
			LastModified: up.LastModified,
			Maintainer:   aurPkg.Maintainer,
		})
		names.Add(up.Name)
		aurPkgsAdded = append(aurPkgsAdded, aurPkg)
	}

	for i := range aurUp.Up {
		up := &aurUp.Up[i]
		// add devel packages if they are not already in the list
		if names.Contains(up.Name) {
			continue
		}

		// check if deps are satisfied for aur packages
		reason := dep.Explicit
		if up.Reason == alpm.PkgReasonDepend {
			reason = dep.Dep
		}

		if filter != nil && !filter(up) {
			continue
		}

		aurPkg := aurdata[up.Name]
		graph = u.grapher.GraphAURTarget(ctx, graph, aurPkg, &dep.InstallInfo{
			Reason:       reason,
			Source:       dep.AUR,
			AURBase:      aurPkg.PackageBase,
			Upgrade:      true,
			Version:      up.RemoteVersion,
			LocalVersion: up.LocalVersion,
			LastModified: up.LastModified,
			Maintainer:   aurPkg.Maintainer,
		})
		aurPkgsAdded = append(aurPkgsAdded, aurPkg)
	}

	u.grapher.AddDepsForPkgs(ctx, aurPkgsAdded, graph)

	if u.cfg.Mode.AtLeastRepo() {
		u.log.OperationInfoln(gotext.Get("Searching databases for updates..."))

		for _, up := range syncUpgrades {
			if filter != nil && !filter(&db.Upgrade{
				Name:          up.Package.Name(),
				RemoteVersion: up.Package.Version(),
				Repository:    up.Package.DB().Name(),
				Base:          up.Package.Base(),
				LocalVersion:  up.LocalVersion,
				Reason:        up.Reason,
			}) {
				continue
			}

			upgradeInfo := up
			graph = u.grapher.GraphSyncPkg(ctx, graph, up.Package, &upgradeInfo)
		}
	}

	return errors.Join(errs...)
}

func (u *UpgradeService) syncReplacedPackageNames(syncUpgrades map[string]db.SyncUpgrade) mapset.Set[string] {
	replaced := mapset.NewThreadUnsafeSet[string]()
	for _, up := range syncUpgrades {
		for _, replacement := range u.dbExecutor.PackageReplaces(up.Package) {
			replaced.Add(replacement.Name)
		}
	}

	return replaced
}

func (u *UpgradeService) graphToUpSlice(graph *topo.Graph[string, *dep.InstallInfo]) (aurUp, repoUp UpSlice) {
	aurUp = UpSlice{Up: make([]Upgrade, 0, graph.Len())}
	repoUp = UpSlice{Up: make([]Upgrade, 0, graph.Len()), Repos: u.dbExecutor.Repos()}

	_ = graph.ForEach(func(name string, info *dep.InstallInfo) error {
		alpmReason := alpm.PkgReasonDepend
		if info.Reason == dep.Explicit {
			alpmReason = alpm.PkgReasonExplicit
		}

		parents := graph.ImmediateDependencies(name)
		extra := ""
		if len(parents) > 0 && !info.Upgrade && info.Reason == dep.MakeDep {
			reducedParents := parents.Slice()[:min(cutOffExtra, len(parents))]
			if len(parents) > cutOffExtra {
				reducedParents = append(reducedParents, "...")
			}
			extra = fmt.Sprintf(" (%s of %s)", dep.ReasonNames[info.Reason], strings.Join(reducedParents, ", "))
		}

		switch info.Source {
		case dep.AUR:
			aurRepo := "aur"
			if info.Devel {
				aurRepo = "devel"
			}
			aurUp.Up = append(aurUp.Up, Upgrade{
				Name:          name,
				RemoteVersion: info.Version,
				Repository:    aurRepo,
				Base:          info.AURBase,
				LocalVersion:  info.LocalVersion,
				Reason:        alpmReason,
				Extra:         extra,
				LastModified:  info.LastModified,
				Maintainer:    info.Maintainer,
			})
		case dep.Sync:
			repoUp.Up = append(repoUp.Up, Upgrade{
				Name:          name,
				RemoteVersion: info.Version,
				Repository:    info.SyncDBName,
				Base:          "",
				LocalVersion:  info.LocalVersion,
				Reason:        alpmReason,
				Extra:         extra,
			})
		}
		return nil
	})

	return aurUp, repoUp
}

func (u *UpgradeService) GraphUpgrades(ctx context.Context,
	graph *topo.Graph[string, *dep.InstallInfo],
	enableDowngrade bool, filter Filter,
) (*topo.Graph[string, *dep.InstallInfo], error) {
	if graph == nil {
		graph = dep.NewGraph()
	}

	err := u.upGraph(ctx, graph, enableDowngrade, filter)
	if err != nil {
		return graph, err
	}

	return graph, nil
}

func (u *UpgradeService) upgradeSelection(graph *topo.Graph[string, *dep.InstallInfo]) UpSlice {
	aurUp, repoUp := u.graphToUpSlice(graph)

	slices.SortFunc(repoUp.Up, repoUp.compare)
	slices.SortFunc(aurUp.Up, aurUp.compare)

	allUp := UpSlice{Repos: append(repoUp.Repos, aurUp.Repos...)}
	for i := range repoUp.Up {
		if repoUp.Up[i].LocalVersion == "" && repoUp.Up[i].Reason != alpm.PkgReasonExplicit {
			allUp.PulledDeps = append(allUp.PulledDeps, repoUp.Up[i])
		} else {
			allUp.Up = append(allUp.Up, repoUp.Up[i])
		}
	}

	for i := range aurUp.Up {
		if aurUp.Up[i].LocalVersion == "" && aurUp.Up[i].Reason != alpm.PkgReasonExplicit {
			allUp.PulledDeps = append(allUp.PulledDeps, aurUp.Up[i])
		} else {
			allUp.Up = append(allUp.Up, aurUp.Up[i])
		}
	}

	return allUp
}

func (u *UpgradeService) runUpgradeSelectHook(graph *topo.Graph[string, *dep.InstallInfo], allUp UpSlice) (
	excluded []string, skipMenu bool, err error,
) {
	if u.lua == nil || !u.lua.HasAutocmd(settingslua.EventUpgradeSelect) {
		return []string{}, false, nil
	}

	var result settingslua.UpgradeSelectResult
	result, err = u.lua.RunUpgradeSelect(upgradeSelectEvent(allUp))
	if err != nil {
		return nil, false, err
	}

	excluded = u.pruneUpgradeNames(graph, result.Exclude)

	return excluded, result.SkipMenu, nil
}

func (u *UpgradeService) pruneUpgradeNames(graph *topo.Graph[string, *dep.InstallInfo], names []string) []string {
	excluded := make([]string, 0, len(names))
	for _, name := range names {
		if !graph.Exists(name) {
			continue
		}

		u.log.Debugln("pruning", name)
		excluded = append(excluded, graph.Prune(name)...)
	}

	return excluded
}

func (u *UpgradeService) printUpgradeSelection(allUp UpSlice) {
	if len(allUp.PulledDeps) > 0 {
		u.log.Printf("%s"+text.Bold(" %d ")+"%s\n", text.Bold(text.Cyan("::")),
			len(allUp.PulledDeps), text.Bold(gotext.Get("%s will also be installed for this operation.",
				gotext.GetN("dependency", "dependencies", len(allUp.PulledDeps)))))
		allUp.PrintDeps(u.log)
	}

	u.log.Printf("%s"+text.Bold(" %d ")+"%s\n", text.Bold(text.Cyan("::")),
		len(allUp.Up), text.Bold(gotext.Get("%s to upgrade/install.", gotext.GetN("package", "packages", len(allUp.Up)))))
	allUp.Print(u.log)
}

// userExcludeUpgrades asks the user which packages to exclude from the upgrade and
// removes them from the graph
func (u *UpgradeService) UserExcludeUpgrades(graph *topo.Graph[string, *dep.InstallInfo]) ([]string, error) {
	if graph.Len() == 0 {
		return []string{}, nil
	}
	allUp := u.upgradeSelection(graph)

	excluded, skipMenu, err := u.runUpgradeSelectHook(graph, allUp)
	if err != nil {
		return nil, err
	}
	if skipMenu {
		return excluded, nil
	}
	if len(excluded) > 0 {
		// The hook pruned packages; refresh the selection to reflect the
		// smaller graph and skip the menu if nothing is left to upgrade.
		allUp = u.upgradeSelection(graph)
		if len(allUp.Up) == 0 {
			return excluded, nil
		}
	}

	u.printUpgradeSelection(allUp)

	u.log.Infoln(gotext.Get("Packages to exclude: (eg: \"1 2 3\", \"1-3\", \"^4\" or repo name)"))
	u.log.Warnln(gotext.Get("Excluding packages may cause partial upgrades and break systems"))

	numbers, err := u.log.GetInput(u.cfg.AnswerUpgrade, settings.NoConfirm)
	if err != nil {
		return nil, err
	}

	// upgrade menu asks you which packages to NOT upgrade so in this case
	// exclude and include are kind of swapped
	exclude, include, otherExclude, otherInclude := intrange.ParseNumberMenu(numbers)

	// true if user doesn't want to include specific repositories/packages
	noIncludes := len(include) == 0 && otherInclude.Cardinality() == 0

	// No exclusions or inclusions specified, return early
	if noIncludes && len(exclude) == 0 && otherExclude.Cardinality() == 0 {
		return excluded, nil
	}

	for i := range allUp.Up {
		up := &allUp.Up[i]
		upgradeID := len(allUp.Up) - i

		// check if user wants to exclude specific things (true) or include specific things
		if noIncludes {
			// exclude repositories mentioned by the user
			if otherExclude.Contains(up.Repository) {
				u.log.Debugln("pruning", up.Name)
				excluded = append(excluded, graph.Prune(up.Name)...)
			}
			// exclude packages mentioned by the user
			if exclude.Get(upgradeID) {
				u.log.Debugln("pruning", up.Name)
				excluded = append(excluded, graph.Prune(up.Name)...)
			}

			// If the user explicitly wants to include a package/repository, exclude everything else
		} else if !include.Get(upgradeID) && !otherInclude.Contains(up.Repository) {
			u.log.Debugln("pruning", up.Name)
			excluded = append(excluded, graph.Prune(up.Name)...)
		}
	}

	return excluded, nil
}

func upgradeSelectEvent(allUp UpSlice) *settingslua.UpgradeSelectEvent {
	return &settingslua.UpgradeSelectEvent{
		Upgrades:           upgradeSelectPackages(allUp.Up, true),
		PulledDependencies: upgradeSelectPackages(allUp.PulledDeps, false),
	}
}

func upgradeSelectPackages(upgrades []Upgrade, selectable bool) []settingslua.UpgradeSelectPackage {
	packages := make([]settingslua.UpgradeSelectPackage, 0, len(upgrades))
	for i := range upgrades {
		up := &upgrades[i]
		id := 0
		if selectable {
			id = len(upgrades) - i
		}

		packages = append(packages, settingslua.UpgradeSelectPackage{
			ID:            id,
			Name:          up.Name,
			Base:          up.Base,
			Repository:    up.Repository,
			LocalVersion:  up.LocalVersion,
			RemoteVersion: up.RemoteVersion,
			Reason:        upgradeSelectReason(up.Reason),
			LastModified:  up.LastModified,
			Maintainer:    up.Maintainer,
		})
	}

	return packages
}

func upgradeSelectReason(reason alpm.PkgReason) string {
	switch reason {
	case alpm.PkgReasonExplicit:
		return "explicit"
	case alpm.PkgReasonDepend:
		return "dependency"
	default:
		return "unknown"
	}
}