jiggablend/internal/storage/storage.go

package storage

import (
	"archive/tar"
	"archive/zip"
	"fmt"
	"io"
	"log"
	"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(),
		s.tempPath(),
	}

	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
}

// tempPath returns the path for temporary files
func (s *Storage) tempPath() string {
	return filepath.Join(s.basePath, "temp")
}

// BasePath returns the storage base path (for cleanup tasks)
func (s *Storage) BasePath() string {
	return s.basePath
}

// TempDir creates a temporary directory under the storage base path
// Returns the path to the temporary directory
func (s *Storage) TempDir(pattern string) (string, error) {
	// Ensure temp directory exists
	if err := os.MkdirAll(s.tempPath(), 0755); err != nil {
		return "", fmt.Errorf("failed to create temp directory: %w", err)
	}

	// Create temp directory under storage base path
	return os.MkdirTemp(s.tempPath(), pattern)
}

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

// SanitizeFilename returns a safe basename for client-supplied names.
// Rejects empty, ".", "..", and root-like basenames.
func SanitizeFilename(name string) (string, error) {
	cleaned := filepath.Clean(strings.ReplaceAll(name, "\\", "/"))
	base := filepath.Base(cleaned)
	// filepath.Base("..") is ".." ; Clean("/../..") style inputs can yield "/" or "."
	if base == "" || base == "." || base == ".." || base == "/" || base == string(os.PathSeparator) {
		return "", fmt.Errorf("invalid filename: %q", name)
	}
	// Reject if the cleaned path still has parent references after basenaming was skipped
	if strings.Contains(cleaned, "..") && base == cleaned {
		return "", fmt.Errorf("invalid filename: %q", name)
	}
	return base, nil
}

// SafePathUnderRoot joins root and a relative path and ensures the result stays under root.
func SafePathUnderRoot(root, rel string) (string, error) {
	if rel == "" {
		return "", fmt.Errorf("empty path")
	}
	// Normalize to forward slashes then clean
	rel = filepath.ToSlash(rel)
	if filepath.IsAbs(rel) || strings.HasPrefix(rel, "/") {
		return "", fmt.Errorf("absolute path not allowed: %q", rel)
	}
	cleanRel := filepath.Clean(rel)
	if cleanRel == ".." || strings.HasPrefix(cleanRel, ".."+string(os.PathSeparator)) {
		return "", fmt.Errorf("path escapes root: %q", rel)
	}
	cleanRoot := filepath.Clean(root)
	full := filepath.Join(cleanRoot, cleanRel)
	cleanFull := filepath.Clean(full)
	sep := string(os.PathSeparator)
	if cleanFull != cleanRoot && !strings.HasPrefix(cleanFull, cleanRoot+sep) {
		return "", fmt.Errorf("path escapes root: %q", rel)
	}
	return cleanFull, nil
}

// SaveUpload saves an uploaded file
func (s *Storage) SaveUpload(jobID int64, filename string, reader io.Reader) (string, error) {
	// Sanitize filename to prevent path traversal
	var err error
	filename, err = SanitizeFilename(filename)
	if err != nil {
		return "", err
	}

	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) {
	// Sanitize filename to prevent path traversal (parity with SaveUpload)
	var err error
	filename, err = SanitizeFilename(filename)
	if err != nil {
		return "", err
	}

	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)
	// Defense in depth: ensure resolved path stays under output dir
	if resolved, err := SafePathUnderRoot(outputPath, filename); err != nil {
		return "", err
	} else {
		filePath = resolved
	}

	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) {
	log.Printf("Extracting ZIP archive: %s -> %s", zipPath, destDir)

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

	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
	fileCount := 0
	dirCount := 0

	log.Printf("ZIP contains %d entries", len(r.File))

	for _, f := range r.File {
		// Sanitize file path to prevent directory traversal
		destPath := filepath.Join(destDir, f.Name)

		cleanDestPath := filepath.Clean(destPath)
		cleanDestDir := filepath.Clean(destDir)
		if !strings.HasPrefix(cleanDestPath, cleanDestDir+string(os.PathSeparator)) && cleanDestPath != cleanDestDir {
			log.Printf("ERROR: Invalid file path in ZIP - target: %s, destDir: %s", cleanDestPath, cleanDestDir)
			return nil, fmt.Errorf("invalid file path in ZIP: %s (target: %s, destDir: %s)", f.Name, cleanDestPath, cleanDestDir)
		}

		// 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)
			}
			dirCount++
			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)
		fileCount++
	}

	log.Printf("ZIP extraction complete: %d files, %d directories extracted to %s", fileCount, dirCount, destDir)
	return extractedFiles, nil
}

