PackageManager/Client/Filesystem/PickFiles.go

package Filesystem

import (
	"os"
	"path/filepath"
	"strings"

	"github.com/vbauerster/mpb"
	bolt "go.etcd.io/bbolt"

	"PackageManager/Client/Database"
	"PackageManager/Client/ProgressBar"
	"PackageManager/Variables"
)

func pickFilesSingle(fileKey, filePath string) error {
	var (
		indexBucket *bolt.Bucket
		picksBucket *bolt.Bucket
		e           error
	)

	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		indexBucket = tx.Bucket(Variables.FsHashIndexBucket)
		picksBucket = tx.Bucket(Variables.FsHashPicksBucket)

		e = AddFileToBucket(picksBucket, fileKey, filePath)
		if e != nil {
			return e
		}
		return RemoveFileFromBucket(indexBucket, fileKey)
	})
	return e
}

func pickFilesRecursive(rootPath string) error {
	var (
		fsStatus    FilesystemStatus
		indexBucket *bolt.Bucket
		picksBucket *bolt.Bucket
		bar         *mpb.Bar
		totalLen    int
		trimF       string
		f           string
		e           error
	)

	fsStatus, e = GetFilesystemDiff(rootPath)
	if e != nil {
		return e
	}

	totalLen = len(fsStatus.NewFiles) + len(fsStatus.ModifiedFiles) + len(fsStatus.MissingFiles)

	if totalLen == 0 {
		return nil
	}

	bar = ProgressBar.InitBar("Adding...", totalLen)

	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		indexBucket = tx.Bucket(Variables.FsHashIndexBucket)
		picksBucket = tx.Bucket(Variables.FsHashPicksBucket)

		if len(fsStatus.NewFiles) > 0 {
			for _, f = range fsStatus.NewFiles {
				bar.Increment()
				trimF = strings.TrimPrefix(f, Variables.RootDir)
				e = AddFileToBucket(picksBucket, trimF, f)
				if e != nil {
					return e
				}
			}
		}

		if len(fsStatus.ModifiedFiles) > 0 {
			for _, f = range fsStatus.ModifiedFiles {
				bar.Increment()
				trimF = strings.TrimPrefix(f, Variables.RootDir)
				e = AddFileToBucket(picksBucket, trimF, f)
				if e != nil {
					return e
				}
			}
		}

		if len(fsStatus.MissingFiles) > 0 {
			for _, f = range fsStatus.MissingFiles {
				bar.Increment()
				e = RemoveFileFromBucket(indexBucket, f)
				if e != nil {
					return e
				}
				e = RemoveFileFromBucket(picksBucket, f)
				if e != nil {
					return e
				}
			}
		}

		return nil
	})

	return e
}

func PickFiles(rootPath string) error {
	var (
		realRootPath string
		rootStat     os.FileInfo
		e            error
	)

	realRootPath = filepath.Join(Variables.RootDir, rootPath)

	rootStat, e = os.Stat(realRootPath)
	if e != nil {
		return e
	}

	if !rootStat.IsDir() {
		return pickFilesSingle(rootPath, realRootPath)
	}

	return pickFilesRecursive(realRootPath)
}

func ResetAllPickedFiles() error {
	var (
		e error
	)

	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		return tx.DeleteBucket(Variables.FsHashPicksBucket)
	})

	return e
}