From b9e38575bccebf7a1b1f1a1ec81acc46daadd571 Mon Sep 17 00:00:00 2001
From: Pijus Kamandulis <pkpjuklas@gmail.com>
Date: Tue, 27 Feb 2024 22:11:33 +0200
Subject: [PATCH] Load state from '-Persist' path if '-InitialData' not
 supplied

---
 README.md                      |  6 ++----
 internal/repositories/state.go | 25 +++++++++++++++++++++++++
 main.go                        |  4 +---
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 05024bd..7005273 100644
--- a/README.md
+++ b/README.md
@@ -23,9 +23,7 @@ Cosmium is available for the following platforms:
 Once downloaded, you can launch Cosmium using the following command:
 
 ```sh
-./cosmium-linux-amd64 \
-  -Persist "./save.json" \
-  -InitialData "./save.json"
+cosmium -Persist "./save.json"
 ```
 
 Connection String Example:
@@ -51,7 +49,7 @@ To disable SSL and run Cosmium on HTTP instead, you can use the `-DisableTls` fl
 * **-DisableAuth**: Disable authentication
 * **-Host**: Hostname (default "localhost")
 * **-InitialData**: Path to JSON containing initial state
-* **-Persist**: Saves data to the given path on application exit
+* **-Persist**: Saves data to the given path on application exit (When `-InitialData` argument is not supplied, it will try to load data from path supplied in `-Persist`)
 * **-Port**: Listen port (default 8081)
 
 These arguments allow you to configure various aspects of Cosmium's behavior according to your requirements.
diff --git a/internal/repositories/state.go b/internal/repositories/state.go
index d92e663..4b62503 100644
--- a/internal/repositories/state.go
+++ b/internal/repositories/state.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"reflect"
 
+	"github.com/pikami/cosmium/api/config"
 	repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 )
 
@@ -19,15 +20,39 @@ var storeState = repositorymodels.State{
 	Documents:   make(map[string]map[string]map[string]repositorymodels.Document),
 }
 
+func InitializeRepository() {
+	if config.Config.InitialDataFilePath != "" {
+		LoadStateFS(config.Config.InitialDataFilePath)
+		return
+	}
+
+	if config.Config.PersistDataFilePath != "" {
+		stat, err := os.Stat(config.Config.PersistDataFilePath)
+		if err != nil {
+			return
+		}
+
+		if stat.IsDir() {
+			fmt.Println("Argument '-Persist' must be a path to file, not a directory.")
+			os.Exit(1)
+		}
+
+		LoadStateFS(config.Config.PersistDataFilePath)
+		return
+	}
+}
+
 func LoadStateFS(filePath string) {
 	data, err := os.ReadFile(filePath)
 	if err != nil {
 		log.Fatalf("Error reading state JSON file: %v", err)
+		return
 	}
 
 	var state repositorymodels.State
 	if err := json.Unmarshal(data, &state); err != nil {
 		log.Fatalf("Error unmarshalling state JSON: %v", err)
+		return
 	}
 
 	fmt.Println("Loaded state:")
diff --git a/main.go b/main.go
index 1274fd7..178db53 100644
--- a/main.go
+++ b/main.go
@@ -13,9 +13,7 @@ import (
 func main() {
 	config.ParseFlags()
 
-	if config.Config.InitialDataFilePath != "" {
-		repositories.LoadStateFS(config.Config.InitialDataFilePath)
-	}
+	repositories.InitializeRepository()
 
 	go api.StartAPI()