// findCommonPrefix finds the common leading directory prefix if all paths share the same first-level directory
// Returns the prefix to strip (with trailing slash) or empty string if no common prefix
func findCommonPrefix(relPaths []string) string {
	if len(relPaths) == 0 {
		return ""
	}

	// Get the first path component of each path
	firstComponents := make([]string, 0, len(relPaths))
	for _, path := range relPaths {
		parts := strings.Split(filepath.ToSlash(path), "/")
		if len(parts) > 0 && parts[0] != "" {
			firstComponents = append(firstComponents, parts[0])
		} else {
			// If any path is at root level, no common prefix
			return ""
		}
	}

	// Check if all first components are the same
	if len(firstComponents) == 0 {
		return ""
	}

	commonFirst := firstComponents[0]
	for _, comp := range firstComponents {
		if comp != commonFirst {
			// Not all paths share the same first directory
			return ""
		}
	}

	// All paths share the same first directory - return it with trailing slash
	return commonFirst + "/"
}

// isBlenderSaveFile checks if a filename is a Blender save file (.blend1, .blend2, etc.)
// Returns true for files like "file.blend1", "file.blend2", but false for "file.blend"
func isBlenderSaveFile(filename string) bool {
	lower := strings.ToLower(filename)
	// Check if it ends with .blend followed by one or more digits
	// Pattern: *.blend[digits]
	if !strings.HasSuffix(lower, ".blend") {
		// Doesn't end with .blend, check if it ends with .blend + digits
		idx := strings.LastIndex(lower, ".blend")
		if idx == -1 {
			return false
		}
		// Check if there are digits after .blend
		suffix := lower[idx+len(".blend"):]
		if len(suffix) == 0 {
			return false
		}
		// All remaining characters must be digits
		for _, r := range suffix {
			if r < '0' || r > '9' {
				return false
			}
		}
		return true
	}
	// Ends with .blend exactly - this is a regular blend file, not a save file
	return false
}

// CreateJobContext creates a tar archive containing all job input files
// Filters out Blender save files (.blend1, .blend2, etc.)
// Uses temporary directories and streaming to handle large files efficiently
func (s *Storage) CreateJobContext(jobID int64) (string, error) {
	jobPath := s.JobPath(jobID)
	contextPath := filepath.Join(jobPath, "context.tar")

	// Create temporary directory for staging
	tmpDir, err := os.MkdirTemp("", "jiggablend-context-*")
	if err != nil {
		return "", fmt.Errorf("failed to create temporary directory: %w", err)
	}
	defer os.RemoveAll(tmpDir)

	// Collect all files from job directory, excluding the context file itself and Blender save files
	var filesToInclude []string
	err = filepath.Walk(jobPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Skip the context file itself if it exists
		if path == contextPath {
			return nil
		}

		// Skip Blender save files
		if isBlenderSaveFile(info.Name()) {
			return nil
		}

		// Get relative path from job directory
		relPath, err := filepath.Rel(jobPath, path)
		if err != nil {
			return err
		}

		// Sanitize path - ensure it doesn't escape the job directory
		cleanRelPath := filepath.Clean(relPath)
		if strings.HasPrefix(cleanRelPath, "..") {
			return fmt.Errorf("invalid file path: %s", relPath)
		}

		filesToInclude = append(filesToInclude, path)
		return nil
	})
	if err != nil {
		return "", fmt.Errorf("failed to walk job directory: %w", err)
	}

	if len(filesToInclude) == 0 {
		return "", fmt.Errorf("no files found to include in context")
	}

	// Create the tar file using streaming
	contextFile, err := os.Create(contextPath)
	if err != nil {
		return "", fmt.Errorf("failed to create context file: %w", err)
	}
	defer contextFile.Close()

	tarWriter := tar.NewWriter(contextFile)
	defer tarWriter.Close()

	// Add each file to the tar archive
	for _, filePath := range filesToInclude {
		file, err := os.Open(filePath)
		if err != nil {
			return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
		}

		// Use a function closure to ensure file is closed even on error
		err = func() error {
			defer file.Close()

			info, err := file.Stat()
			if err != nil {
				return fmt.Errorf("failed to stat file %s: %w", filePath, err)
			}

			// Get relative path for tar header
			relPath, err := filepath.Rel(jobPath, filePath)
			if err != nil {
				return fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
			}

			// Normalize path separators for tar (use forward slashes)
			tarPath := filepath.ToSlash(relPath)

			// Create tar header
			header, err := tar.FileInfoHeader(info, "")
			if err != nil {
				return fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
			}
			header.Name = tarPath

			// Write header
			if err := tarWriter.WriteHeader(header); err != nil {
				return fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
			}

			// Copy file contents using streaming
			if _, err := io.Copy(tarWriter, file); err != nil {
				return fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
			}

			return nil
		}()

		if err != nil {
			return "", err
		}
	}

	// Ensure all data is flushed
	if err := tarWriter.Close(); err != nil {
		return "", fmt.Errorf("failed to close tar writer: %w", err)
	}
	if err := contextFile.Close(); err != nil {
		return "", fmt.Errorf("failed to close context file: %w", err)
	}

	return contextPath, nil
}

