Agate/manager.go

package agate

import (
	"archive/zip"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/uuid"

	"unprism.ru/KRBL/agate/archive"
	"unprism.ru/KRBL/agate/hash"
	"unprism.ru/KRBL/agate/store"
)

type SnapshotManagerData struct {
	metadataStore store.MetadataStore
	blobStore     store.BlobStore
}

func CreateSnapshotManager(metadataStore store.MetadataStore, blobStore store.BlobStore) (SnapshotManager, error) {
	if metadataStore == nil || blobStore == nil {
		return nil, errors.New("parameters can't be nil")
	}

	return &SnapshotManagerData{metadataStore, blobStore}, nil
}

func (data *SnapshotManagerData) CreateSnapshot(ctx context.Context, sourceDir string, name string, parentID string) (*store.Snapshot, error) {
	// Validate source directory
	info, err := os.Stat(sourceDir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, ErrSourceNotFound
		}
		return nil, fmt.Errorf("failed to access source directory: %w", err)
	}

	if !info.IsDir() {
		return nil, ErrSourceNotDirectory
	}

	// Check if parent exists if specified
	if parentID != "" {
		_, err := data.metadataStore.GetSnapshotMetadata(ctx, parentID)
		if err != nil {
			if errors.Is(err, ErrNotFound) {
				return nil, ErrParentNotFound
			}
			return nil, fmt.Errorf("failed to check parent snapshot: %w", err)
		}
	}

	// Generate a unique ID for the snapshot
	snapshotID := uuid.New().String()

	// Create a temporary file for the archive
	tempFile, err := os.CreateTemp("", "agate-snapshot-*.zip")
	if err != nil {
		return nil, fmt.Errorf("failed to create temporary file: %w", err)
	}
	tempFilePath := tempFile.Name()
	tempFile.Close()              // Close it as CreateArchive will reopen it
	defer os.Remove(tempFilePath) // Clean up temp file after we're done

	// Create archive of the source directory
	if err := archive.CreateArchive(sourceDir, tempFilePath); err != nil {
		return nil, fmt.Errorf("failed to create archive: %w", err)
	}

	// Scan the directory to collect file information
	var files []store.FileInfo
	err = filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip the root directory itself
		if path == sourceDir {
			return nil
		}

		// Create relative path
		relPath, err := filepath.Rel(sourceDir, path)
		if err != nil {
			return fmt.Errorf("failed to get relative path: %w", err)
		}

		fileInfo := store.FileInfo{
			Path:  filepath.ToSlash(relPath),
			Size:  info.Size(),
			IsDir: info.IsDir(),
		}

		// Calculate hash for files (not directories)
		if !info.IsDir() {
			hash, err := hash.CalculateFileHash(path)
			if err != nil {
				return fmt.Errorf("failed to calculate hash for %s: %w", path, err)
			}
			fileInfo.SHA256 = hash
		}

		files = append(files, fileInfo)
		return nil
	})

	if err != nil {
		return nil, fmt.Errorf("failed to scan directory: %w", err)
	}

	// Open the archive file for reading
	archiveFile, err := os.Open(tempFilePath)
	if err != nil {
		return nil, fmt.Errorf("failed to open archive file: %w", err)
	}
	defer archiveFile.Close()

	// Store the blob
	_, err = data.blobStore.StoreBlob(ctx, snapshotID, archiveFile)
	if err != nil {
		return nil, fmt.Errorf("failed to store blob: %w", err)
	}

	// Create snapshot metadata
	snapshot := store.Snapshot{
		ID:           snapshotID,
		Name:         name,
		ParentID:     parentID,
		CreationTime: time.Now(),
		Files:        files,
	}

	// Save metadata
	if err := data.metadataStore.SaveSnapshotMetadata(ctx, snapshot); err != nil {
		// If metadata save fails, try to clean up the blob
		_ = data.blobStore.DeleteBlob(ctx, snapshotID)
		return nil, fmt.Errorf("failed to save snapshot metadata: %w", err)
	}

	return &snapshot, nil
}

func (data *SnapshotManagerData) GetSnapshotDetails(ctx context.Context, snapshotID string) (*store.Snapshot, error) {
	if snapshotID == "" {
		return nil, errors.New("snapshot ID cannot be empty")
	}

	// Retrieve snapshot metadata from the store
	snapshot, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to retrieve snapshot details: %w", err)
	}

	return snapshot, nil
}

func (data *SnapshotManagerData) ListSnapshots(ctx context.Context) ([]store.SnapshotInfo, error) {
	// Retrieve list of snapshots from the metadata store
	snapshots, err := data.metadataStore.ListSnapshotsMetadata(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list snapshots: %w", err)
	}

	return snapshots, nil
}

