refactor(main.go): optimized and cleaned up code

- Added option to select working mode.
- Improved flag parsing using pflag.
- Created separate functions for each mode.

4 files changed, 144 insertions, 134 deletions

diff --git a/README.md b/README.md
index c761336..5c78c9a 100644
--- a/README.md
+++ b/README.md
@@ -32,10 +32,9 @@ You can run the applications with or without providing any options. Running it w
 ```
 Options:
   -d, --daemon       Run the application in daemon mode
+  -m, --mode         Specify the mode (temp, usage or custom)
   -f, --file         Specify the CSV file containing the pattern data (This is required in daemon mode)
-  -t, --temperature  Display the CPU temperature
   -c, --celcius      Display the CPU temperature in celcius
-  -u, --usage        Display the CPU usage
 
 Commands:
   -h, --help         Print help
@@ -45,10 +44,10 @@ Commands:
 Run the application in daemon mode to display the pattern from a CSV file:
 
 ```bash
-./deepcool-display-linux -d -f /path/to/pattern.csv
+./deepcool-display-linux -d -m custom -f /path/to/pattern.csv
 ```
 
-In daemon mode, the \`-f\` or \`--file\` flag is required to specify the CSV file containing the pattern.
+In custom mode, the \`-f\` or \`--file\` flag is required to specify the CSV file containing the pattern.
 
 ### Exporting patterns to CSV
 
diff --git a/go.mod b/go.mod
index d8b8127..4602ac5 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ toolchain go1.23.4
 
 require (
 	github.com/google/gousb v1.1.3
+	github.com/spf13/pflag v1.0.5
 	github.com/wailsapp/wails/v2 v2.9.2
 )
 
diff --git a/go.sum b/go.sum
index 4508a36..55a6b15 100644
--- a/go.sum
+++ b/go.sum
@@ -48,6 +48,8 @@ github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
 github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/samber/lo v1.38.1 h1:j2XEAqXKb09Am4ebOg31SpvzUTTs6EN3VfgeLUhPdXM=
 github.com/samber/lo v1.38.1/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
diff --git a/main.go b/main.go
index 213aed6..b6954a3 100644
--- a/main.go
+++ b/main.go
@@ -4,8 +4,8 @@ import (
 	"context"
 	"deepcool-display-linux/modules"
 	"embed"
-	"flag"
 	"fmt"
+	flag "github.com/spf13/pflag"
 	"github.com/wailsapp/wails/v2"
 	"github.com/wailsapp/wails/v2/pkg/options"
 	"github.com/wailsapp/wails/v2/pkg/options/assetserver"
@@ -16,16 +16,17 @@ import (
 //go:embed all:frontend/dist
 var assets embed.FS
 
+const (
+	version = "0.1.0"
+)
+
 var (
-	daemonFlag  bool
-	err         error
-	filename    string
-	tempFlag    bool
-	celsiusFlag bool
-	usageFlag   bool
-	helpFlag    bool
-	versionFlag bool
-	version     = "0.1.0"
+	daemonFlag  = flag.BoolP("daemon", "d", false, "Run in daemon mode")
+	modeFlag    = flag.StringP("mode", "m", "", "Specify the mode (temp, usage or custom)")
+	celsiusFlag = flag.BoolP("celsius", "c", false, "Show CPU temperature in Celsius")
+	helpFlag    = flag.BoolP("help", "h", false, "Show help message")
+	versionFlag = flag.BoolP("version", "v", false, "Show app version")
+	filename    = flag.StringP("file", "f", "", "Specify CSV file path for pattern")
 )
 
 func printHelp() {
@@ -37,10 +38,9 @@ Usage:
 Options:
     -h, --help     Show this help message
     -d, --daemon   Run in daemon mode
+    -m, --mode     Specify the mode (temp, usage or custom)
     -f, --file     Specify CSV file path for pattern
-    -t, --temp     Show CPU temperature
     -c, --celsius  Show CPU temperature in Celsius
-    -u, --usage    Show CPU usage
     -v, --version  Show the version of the app
 
 Modes:
@@ -56,136 +56,144 @@ For more information, visit: https://github.com/blaster4385/deepcool-display-lin
 `, os.Args[0], os.Args[0], os.Args[0])
 }
 