// CreateJobContextFromDir creates a context archive (tar) from files in a source directory
// This is used during upload to immediately create the context archive as the primary artifact
// excludeFiles is a set of relative paths (from sourceDir) to exclude from the context
func (s *Storage) CreateJobContextFromDir(sourceDir string, jobID int64, excludeFiles ...string) (string, error) {
	jobPath := s.JobPath(jobID)
	contextPath := filepath.Join(jobPath, "context.tar")

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

	// Build set of files to exclude (normalize paths)
	excludeSet := make(map[string]bool)
	for _, excludeFile := range excludeFiles {
		// Normalize the exclude path
		excludePath := filepath.Clean(excludeFile)
		excludeSet[excludePath] = true
		// Also add with forward slash for cross-platform compatibility
		excludeSet[filepath.ToSlash(excludePath)] = true
	}

	// Collect all files from source directory, excluding Blender save files and excluded files
	var filesToInclude []string
	err := filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Skip Blender save files
		if isBlenderSaveFile(info.Name()) {
			return nil
		}

		// Get relative path from source directory
		relPath, err := filepath.Rel(sourceDir, path)
		if err != nil {
			return err
		}

		// Sanitize path - ensure it doesn't escape the source directory
		cleanRelPath := filepath.Clean(relPath)
		if strings.HasPrefix(cleanRelPath, "..") {
			return fmt.Errorf("invalid file path: %s", relPath)
		}

		// Check if this file should be excluded
		if excludeSet[cleanRelPath] || excludeSet[filepath.ToSlash(cleanRelPath)] {
			return nil
		}

		filesToInclude = append(filesToInclude, path)
		return nil
	})
	if err != nil {
		return "", fmt.Errorf("failed to walk source directory: %w", err)
	}

	if len(filesToInclude) == 0 {
		return "", fmt.Errorf("no files found to include in context archive")
	}

	// Collect relative paths to find common prefix
	relPaths := make([]string, 0, len(filesToInclude))
	for _, filePath := range filesToInclude {
		relPath, err := filepath.Rel(sourceDir, filePath)
		if err != nil {
			return "", fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
		}
		relPaths = append(relPaths, relPath)
	}

	// Find and strip common leading directory if all files share one
	commonPrefix := findCommonPrefix(relPaths)

	// Validate that there's exactly one .blend file at the root level after prefix stripping
	blendFilesAtRoot := 0
	for _, relPath := range relPaths {
		tarPath := filepath.ToSlash(relPath)
		// Strip common prefix if present
		if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
			tarPath = strings.TrimPrefix(tarPath, commonPrefix)
		}

		// Check if it's a .blend file at root (no path separators after prefix stripping)
		if strings.HasSuffix(strings.ToLower(tarPath), ".blend") {
			// Check if it's at root level (no directory separators)
			if !strings.Contains(tarPath, "/") {
				blendFilesAtRoot++
			}
		}
	}

	if blendFilesAtRoot == 0 {
		return "", fmt.Errorf("no .blend file found at root level in context archive - .blend files must be at the root level of the uploaded archive, not in subdirectories")
	}
	if blendFilesAtRoot > 1 {
		return "", fmt.Errorf("multiple .blend files found at root level in context archive (found %d, expected 1)", blendFilesAtRoot)
	}

	// Create the tar file using streaming
	contextFile, err := os.Create(contextPath)
	if err != nil {
		return "", fmt.Errorf("failed to create context file: %w", err)
	}
	defer contextFile.Close()

	tarWriter := tar.NewWriter(contextFile)
	defer tarWriter.Close()

	// Add each file to the tar archive
	for i, filePath := range filesToInclude {
		file, err := os.Open(filePath)
		if err != nil {
			return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
		}

		// Use a function closure to ensure file is closed even on error
		err = func() error {
			defer file.Close()

			info, err := file.Stat()
			if err != nil {
				return fmt.Errorf("failed to stat file %s: %w", filePath, err)
			}

			// Get relative path and strip common prefix if present
			relPath := relPaths[i]
			tarPath := filepath.ToSlash(relPath)

			// Strip common prefix if found
			if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
				tarPath = strings.TrimPrefix(tarPath, commonPrefix)
			}

			// Create tar header
			header, err := tar.FileInfoHeader(info, "")
			if err != nil {
				return fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
			}
			header.Name = tarPath

			// Write header
			if err := tarWriter.WriteHeader(header); err != nil {
				return fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
			}

			// Copy file contents using streaming
			if _, err := io.Copy(tarWriter, file); err != nil {
				return fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
			}

			return nil
		}()

		if err != nil {
			return "", err
		}
	}

	// Ensure all data is flushed
	if err := tarWriter.Close(); err != nil {
		return "", fmt.Errorf("failed to close tar writer: %w", err)
	}
	if err := contextFile.Close(); err != nil {
		return "", fmt.Errorf("failed to close context file: %w", err)
	}

	return contextPath, nil
}

