brunel/packages/packages.go

package packages

import (
	"brunel/config"
	"brunel/deb"
	"brunel/domain"
	"brunel/helpers"
	"compress/bzip2"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"slices"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/errgroup"
	"pault.ag/go/debian/version"

	"github.com/alphadose/haxmap"
	"github.com/klauspost/compress/gzip"
	"github.com/ulikunitz/xz"
)

// Constants for frequently used strings
const (
	DMOSuffix              = "-dmo"
	DebianInstallerSection = "debian-installer"
	PackageKey             = "Package"
	SourceKey              = "Source"
	VersionKey             = "Version"
	ArchitectureKey        = "Architecture"
	DescriptionKey         = "Description"
	SectionKey             = "Section"
)

var LastUpdateTime time.Time
var currentPackages = haxmap.New[string, domain.SourcePackage]()

func ProcessPackages() error {
	start := time.Now()
	internalPackages := haxmap.New[string, domain.SourcePackage]()
	externalPackages := haxmap.New[string, domain.SourcePackage]()

	var eg errgroup.Group
	eg.Go(func() error {
		return LoadPackages(internalPackages, config.Configs.LocalPackageFiles, true)
	})

	eg.Go(func() error {
		return LoadPackages(externalPackages, config.Configs.ExternalPackageFiles, false)
	})

	if err := eg.Wait(); err != nil {
		return err
	}
	slog.Info("packages loaded in " + time.Since(start).String())

	// Merge internal and external packages
	mergePackages(internalPackages, externalPackages)

	// **Reintroduce combinePackages here**
	combinePackages(internalPackages)
	combinePackages(externalPackages)

	// Determine which packages have been updated
	updatedPackages := determineUpdatedPackages(internalPackages, externalPackages)

	// Update binary package flags based on updated packages
	updateBinaryPackageFlags(updatedPackages)

	updatedPackages.ForEach(func(key string, pkg domain.SourcePackage) bool {
		internalPackages.Set(key, pkg)
		return true
	})

	currentPackages = internalPackages

	LastUpdateTime = time.Now()
	if err := helpers.DBInst.DropPackages(); err != nil {
		return fmt.Errorf("dropping packages: %w", err)
	}
	if err := SaveToDb(); err != nil {
		return fmt.Errorf("saving to database: %w", err)
	}

	slog.Info("packages processed and saved in " + time.Since(start).String())
	return nil
}

func GetPackages() *haxmap.Map[string, domain.SourcePackage] {
	return currentPackages
}

func UpdatePackage(pkg domain.PackageInfo) error {
	curr, ok := currentPackages.Get(pkg.Source)
	if !ok {
		return fmt.Errorf("package %s not found", pkg.Source)
	}
	curr.Packages.Set(pkg.PackageName, pkg)
	currentPackages.Set(pkg.Source, curr)
	return saveSingleToDb(curr)
}

func UpdateSourcePackage(pkg domain.SourcePackage) error {
	currentPackages.Set(pkg.Name, pkg)
	return saveSingleToDb(pkg)
}

func saveSingleToDb(pkg domain.SourcePackage) error {
	if err := helpers.DBInst.UpdatePackage(pkg); err != nil {
		return fmt.Errorf("updating package: %w", err)
	}
	LastUpdateTime = time.Now()
	if err := helpers.DBInst.UpdateLastUpdateTime(LastUpdateTime); err != nil {
		return fmt.Errorf("updating last update time: %w", err)
	}
	return nil
}

func SaveToDb() error {
	if err := helpers.DBInst.SavePackages(currentPackages); err != nil {
		slog.Error("Error saving packages to database", "error", err)
		return fmt.Errorf("saving packages to database: %w", err)
	}
	LastUpdateTime = time.Now()
	return helpers.DBInst.UpdateLastUpdateTime(LastUpdateTime)
}

func LoadFromDb() error {
	packages, err := helpers.DBInst.GetPackages()
	if err != nil {
		slog.Error("Error getting packages from database", "error", err)
		return nil
	}
	slices.SortStableFunc(packages, func(a, b domain.SourcePackage) int {
		if a.Name == b.Name {
			return 0
		}
		if a.Name > b.Name {
			return 1
		}
		return -1
	})
	currentPackages.Clear()
	for _, pkg := range packages {
		currentPackages.Set(pkg.Name, pkg)
	}
	return nil
}

