From 9ecbbff6faf2fd56273b84d4f3a0613d3ba1613d Mon Sep 17 00:00:00 2001
From: PhatPhuckDave <david.majdandzic@hotmail.com>
Date: Sun, 20 Jul 2025 11:47:11 +0200
Subject: [PATCH] Implement special flags for dump and reset db

---
 main.go       | 31 +++++++++++++++++++++++++++++++
 utils/db.go   | 25 +++++++++++++++++++++++++
 utils/file.go | 20 ++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/main.go b/main.go
index 40db827..6ba8230 100644
--- a/main.go
+++ b/main.go
@@ -62,6 +62,15 @@ func main() {
 		return
 	}
 
+	workdone, err := HandleSpecialArgs(args, err, db)
+	if err != nil {
+		logger.Error("Failed to handle special args: %v", err)
+		return
+	}
+	if workdone {
+		return
+	}
+
 	// The plan is:
 	// Load all commands
 	commands, err := utils.LoadCommands(args)
@@ -268,6 +277,28 @@ func main() {
 	}
 }
 
+func HandleSpecialArgs(args []string, err error, db utils.DB) (bool, error) {
+	switch args[0] {
+	case "reset":
+		err = utils.ResetAllFiles(db)
+		if err != nil {
+			logger.Error("Failed to reset all files: %v", err)
+			return true, err
+		}
+		logger.Info("All files reset")
+		return true, nil
+	case "dump":
+		err = db.RemoveAllFiles()
+		if err != nil {
+			logger.Error("Failed to remove all files from database: %v", err)
+			return true, err
+		}
+		logger.Info("All files removed from database")
+		return true, nil
+	}
+	return false, nil
+}
+
 func CreateExampleConfig() {
 	commands := []utils.ModifyCommand{
 		{
diff --git a/utils/db.go b/utils/db.go
index 29ede0d..1c67b1e 100644
--- a/utils/db.go
+++ b/utils/db.go
@@ -15,6 +15,8 @@ type DB interface {
 	Raw(sql string, args ...any) *gorm.DB
 	SaveFile(filePath string, fileData []byte) error
 	GetFile(filePath string) ([]byte, error)
+	GetAllFiles() ([]FileSnapshot, error)
+	RemoveAllFiles() error
 }
 
 type FileSnapshot struct {
@@ -93,3 +95,26 @@ func (db *DBWrapper) GetFile(filePath string) ([]byte, error) {
 	log.Debug("File found in database")
 	return fileSnapshot.FileData, nil
 }
+
+func (db *DBWrapper) GetAllFiles() ([]FileSnapshot, error) {
+	log := cylogger.Default.WithPrefix("GetAllFiles")
+	log.Debug("Getting all files from database")
+	var fileSnapshots []FileSnapshot
+	err := db.db.Model(&FileSnapshot{}).Find(&fileSnapshots).Error
+	if err != nil {
+		return nil, err
+	}
+	log.Debug("Found %d files in database", len(fileSnapshots))
+	return fileSnapshots, nil
+}
+
+func (db *DBWrapper) RemoveAllFiles() error {
+	log := cylogger.Default.WithPrefix("RemoveAllFiles")
+	log.Debug("Removing all files from database")
+	err := db.db.Exec("DELETE FROM file_snapshots").Error
+	if err != nil {
+		return err
+	}
+	log.Debug("All files removed from database")
+	return nil
+}
diff --git a/utils/file.go b/utils/file.go
index b576537..da21c2f 100644
--- a/utils/file.go
+++ b/utils/file.go
@@ -71,3 +71,23 @@ func ResetWhereNecessary(associations map[string]FileCommandAssociation, db DB)
 	log.Debug("Done")
 	return nil
 }
+
+func ResetAllFiles(db DB) error {
+	log := cylogger.Default.WithPrefix("ResetAllFiles")
+	log.Debug("Start")
+	fileSnapshots, err := db.GetAllFiles()
+	if err != nil {
+		return err
+	}
+	log.Debug("Found %d files in database", len(fileSnapshots))
+	for _, fileSnapshot := range fileSnapshots {
+		log.Debug("Resetting file %q", fileSnapshot.FilePath)
+		err = os.WriteFile(fileSnapshot.FilePath, fileSnapshot.FileData, 0644)
+		if err != nil {
+			return err
+		}
+		log.Debug("File %q written to disk", fileSnapshot.FilePath)
+	}
+	log.Debug("Done")
+	return nil
+}