// CreateContextArchiveFromDirToPath creates a context archive from files in sourceDir at destPath.
// This is used for pre-job upload sessions where the archive is staged before a job ID exists.
func (s *Storage) CreateContextArchiveFromDirToPath(sourceDir, destPath string, excludeFiles ...string) (string, error) {
	excludeSet := make(map[string]bool)
	for _, excludeFile := range excludeFiles {
		excludePath := filepath.Clean(excludeFile)
		excludeSet[excludePath] = true
		excludeSet[filepath.ToSlash(excludePath)] = true
	}

	var filesToInclude []string
	err := filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		if isBlenderSaveFile(info.Name()) {
			return nil
		}

		relPath, err := filepath.Rel(sourceDir, path)
		if err != nil {
			return err
		}
		cleanRelPath := filepath.Clean(relPath)
		if strings.HasPrefix(cleanRelPath, "..") {
			return fmt.Errorf("invalid file path: %s", relPath)
		}
		if excludeSet[cleanRelPath] || excludeSet[filepath.ToSlash(cleanRelPath)] {
			return nil
		}

		filesToInclude = append(filesToInclude, path)
		return nil
	})
	if err != nil {
		return "", fmt.Errorf("failed to walk source directory: %w", err)
	}
	if len(filesToInclude) == 0 {
		return "", fmt.Errorf("no files found to include in context archive")
	}

	relPaths := make([]string, 0, len(filesToInclude))
	for _, filePath := range filesToInclude {
		relPath, err := filepath.Rel(sourceDir, filePath)
		if err != nil {
			return "", fmt.Errorf("failed to get relative path: %w", err)
		}
		relPaths = append(relPaths, relPath)
	}

	commonPrefix := findCommonPrefix(relPaths)
	blendFilesAtRoot := 0
	for _, relPath := range relPaths {
		tarPath := filepath.ToSlash(relPath)
		if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
			tarPath = strings.TrimPrefix(tarPath, commonPrefix)
		}
		if strings.HasSuffix(strings.ToLower(tarPath), ".blend") && !strings.Contains(tarPath, "/") {
			blendFilesAtRoot++
		}
	}
	if blendFilesAtRoot == 0 {
		return "", fmt.Errorf("no .blend file found at root level in context archive - .blend files must be at the root level of the uploaded archive, not in subdirectories")
	}
	if blendFilesAtRoot > 1 {
		return "", fmt.Errorf("multiple .blend files found at root level in context archive (found %d, expected 1)", blendFilesAtRoot)
	}

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

	tarWriter := tar.NewWriter(contextFile)
	defer tarWriter.Close()
	copyBuf := make([]byte, 32*1024)

	for i, filePath := range filesToInclude {
		file, err := os.Open(filePath)
		if err != nil {
			return "", fmt.Errorf("failed to open file: %w", err)
		}

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

		tarPath := filepath.ToSlash(relPaths[i])
		if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
			tarPath = strings.TrimPrefix(tarPath, commonPrefix)
		}

		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			file.Close()
			return "", fmt.Errorf("failed to create tar header: %w", err)
		}
		header.Name = tarPath

		if err := tarWriter.WriteHeader(header); err != nil {
			file.Close()
			return "", fmt.Errorf("failed to write tar header: %w", err)
		}

		if _, err := io.CopyBuffer(tarWriter, file, copyBuf); err != nil {
			file.Close()
			return "", fmt.Errorf("failed to write file to tar: %w", err)
		}
		file.Close()
	}

	if err := tarWriter.Close(); err != nil {
		return "", fmt.Errorf("failed to close tar writer: %w", err)
	}
	if err := contextFile.Close(); err != nil {
		return "", fmt.Errorf("failed to close context file: %w", err)
	}

	return destPath, nil
}