func LoadPackages(packages *haxmap.Map[string, domain.SourcePackage], packageFiles []config.PackageFile, isInternal bool) error {
	sortedFiles := sortPackageFiles(packageFiles, isInternal)
	var eg errgroup.Group
	var mu sync.Mutex
	packageResults := make([][]fetchResult, len(sortedFiles))

	for i, pkg := range sortedFiles {
		i, pkg := i, pkg // capture range variables
		eg.Go(func() error {
			results, err := fetchPackagesForFile(pkg)
			if err != nil {
				return err
			}
			mu.Lock()
			packageResults[i] = results
			mu.Unlock()
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return fmt.Errorf("fetching package files: %w", err)
	}

	for _, results := range packageResults {
		for _, result := range results {
			processPackageResult(packages, result)
		}
	}

	return nil
}

func sortPackageFiles(packageFiles []config.PackageFile, isInternal bool) []config.PackageFile {
	sorted := make([]config.PackageFile, len(packageFiles))
	copy(sorted, packageFiles)
	slices.SortStableFunc(sorted, func(a, b config.PackageFile) int {
		if a.Priority == b.Priority {
			return 0
		}
		if (isInternal && a.Priority < b.Priority) || (!isInternal && a.Priority > b.Priority) {
			return 1
		}
		return -1
	})
	return sorted
}

type fetchResult struct {
	repo     string
	packages *haxmap.Map[string, domain.PackageInfo]
}

func fetchPackagesForFile(pkg config.PackageFile) ([]fetchResult, error) {
	var results []fetchResult
	for _, repo := range pkg.Subrepos {
		packages, err := fetchPackageFile(pkg, repo)
		if err != nil {
			return nil, fmt.Errorf("fetching package file for repo %s: %w", repo, err)
		}
		results = append(results, fetchResult{repo: repo, packages: packages})
	}
	return results, nil
}

func processPackageResult(packages *haxmap.Map[string, domain.SourcePackage], result fetchResult) {
	result.packages.ForEach(func(newKey string, newPkg domain.PackageInfo) bool {
		pk, ok := packages.Get(newPkg.Source)
		if !ok {
			newMap := haxmap.New[string, domain.PackageInfo]()
			newMap.Set(newKey, newPkg)
			packages.Set(newPkg.Source, domain.SourcePackage{
				Name:     newPkg.Source,
				Packages: newMap,
				Status:   domain.Current,
				Version:  newPkg.Version,
			})
			return true
		}

		pk.Version = getHighestVer(pk.Version, newPkg.Version)
		pkg, ok := pk.Packages.Get(newKey)
		if !ok {
			pk.Packages.Set(newKey, newPkg)
		} else {
			mVer, _ := version.Parse(pkg.Version)
			extVer, _ := version.Parse(strings.Split(newPkg.Version, "+b")[0])
			cmpVal := version.Compare(extVer, mVer)
			if cmpVal >= 0 {
				pk.Packages.Set(newKey, newPkg)
			}
		}

		packages.Set(newPkg.Source, pk)
		return true
	})
}

func ProcessMissingPackages(internalPackages *haxmap.Map[string, domain.SourcePackage], externalPackages *haxmap.Map[string, domain.SourcePackage]) {
	externalPackages.ForEach(func(k string, src domain.SourcePackage) bool {
		_, ok := internalPackages.Get(k)
		if !ok && src.Packages.Len() > 0 {
			src.Status = domain.Missing
			internalPackages.Set(k, src)
		}
		return true
	})
}

func ProcessStalePackages(internalPackages *haxmap.Map[string, domain.SourcePackage], externalPackages *haxmap.Map[string, domain.SourcePackage]) {
	externalPackages.ForEach(func(newPackage string, newSource domain.SourcePackage) bool {
		matchedPackage, ok := internalPackages.Get(newPackage)
		if !ok || matchedPackage.Packages.Len() == 0 {
			return true
		}
		ver := getHighestVer(newSource.Version, matchedPackage.Version)
		if ver != matchedPackage.Version {
			matchedPackage.NewVersion = ver
			matchedPackage.Status = domain.Stale
			internalPackages.Set(newPackage, matchedPackage)
		}

		return true
	})
}

func fetchPackageFile(pkg config.PackageFile, selectedRepo string) (*haxmap.Map[string, domain.PackageInfo], error) {
	url := fmt.Sprintf("%s%s/%s.%s?nocache=%d", pkg.Url, selectedRepo, pkg.Packagepath, pkg.Compression, time.Now().Unix())
	resp, err := http.Get(url)
	if err != nil {
		return nil, fmt.Errorf("fetching package file: %w", err)
	}
	defer resp.Body.Close()

	rdr, err := getCompressedReader(resp.Body, pkg.Compression)
	if err != nil {
		return nil, fmt.Errorf("creating compressed reader: %w", err)
	}

	packages := haxmap.New[string, domain.PackageInfo]()
	sreader := deb.NewControlFileReader(rdr, false, false)

	for {
		stanza, err := sreader.ReadStanza()
		if err != nil {
			if errors.Is(err, deb.ErrMalformedStanza) {
				slog.Warn("Malformed stanza encountered", "error", err)
				continue
			}
			return nil, err
		}
		if stanza == nil {
			break
		}

		if stanza[SectionKey] == DebianInstallerSection {
			continue
		}

		name := stanza[PackageKey]
		if shouldSkipPackage(name, pkg) {
			continue
		}

		source, sourceVersion := parseSource(stanza[SourceKey])
		if source == "" {
			source = name
		}

		versionStr := chooseVersion(sourceVersion, stanza[VersionKey])
		ver, err := version.Parse(versionStr)
		if err != nil {
			slog.Warn("Invalid version format", "version", versionStr, "error", err)
			continue
		}

		existingPkg, exists := packages.Get(name)
		if !exists {
			// If the package doesn't exist, add it
			packages.Set(name, domain.PackageInfo{
				PackageName:  name,
				Version:      ver.String(),
				Source:       source,
				Architecture: stanza[ArchitectureKey],
				Description:  stanza[DescriptionKey],
				Status:       domain.Current,
			})
		} else {
			// Compare existing version with the new version
			cmp, err := compareVersions(ver.String(), existingPkg.Version)
			if err != nil {
				slog.Warn("Version comparison failed", "package", name, "error", err)
				continue
			}
			if cmp >= 0 {
				// Update the package if the new version is higher or equal
				packages.Set(name, domain.PackageInfo{
					PackageName:  name,
					Version:      ver.String(),
					Source:       source,
					Architecture: stanza[ArchitectureKey],
					Description:  stanza[DescriptionKey],
					Status:       existingPkg.Status,
				})
			}
		}
	}

	return packages, nil
}

func getCompressedReader(body io.Reader, compression string) (io.Reader, error) {
	switch compression {
	case "bz2":
		return bzip2.NewReader(body), nil
	case "xz":
		return xz.NewReader(body)
	case "gz":
		return gzip.NewReader(body)
	default:
		return body, nil
	}
}

func shouldSkipPackage(name string, pkg config.PackageFile) bool {
	if pkg.UseWhitelist && len(pkg.Whitelist) > 0 && !nameContains(name, pkg.Whitelist) {
		return true
	}
	return nameContains(name, pkg.Blacklist)
}

func parseSource(source string) (string, string) {
	sourceSplit := strings.Split(source, " ")
	if len(sourceSplit) > 1 {
		return sourceSplit[0], strings.Trim(sourceSplit[1], "()")
	}
	return source, ""
}

func chooseVersion(sourceVersion, stanzaVersion string) string {
	if sourceVersion != "" {
		return sourceVersion
	}
	return stanzaVersion
}

func nameContains(name string, match []string) bool {
	for _, m := range match {
		if strings.Contains(name, m) {
			return true
		}
	}
	return false
}

func GetPackagesCount() domain.PackagesCount {
	count := domain.PackagesCount{
		Stale:    0,
		Missing:  0,
		Current:  0,
		Error:    0,
		Queued:   0,
		Building: 0,
	}
	currentPackages.ForEach(func(k string, pkg domain.SourcePackage) bool {
		switch pkg.Status {
		case domain.Stale:
			count.Stale++
		case domain.Missing:
			count.Missing++
		case domain.Built:
			count.Current++
		case domain.Current:
			count.Current++
		case domain.Error:
			count.Error++
		}

		return true
	})

	return count
}

func mergePackages(internal, external *haxmap.Map[string, domain.SourcePackage]) {
	external.ForEach(func(name string, extPkg domain.SourcePackage) bool {
		intPkg, exists := internal.Get(name)
		if exists {
			// Compare internal package version with external package version
			cmp, err := compareVersions(intPkg.Version, extPkg.Version)
			if err != nil {
				slog.Warn("Version comparison failed in mergePackages",
					"package", name,
					"internal_version", intPkg.Version,
					"external_version", extPkg.Version,
					"error", err)
				// Decide whether to continue or handle the error differently
				return true // Continue with next package
			}

			if cmp < 0 {
				// External package has a newer version
				intPkg.NewVersion = extPkg.Version
				intPkg.Status = domain.Stale
				// Merge binary packages from external to internal
				extPkg.Packages.ForEach(func(binName string, binPkg domain.PackageInfo) bool {
					intPkg.Packages.Set(binName, binPkg)
					return true
				})
				internal.Set(name, intPkg)
			} else {
				// Internal package has the same or newer version
				// Optionally, you might still want to merge binary packages
				extPkg.Packages.ForEach(func(binName string, binPkg domain.PackageInfo) bool {
					// Only add binary package if it doesn't exist or has a newer version
					existingBinPkg, binExists := intPkg.Packages.Get(binName)
					if !binExists || binPkg.Version > existingBinPkg.Version {
						intPkg.Packages.Set(binName, binPkg)
					}
					return true
				})
				internal.Set(name, intPkg)
			}
		} else {
			// Package exists only externally; mark as missing
			extPkg.Status = domain.Missing
			internal.Set(name, extPkg)
		}
		return true
	})
}

func determineUpdatedPackages(internal, external *haxmap.Map[string, domain.SourcePackage]) *haxmap.Map[string, domain.SourcePackage] {
	updated := haxmap.New[string, domain.SourcePackage]()

	// Now, update the packages that have newer versions
	external.ForEach(func(name string, extPkg domain.SourcePackage) bool {
		if intPkg, exists := internal.Get(name); exists {
			cmp, err := compareVersions(intPkg.Version, extPkg.Version)
			if err != nil {
				slog.Warn("Version comparison failed", "package", name, "v1", intPkg.Version, "v2", extPkg.Version, "error", err)
				return true
			}
			if cmp < 0 {
				intPkg.NewVersion = extPkg.Version
				intPkg.Status = domain.Stale
				updated.Set(name, intPkg)
			}
		} else {
			// If the package doesn't exist internally, add it
			updated.Set(name, extPkg)
		}
		return true
	})

	return updated
}

func compareVersions(v1, v2 string) (int, error) {
	parsedV1, err1 := version.Parse(v1)
	if err1 != nil {
		return 0, fmt.Errorf("invalid version '%s': %w", v1, err1)
	}
	parsedV2, err2 := version.Parse(v2)
	if err2 != nil {
		return 0, fmt.Errorf("invalid version '%s': %w", v2, err2)
	}
	return version.Compare(parsedV1, parsedV2), nil
}

func updateBinaryPackageFlags(pkgs *haxmap.Map[string, domain.SourcePackage]) {
	pkgs.ForEach(func(name string, pkg domain.SourcePackage) bool {
		for _, p := range config.Configs.I386List {
			if pkg.Name == p || pkg.Name == p+DMOSuffix {
				pkg.Has32bit = true
				pkgs.Set(name, pkg)
				return true
			}
		}
		pkg.Has32bit = false
		pkgs.Set(name, pkg)
		return true
	})
}

func combinePackages(packages *haxmap.Map[string, domain.SourcePackage]) {
	dmoPackages := haxmap.New[string, domain.SourcePackage]()

	// Identify and collect -dmo packages
	packages.ForEach(func(k string, v domain.SourcePackage) bool {
		if strings.HasSuffix(k, DMOSuffix) {
			dmoPackages.Set(k, v)
			packages.Del(k)
		}
		return true
	})

	// Combine -dmo packages with base packages or add them back
	dmoPackages.ForEach(func(dmoName string, dmoPkg domain.SourcePackage) bool {
		baseName := strings.TrimSuffix(dmoName, DMOSuffix)
		if basePkg, hasBase := packages.Get(baseName); hasBase {
			combinedPkg := combineDMOPackages(dmoPkg, basePkg)
			packages.Set(dmoName, combinedPkg)
			packages.Del(baseName)
		} else {
			packages.Set(dmoName, dmoPkg)
		}
		return true
	})
}

func combineDMOPackages(dmoPkg, basePkg domain.SourcePackage) domain.SourcePackage {
	combinedPkg := dmoPkg
	basePkg.Packages.ForEach(func(pkgName string, pkgInfo domain.PackageInfo) bool {
		if _, exists := combinedPkg.Packages.Get(pkgName); !exists {
			combinedPkg.Packages.Set(pkgName, pkgInfo)
		}
		return true
	})
	return combinedPkg
}

func getHighestVer(ver string, newVer string) string {
	mVer, _ := version.Parse(ver)
	extVer, _ := version.Parse(newVer)
	cmpVal := version.Compare(mVer, extVer)
	if cmpVal < 0 {
		return newVer
	}
	return ver
}