From fb2ae39b476f5d3cab2f8a069de88c63f63344be Mon Sep 17 00:00:00 2001
From: Alexander Lazarenko <kerblif@unprism.ru>
Date: Sat, 10 May 2025 13:12:17 +0300
Subject: [PATCH] Add support for tracking and persisting snapshot IDs

Introduce `currentSnapshotID` to track the active snapshot and persist its state in a file. This change ensures the current snapshot ID is restored during initialization and maintained consistently across snapshot operations.
---
 api.go | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 78 insertions(+), 11 deletions(-)

diff --git a/api.go b/api.go
index e08690b..1cb8651 100644
--- a/api.go
+++ b/api.go
@@ -36,10 +36,12 @@ type AgateOptions struct {
 
 // Agate is the main entry point for the snapshot library.
 type Agate struct {
-	manager     SnapshotManager
-	options     AgateOptions
-	metadataDir string
-	blobsDir    string
+	manager           SnapshotManager
+	options           AgateOptions
+	metadataDir       string
+	blobsDir          string
+	currentSnapshotID string
+	currentIDFile     string
 }
 
 // New initializes a new instance of the Agate library with the given options.
@@ -90,16 +92,38 @@ func New(options AgateOptions) (*Agate, error) {
 		return nil, fmt.Errorf("failed to create snapshot manager: %w", err)
 	}
 
-	return &Agate{
-		manager:     manager,
-		options:     options,
-		metadataDir: metadataDir,
-		blobsDir:    blobsDir,
-	}, nil
+	// Create a file path for storing the current snapshot ID
+	currentIDFile := filepath.Join(options.WorkDir, "current_snapshot_id")
+
+	agate := &Agate{
+		manager:       manager,
+		options:       options,
+		metadataDir:   metadataDir,
+		blobsDir:      blobsDir,
+		currentIDFile: currentIDFile,
+	}
+
+	// Load the current snapshot ID if it exists
+	if _, err := os.Stat(currentIDFile); err == nil {
+		data, err := os.ReadFile(currentIDFile)
+		if err == nil && len(data) > 0 {
+			agate.currentSnapshotID = string(data)
+		}
+	}
+
+	// Call OpenFunc if provided to initialize resources in the active directory
+	if options.OpenFunc != nil {
+		if err := options.OpenFunc(options.BlobStore.GetActiveDir()); err != nil {
+			return nil, fmt.Errorf("failed to open resources during initialization: %w", err)
+		}
+	}
+
+	return agate, nil
 }
 
 // SaveSnapshot creates a new snapshot from the current state of the active directory.
 // If parentID is provided, it will be set as the parent of the new snapshot.
+// If parentID is empty, it will use the ID of the snapshot currently loaded in the active directory.
 // Returns the ID of the created snapshot.
 func (a *Agate) SaveSnapshot(ctx context.Context, name string, parentID string) (string, error) {
 	// Call CloseFunc if provided
@@ -109,12 +133,23 @@ func (a *Agate) SaveSnapshot(ctx context.Context, name string, parentID string)
 		}
 	}
 
+	// If parentID is not provided, use the current snapshot ID
+	effectiveParentID := parentID
+
 	// Create the snapshot
-	snapshot, err := a.manager.CreateSnapshot(ctx, a.options.BlobStore.GetActiveDir(), name, parentID)
+	snapshot, err := a.manager.CreateSnapshot(ctx, a.options.BlobStore.GetActiveDir(), name, effectiveParentID)
 	if err != nil {
 		return "", fmt.Errorf("failed to create snapshot: %w", err)
 	}
 
+	// Update the current snapshot ID to the newly created snapshot
+	a.currentSnapshotID = snapshot.ID
+
+	// Save the current snapshot ID to a file
+	if err := a.saveCurrentSnapshotID(); err != nil {
+		return "", fmt.Errorf("failed to save current snapshot ID: %w", err)
+	}
+
 	// Call OpenFunc if provided
 	if a.options.OpenFunc != nil {
 		if err := a.options.OpenFunc(a.options.BlobStore.GetActiveDir()); err != nil {
@@ -139,6 +174,14 @@ func (a *Agate) RestoreSnapshot(ctx context.Context, snapshotID string) error {
 		return fmt.Errorf("failed to extract snapshot: %w", err)
 	}
 
+	// Save the ID of the snapshot that was restored
+	a.currentSnapshotID = snapshotID
+
+	// Save the current snapshot ID to a file
+	if err := a.saveCurrentSnapshotID(); err != nil {
+		return fmt.Errorf("failed to save current snapshot ID: %w", err)
+	}
+
 	// Call OpenFunc if provided
 	if a.options.OpenFunc != nil {
 		if err := a.options.OpenFunc(a.options.BlobStore.GetActiveDir()); err != nil {
@@ -163,6 +206,16 @@ func (a *Agate) RestoreSnapshotToDir(ctx context.Context, snapshotID string, dir
 		return fmt.Errorf("failed to extract snapshot: %w", err)
 	}
 
+	// If restoring to the active directory, save the snapshot ID
+	if dir == a.options.BlobStore.GetActiveDir() {
+		a.currentSnapshotID = snapshotID
+
+		// Save the current snapshot ID to a file
+		if err := a.saveCurrentSnapshotID(); err != nil {
+			return fmt.Errorf("failed to save current snapshot ID: %w", err)
+		}
+	}
+
 	// Call OpenFunc if provided
 	if a.options.OpenFunc != nil {
 		if err := a.options.OpenFunc(dir); err != nil {
@@ -188,6 +241,20 @@ func (a *Agate) DeleteSnapshot(ctx context.Context, snapshotID string) error {
 	return a.manager.DeleteSnapshot(ctx, snapshotID)
 }
 
+// saveCurrentSnapshotID saves the current snapshot ID to a file in the WorkDir
+func (a *Agate) saveCurrentSnapshotID() error {
+	if a.currentSnapshotID == "" {
+		// If there's no current snapshot ID, remove the file if it exists
+		if _, err := os.Stat(a.currentIDFile); err == nil {
+			return os.Remove(a.currentIDFile)
+		}
+		return nil
+	}
+
+	// Write the current snapshot ID to the file
+	return os.WriteFile(a.currentIDFile, []byte(a.currentSnapshotID), 0644)
+}
+
 // Close releases all resources used by the Agate instance.
 func (a *Agate) Close() error {
 	// Currently, we don't have a way to close the manager directly