Agate/manager.go

package agate

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

	"github.com/google/uuid"

	"gitea.unprism.ru/KRBL/Agate/archive"
	"gitea.unprism.ru/KRBL/Agate/hash"
	"gitea.unprism.ru/KRBL/Agate/interfaces"
	"gitea.unprism.ru/KRBL/Agate/store"
)

type SnapshotManagerData struct {
	metadataStore store.MetadataStore
	blobStore     store.BlobStore
	logger        *log.Logger
}

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

	// Ensure logger is never nil.
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}

	return &SnapshotManagerData{
		metadataStore: metadataStore,
		blobStore:     blobStore,
		logger:        logger,
	}, 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 {
			fmt.Println("failed to check parent snapshot: %w", err)
			parentID = ""
		}
	}

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

	// Create a temporary file for the archive in the working directory
	tempFilePath := filepath.Join(data.blobStore.GetBaseDir(), "temp-"+snapshotID+".zip")
	tempFile, err := os.Create(tempFilePath)
	if err != nil {
		return nil, fmt.Errorf("failed to create temporary file in working directory: %w", err)
	}
	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, opts store.ListOptions) ([]store.SnapshotInfo, error) {
	// Retrieve list of snapshots from the metadata store with the provided options
	snapshots, err := data.metadataStore.ListSnapshotsMetadata(ctx, opts)
	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")
	}

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

	parentID := snapshot.ParentID

	opts := store.ListOptions{}
	allSnapshots, err := data.metadataStore.ListSnapshotsMetadata(ctx, opts)
	if err != nil {
		return fmt.Errorf("failed to list snapshots: %w", err)
	}

	for _, info := range allSnapshots {
		if info.ParentID == snapshotID {
			if err := data.metadataStore.UpdateSnapshotParentID(ctx, info.ID, parentID); err != nil {
				data.logger.Printf("WARNING: failed to update parent reference for snapshot %s: %v", info.ID, err)
			} else {
				data.logger.Printf("Updated parent reference for snapshot %s from %s to %s", info.ID, snapshotID, parentID)
			}
		}
	}

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

	if err := data.blobStore.DeleteBlob(ctx, snapshotID); err != nil {
		data.logger.Printf("WARNING: failed to delete snapshot blob: %v", 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, cleanTarget bool) error {
	if snapshotID == "" {
		return errors.New("snapshot ID cannot be empty")
	}

	// If no specific path is provided, use the active directory
	if path == "" {
		path = data.blobStore.GetActiveDir()
	}

	// 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)
	}

	// If cleanTarget is true, clean the target directory before extraction
	if cleanTarget {
		// Remove the directory and recreate it
		if err := os.RemoveAll(path); err != nil {
			return fmt.Errorf("failed to clean target directory: %w", err)
		}
		if err := os.MkdirAll(path, 0755); err != nil {
			return fmt.Errorf("failed to create target directory: %w", err)
		}
	} else {
		// Just 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")
	}

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

	snapshot.Name = newName

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

	return nil
}

// diffArchiveReader is a wrapper around an *os.File that handles cleanup of temporary files.
type diffArchiveReader struct {
	*os.File
	tempArchive string
	tempStaging string
}

// Close closes the file and removes the temporary archive and staging directory.
func (r *diffArchiveReader) Close() error {
	err := r.File.Close()
	os.Remove(r.tempArchive)
	os.RemoveAll(r.tempStaging)
	return err
}

func (data *SnapshotManagerData) StreamSnapshotDiff(ctx context.Context, snapshotID, parentID string, offset int64) (io.ReadCloser, error) {
	targetSnap, err := data.metadataStore.GetSnapshotMetadata(ctx, snapshotID)
	if err != nil {
		return nil, fmt.Errorf("failed to get target snapshot metadata: %w", err)
	}

	parentFiles := make(map[string]string)
	if parentID != "" {
		parentSnap, err := data.metadataStore.GetSnapshotMetadata(ctx, parentID)
		if err == nil {
			for _, file := range parentSnap.Files {
				if !file.IsDir {
					parentFiles[file.Path] = file.SHA256
				}
			}
		} else {
			data.logger.Printf("Warning: failed to get parent snapshot %s, creating full diff: %v", parentID, err)
		}
	}

	var filesToInclude []string
	for _, file := range targetSnap.Files {
		if file.IsDir {
			continue
		}
		if parentHash, ok := parentFiles[file.Path]; !ok || parentHash != file.SHA256 {
			filesToInclude = append(filesToInclude, file.Path)
		}
	}

	if len(filesToInclude) == 0 {
		return io.NopCloser(bytes.NewReader(nil)), nil
	}

	tempStagingDir, err := os.MkdirTemp(data.blobStore.GetBaseDir(), "diff-staging-*")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp staging directory: %w", err)
	}

	targetBlobPath, err := data.blobStore.GetBlobPath(ctx, snapshotID)
	if err != nil {
		os.RemoveAll(tempStagingDir)
		return nil, err
	}

	for _, filePath := range filesToInclude {
		destPath := filepath.Join(tempStagingDir, filePath)
		if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
			os.RemoveAll(tempStagingDir)
			return nil, fmt.Errorf("failed to create dir for diff file: %w", err)
		}

		fileWriter, err := os.Create(destPath)
		if err != nil {
			os.RemoveAll(tempStagingDir)
			return nil, err
		}

		err = archive.ExtractFileFromArchive(targetBlobPath, filePath, fileWriter)
		fileWriter.Close()
		if err != nil {
			os.RemoveAll(tempStagingDir)
			return nil, fmt.Errorf("failed to extract file %s for diff: %w", filePath, err)
		}
	}

	tempArchivePath := filepath.Join(data.blobStore.GetBaseDir(), "diff-"+snapshotID+".zip")
	if err := archive.CreateArchive(tempStagingDir, tempArchivePath); err != nil {
		os.RemoveAll(tempStagingDir)
		os.Remove(tempArchivePath)
		return nil, fmt.Errorf("failed to create diff archive: %w", err)
	}

	archiveFile, err := os.Open(tempArchivePath)
	if err != nil {
		os.RemoveAll(tempStagingDir)
		os.Remove(tempArchivePath)
		return nil, err
	}

	if offset > 0 {
		if _, err := archiveFile.Seek(offset, io.SeekStart); err != nil {
			archiveFile.Close()
			os.RemoveAll(tempStagingDir)
			os.Remove(tempArchivePath)
			return nil, fmt.Errorf("failed to seek in diff archive: %w", err)
		}
	}

	return &diffArchiveReader{
		File:        archiveFile,
		tempArchive: tempArchivePath,
		tempStaging: tempStagingDir,
	}, nil
}