+func handleError(err error) {
+	fmt.Printf("Error: %v\n", err)
+	os.Exit(1)
+}
+
 func main() {
-	flag.BoolVar(&daemonFlag, "daemon", false, "Run as daemon")
-	flag.BoolVar(&daemonFlag, "d", false, "Run as daemon")
-	flag.StringVar(&filename, "file", "", "CSV file")
-	flag.StringVar(&filename, "f", "", "CSV file")
-	flag.BoolVar(&tempFlag, "temp", false, "Show CPU temperature")
-	flag.BoolVar(&tempFlag, "t", false, "Show CPU temperature")
-	flag.BoolVar(&celsiusFlag, "celsius", false, "Show CPU temperature in Celsius")
-	flag.BoolVar(&celsiusFlag, "c", false, "Show CPU temperature in Celsius")
-	flag.BoolVar(&usageFlag, "usage", false, "Show CPU usage")
-	flag.BoolVar(&usageFlag, "u", false, "Show CPU usage")
-	flag.BoolVar(&helpFlag, "help", false, "Show help message")
-	flag.BoolVar(&helpFlag, "h", false, "Show help message")
-	flag.BoolVar(&versionFlag, "version", false, "Show app version")
-	flag.BoolVar(&versionFlag, "v", false, "Show app version")
 	flag.Parse()
 
-	if helpFlag {
+	if *helpFlag {
 		printHelp()
-		os.Exit(0)
+		return
 	}
-
-	if versionFlag {
+	if *versionFlag {
 		fmt.Printf("Version: %s\n", version)
-		os.Exit(0)
+		return
 	}
 
 	app := NewApp()
-	if daemonFlag {
-		ctx := context.Background()
-		app.startup(ctx)
-		if filename != "" {
-			grid, err := app.ParseCSV(filename)
-			if err != nil {
-				fmt.Printf("Error parsing CSV file: %v\n", err)
-				os.Exit(1)
-			}
-			err = app.SendPattern(grid)
-			if err != nil {
-				fmt.Printf("Error sending pattern: %v\n", err)
-				os.Exit(1)
-			}
-		} else if tempFlag {
-			ticker := time.NewTicker(3 * time.Second)
-			defer ticker.Stop()
-
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-
-			go func() {
-				for {
-					select {
-					case <-ctx.Done():
-						return
-					case <-ticker.C:
-						var temp float64
-						if celsiusFlag {
-							temp, err = modules.GetCPUTemperature(false)
-							grid, err := modules.CreateNumberGrid(int(temp), "celsius", 5)
-							if err != nil {
-								fmt.Printf("Error creating number grid: %v\n", err)
-								os.Exit(1)
-							}
-							err = app.SendPattern(grid)
-							if err != nil {
-								fmt.Printf("Error sending pattern: %v\n", err)
-								os.Exit(1)
-							}
-						} else {
-							temp, err = modules.GetCPUTemperature(true)
-							grid, err := modules.CreateNumberGrid(int(temp), "fahrenheit", 5)
-							if err != nil {
-								fmt.Printf("Error creating number grid: %v\n", err)
-								os.Exit(1)
-							}
-							err = app.SendPattern(grid)
-						}
-					}
+
+	if *daemonFlag {
+		daemonMode(app)
+	} else {
+		guiMode(app)
+	}
+}
+
+func daemonMode(app *App) {
+	contextBackground, contextCancel := context.WithCancel(context.Background())
+	defer contextCancel()
+
+	if *modeFlag == "custom" {
+		if *filename == "" {
+			handleError(fmt.Errorf("filename is required for custom mode"))
+		}
+		handleDaemonPattern(app, *filename, contextBackground)
+		return
+	}
+	if *modeFlag == "temp" {
+		daemonTemperatureDisplay(app, contextBackground)
+		return
+	}
+	if *modeFlag == "usage" {
+		daemonUsageDisplay(app, contextBackground)
+		return
+	}
+
+	select {}
+}
+
+func handleDaemonPattern(app *App, filename string, ctx context.Context) {
+	grid, err := app.ParseCSV(filename)
+	if err != nil {
+		handleError(fmt.Errorf("parsing CSV file: %w", err))
+	}
+	err = app.SendPattern(grid)
+	if err != nil {
+		handleError(fmt.Errorf("sending pattern: %w", err))
+	}
+	select {}
+}
+
+func daemonTemperatureDisplay(app *App, ctx context.Context) {
+	ticker := time.NewTicker(3 * time.Second)
+	defer ticker.Stop()
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				temp, err := modules.GetCPUTemperature(!*celsiusFlag)
+				if err != nil {
+					handleError(fmt.Errorf("getting temperature: %w", err))
+				}
+				var tempSymbol string
+				if *celsiusFlag {
+					tempSymbol = "celsius"
+				} else {
+					tempSymbol = "fahrenheit"
 				}
-			}()
-		} else if usageFlag {
-			ticker := time.NewTicker(3 * time.Second)
-			defer ticker.Stop()
-
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-
-			go func() {
-				for {
-					select {
-					case <-ctx.Done():
-						return
-					case <-ticker.C:
-						usage, err := modules.GetCPUUsage()
-						if err != nil {
-							fmt.Printf("Error getting CPU usage: %v\n", err)
-							os.Exit(1)
-						}
-						grid, err := modules.CreateNumberGrid(int(usage), "percent", 5)
-						if err != nil {
-							fmt.Printf("Error creating number grid: %v\n", err)
-							os.Exit(1)
-						}
-						err = app.SendPattern(grid)
-						if err != nil {
-							fmt.Printf("Error sending pattern: %v\n", err)
-							os.Exit(1)
-						}
-					}
+				grid, err := modules.CreateNumberGrid(int(temp), tempSymbol, 5)
+				if err != nil {
+					handleError(fmt.Errorf("creating temperature grid: %w", err))
 				}
-			}()
+				err = app.SendPattern(grid)
+				if err != nil {
+					handleError(fmt.Errorf("sending pattern: %w", err))
+				}
+			}
 		}
-		select {}
-	} else {
-		err := wails.Run(&options.App{
-			Title:  "deepcool-display-linux",
-			Width:  1024,
-			Height: 768,
-			AssetServer: &assetserver.Options{
-				Assets: assets,
-			},
-			BackgroundColour: &options.RGBA{R: 40, G: 40, B: 40, A: 1},
-			OnStartup:        app.startup,
-			Bind: []interface{}{
-				app,
-			},
-		})
-		if err != nil {
-			fmt.Printf("Error starting application: %v\n", err)
-			os.Exit(1)
+	}()
+	select {}
+}
+
+func daemonUsageDisplay(app *App, ctx context.Context) {
+	ticker := time.NewTicker(3 * time.Second)
+	defer ticker.Stop()
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				usage, err := modules.GetCPUUsage()
+				if err != nil {
+					handleError(fmt.Errorf("getting CPU usage: %w", err))
+				}
+				grid, err := modules.CreateNumberGrid(int(usage), "percent", 5)
+				if err != nil {
+					handleError(fmt.Errorf("creating usage grid: %w", err))
+				}
+				err = app.SendPattern(grid)
+				if err != nil {
+					handleError(fmt.Errorf("sending pattern: %w", err))
+				}
+			}
 		}
+	}()
+	select {}
+}
+
+func guiMode(app *App) {
+	err := wails.Run(&options.App{
+		Title:  "deepcool-display-linux",
+		Width:  1024,
+		Height: 768,
+		AssetServer: &assetserver.Options{
+			Assets: assets,
+		},
+		BackgroundColour: &options.RGBA{R: 40, G: 40, B: 40, A: 1},
+		OnStartup:        app.startup,
+		Bind: []interface{}{
+			app,
+		},
+	})
+	if err != nil {
+		handleError(fmt.Errorf("starting application: %w", err))
 	}
 }