func (data *SnapshotManagerData) DeleteSnapshot(ctx context.Context, snapshotID string) error {
	if snapshotID == "" {
		return errors.New("snapshot ID cannot be empty")
	}

	// First check if the snapshot exists
	_, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			// If snapshot doesn't exist, return success (idempotent operation)
			return nil
		}
		return fmt.Errorf("failed to check if snapshot exists: %w", err)
	}

	// Delete the metadata first
	if err := data.metadataStore.DeleteSnapshotMetadata(ctx, snapshotID); err != nil {
		return fmt.Errorf("failed to delete snapshot metadata: %w", err)
	}

	// Then delete the blob
	if err := data.blobStore.DeleteBlob(ctx, snapshotID); err != nil {
		// Note: We don't return here because we've already deleted the metadata
		// and the blob store should handle the case where the blob doesn't exist
		// Log the error instead
		fmt.Printf("Warning: failed to delete snapshot blob: %v\n", err)
	}

	return nil
}

func (data *SnapshotManagerData) OpenFile(ctx context.Context, snapshotID string, filePath string) (io.ReadCloser, error) {
	if snapshotID == "" {
		return nil, errors.New("snapshot ID cannot be empty")
	}
	if filePath == "" {
		return nil, errors.New("file path cannot be empty")
	}

	// First check if the snapshot exists
	_, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to check if snapshot exists: %w", err)
	}

	// Get the blob path from the blob store
	blobPath, err := data.blobStore.GetBlobPath(ctx, snapshotID)
	if err != nil {
		return nil, fmt.Errorf("failed to get blob path: %w", err)
	}

	// Create a pipe to stream the file content
	pr, pw := io.Pipe()

	// Extract the file in a goroutine to avoid blocking
	go func() {
		// Close the write end of the pipe when done
		defer pw.Close()

		// Extract the file from the archive
		err := archive.ExtractFileFromArchive(blobPath, filePath, pw)
		if err != nil {
			if errors.Is(err, archive.ErrFileNotFoundInArchive) {
				pw.CloseWithError(ErrFileNotFound)
				return
			}
			pw.CloseWithError(fmt.Errorf("failed to extract file from archive: %w", err))
		}
	}()

	return pr, nil
}

func (data *SnapshotManagerData) ExtractSnapshot(ctx context.Context, snapshotID string, path string) error {
	if snapshotID == "" {
		return errors.New("snapshot ID cannot be empty")
	}
	if path == "" {
		return errors.New("target path cannot be empty")
	}

	// First check if the snapshot exists
	_, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			return ErrNotFound
		}
		return fmt.Errorf("failed to check if snapshot exists: %w", err)
	}

	// Get the blob path from the blob store
	blobPath, err := data.blobStore.GetBlobPath(ctx, snapshotID)
	if err != nil {
		return fmt.Errorf("failed to get blob path: %w", err)
	}

	// Ensure the target directory exists
	if err := os.MkdirAll(path, 0755); err != nil {
		return fmt.Errorf("failed to create target directory: %w", err)
	}

	// Extract the archive to the target directory
	if err := extractArchive(blobPath, path); err != nil {
		return fmt.Errorf("failed to extract snapshot: %w", err)
	}

	return nil
}

// extractArchive extracts a ZIP archive to a target directory
func extractArchive(archivePath, targetDir string) error {
	// Open the ZIP archive
	zipReader, err := zip.OpenReader(archivePath)
	if err != nil {
		return fmt.Errorf("failed to open archive file: %w", err)
	}
	defer zipReader.Close()

	// Extract each file
	for _, file := range zipReader.File {
		// Construct the full path for the file
		filePath := filepath.Join(targetDir, file.Name)

		// Check for path traversal attacks
		if !strings.HasPrefix(filePath, targetDir) {
			return fmt.Errorf("invalid file path in archive: %s", file.Name)
		}

		if file.FileInfo().IsDir() {
			// Create directory
			if err := os.MkdirAll(filePath, 0755); err != nil {
				return fmt.Errorf("failed to create directory %s: %w", filePath, err)
			}
			continue
		}

		// Ensure the parent directory exists
		if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
			return fmt.Errorf("failed to create parent directory for %s: %w", filePath, err)
		}

		// Create the file
		outFile, err := os.Create(filePath)
		if err != nil {
			return fmt.Errorf("failed to create file %s: %w", filePath, err)
		}

		// Open the file in the archive
		inFile, err := file.Open()
		if err != nil {
			outFile.Close()
			return fmt.Errorf("failed to open file %s in archive: %w", file.Name, err)
		}

		// Copy the content
		_, err = io.Copy(outFile, inFile)
		outFile.Close()
		inFile.Close()

		if err != nil {
			return fmt.Errorf("failed to copy content for %s: %w", filePath, err)
		}
	}

	return nil
}

func (data *SnapshotManagerData) UpdateSnapshotMetadata(ctx context.Context, snapshotID string, newName string) error {
	if snapshotID == "" {
		return errors.New("snapshot ID cannot be empty")
	}
	if newName == "" {
		return errors.New("new name cannot be empty")
	}

	// Get the current snapshot metadata
	snapshot, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			return ErrNotFound
		}
		return fmt.Errorf("failed to get snapshot metadata: %w", err)
	}

	// Update the name
	snapshot.Name = newName

	// Save the updated metadata
	if err := data.metadataStore.SaveSnapshotMetadata(ctx, *snapshot); err != nil {
		return fmt.Errorf("failed to update snapshot metadata: %w", err)
	}

	return nil
}