jiggablend/internal/storage/storage.go

package storage

import (
	"archive/zip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// Storage handles file storage operations
type Storage struct {
	basePath string
}

// NewStorage creates a new storage instance
func NewStorage(basePath string) (*Storage, error) {
	s := &Storage{basePath: basePath}
	if err := s.init(); err != nil {
		return nil, err
	}
	return s, nil
}

// init creates necessary directories
func (s *Storage) init() error {
	dirs := []string{
		s.basePath,
		s.uploadsPath(),
		s.outputsPath(),
	}

	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("failed to create directory %s: %w", dir, err)
		}
	}

	return nil
}

// uploadsPath returns the path for uploads
func (s *Storage) uploadsPath() string {
	return filepath.Join(s.basePath, "uploads")
}

// outputsPath returns the path for outputs
func (s *Storage) outputsPath() string {
	return filepath.Join(s.basePath, "outputs")
}

// JobPath returns the path for a specific job's files
func (s *Storage) JobPath(jobID int64) string {
	return filepath.Join(s.basePath, "jobs", fmt.Sprintf("%d", jobID))
}

// SaveUpload saves an uploaded file
func (s *Storage) SaveUpload(jobID int64, filename string, reader io.Reader) (string, error) {
	jobPath := s.JobPath(jobID)
	if err := os.MkdirAll(jobPath, 0755); err != nil {
		return "", fmt.Errorf("failed to create job directory: %w", err)
	}

	filePath := filepath.Join(jobPath, filename)
	file, err := os.Create(filePath)
	if err != nil {
		return "", fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	if _, err := io.Copy(file, reader); err != nil {
		return "", fmt.Errorf("failed to write file: %w", err)
	}

	return filePath, nil
}

// SaveOutput saves an output file
func (s *Storage) SaveOutput(jobID int64, filename string, reader io.Reader) (string, error) {
	outputPath := filepath.Join(s.outputsPath(), fmt.Sprintf("%d", jobID))
	if err := os.MkdirAll(outputPath, 0755); err != nil {
		return "", fmt.Errorf("failed to create output directory: %w", err)
	}

	filePath := filepath.Join(outputPath, filename)
	file, err := os.Create(filePath)
	if err != nil {
		return "", fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	if _, err := io.Copy(file, reader); err != nil {
		return "", fmt.Errorf("failed to write file: %w", err)
	}

	return filePath, nil
}

// GetFile returns a file reader for the given path
func (s *Storage) GetFile(filePath string) (*os.File, error) {
	return os.Open(filePath)
}

// DeleteFile deletes a file
func (s *Storage) DeleteFile(filePath string) error {
	return os.Remove(filePath)
}

// DeleteJobFiles deletes all files for a job
func (s *Storage) DeleteJobFiles(jobID int64) error {
	jobPath := s.JobPath(jobID)
	if err := os.RemoveAll(jobPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to delete job files: %w", err)
	}

	outputPath := filepath.Join(s.outputsPath(), fmt.Sprintf("%d", jobID))
	if err := os.RemoveAll(outputPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to delete output files: %w", err)
	}

	return nil
}

// FileExists checks if a file exists
func (s *Storage) FileExists(filePath string) bool {
	_, err := os.Stat(filePath)
	return err == nil
}

// GetFileSize returns the size of a file
func (s *Storage) GetFileSize(filePath string) (int64, error) {
	info, err := os.Stat(filePath)
	if err != nil {
		return 0, err
	}
	return info.Size(), nil
}

// ExtractZip extracts a ZIP file to the destination directory
// Returns a list of all extracted file paths
func (s *Storage) ExtractZip(zipPath, destDir string) ([]string, error) {
	r, err := zip.OpenReader(zipPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open ZIP file: %w", err)
	}
	defer r.Close()

	var extractedFiles []string

	for _, f := range r.File {
		// Sanitize file path to prevent directory traversal
		destPath := filepath.Join(destDir, f.Name)
		if !strings.HasPrefix(destPath, filepath.Clean(destDir)+string(os.PathSeparator)) {
			return nil, fmt.Errorf("invalid file path in ZIP: %s", f.Name)
		}

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

		// Create parent directories
		if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
			return nil, fmt.Errorf("failed to create parent directory: %w", err)
		}

		// Extract file
		rc, err := f.Open()
		if err != nil {
			return nil, fmt.Errorf("failed to open file in ZIP: %w", err)
		}

		outFile, err := os.Create(destPath)
		if err != nil {
			rc.Close()
			return nil, fmt.Errorf("failed to create file: %w", err)
		}

		_, err = io.Copy(outFile, rc)
		outFile.Close()
		rc.Close()

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

		extractedFiles = append(extractedFiles, destPath)
	}

	return extractedFiles, nil
}