query-orchestration/internal/server/api/batch_worker.go

package api

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log/slog"
	"path"
	"strings"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document/batch"
	"queryorchestration/internal/document/batch/foldercleanup"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
)

type batchWorkerService interface {
	AddFailedFilename(ctx context.Context, batchID uuid.UUID, filename string) error
	CheckDuplicate(ctx context.Context, clientID string, hash string) (*uuid.UUID, error)
	MarkCompleted(ctx context.Context, batchID uuid.UUID) error
	MarkFailed(ctx context.Context, batchID uuid.UUID) error
	RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string, outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error
	SetTotalDocuments(ctx context.Context, batchID uuid.UUID, total int32) error
	UpdateProgress(ctx context.Context, clientID string, batchID uuid.UUID, processed int32, failed int32, invalid int32) error
}

// Config key constants for batch worker
const (
	ConfigKeyBatchService  = "batchService"
	ConfigKeyS3Client      = "s3Client"
	ConfigKeyBucket        = "bucket"
	ConfigKeyUploadHandler = "uploadHandler"
	ConfigKeyFolderCleanup = "folderCleanupService"
)

// acceptedExtensions defines the file extensions allowed in batch uploads.
// Macro-enabled extensions (.docm, .xlsm, .pptm) are intentionally excluded --
// documents with macros are rejected, so users must save as standard formats first.
var acceptedExtensions = map[string]bool{
	".pdf":  true,
	".tiff": true,
	".tif":  true,
	".jpeg": true,
	".jpg":  true,
	".png":  true,
	".bmp":  true,
	".docx": true,
	".xlsx": true,
	".pptx": true,
	".txt":  true,
	".eml":  true,
}

// isAcceptedExtension checks if a filename has a supported file extension.
func isAcceptedExtension(filename string) bool {
	ext := strings.ToLower(path.Ext(filename))
	return acceptedExtensions[ext]
}

// sanitizeZipPath cleans a path from a zip file to prevent path traversal attacks
// while preserving the internal folder structure.
// Returns empty string if the path is invalid or attempts traversal.
func sanitizeZipPath(p string) string {
	// Reject paths with backslashes (Windows path traversal)
	if strings.Contains(p, "\\") {
		return ""
	}

	// Clean the path to resolve . and .. elements
	cleaned := path.Clean(p)

	// Reject paths that try to escape (start with .. or /)
	if strings.HasPrefix(cleaned, "../") || strings.HasPrefix(cleaned, "/") || cleaned == ".." {
		return ""
	}

	// Check for net upward traversal
	// We need to ensure that we don't escape the root level
	parts := strings.Split(p, "/")
	depth := 0
	for _, part := range parts {
		if part == ".." {
			depth--
			if depth < 0 {
				// Trying to escape root
				return ""
			}
		} else if part != "." && part != "" {
			depth++
		}
	}

	// The cleaned path is safe and preserves internal structure
	return cleaned
}

// processBatchWork is the background worker function that processes batch uploads
func processBatchWork(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	// Extract the batch service from config
	batchService, ok := config[ConfigKeyBatchService].(*batch.Service)
	if !ok {
		return fmt.Errorf("batch service not found in config")
	}

	// Extract the S3 client from config
	s3Client, ok := config[ConfigKeyS3Client].(*s3.Client)
	if !ok {
		return fmt.Errorf("S3 client not found in config")
	}

	// Extract the bucket name from config
	bucket, ok := config[ConfigKeyBucket].(string)
	if !ok {
		return fmt.Errorf("bucket not found in config")
	}

	// Extract the document upload handler from config
	uploadHandler, ok := config[ConfigKeyUploadHandler].(func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error)
	if !ok {
		return fmt.Errorf("upload handler not found in config")
	}

	cleanupService, _ := config[ConfigKeyFolderCleanup].(*foldercleanup.Service)
	if cleanupService != nil {
		if finalized, err := cleanupService.FinalizeReadyBatches(ctx, foldercleanup.DefaultReadyBatchLimit); err != nil {
			logger.Error("Failed to finalize ready batch folders", "error", err)
		} else if finalized > 0 {
			logger.Info("Finalized ready batch folders", "count", finalized)
		}
	}

	// Query unprocessed batches
	batches, err := batchService.ListUnprocessed(ctx)
	if err != nil {
		logger.Error("Failed to query unprocessed batches", "error", err)
		return err
	}

	logger.Info("Background task executed", "unprocessed_batches", len(batches))

	// Process each batch
	for _, batchSummary := range batches {
		// Get full batch details with archive information
		batchDetails, err := batchService.Get(ctx, batchSummary.ClientID, batchSummary.ID)
		if err != nil {
			logger.Error("Failed to get batch details",
				"batch_id", batchSummary.ID,
				"client_id", batchSummary.ClientID,
				"error", err)
			continue
		}

		if err := processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails, cleanupService); err != nil {
			logger.Error("Failed to process batch",
				"batch_id", batchDetails.ID,
				"client_id", batchDetails.ClientID,
				"error", err)
			// Continue processing other batches even if one fails
			continue
		}
	}

	if cleanupService != nil {
		if deleted, err := cleanupService.DeleteSafeAutoCreatedFolders(ctx, foldercleanup.DefaultFolderDeleteLimit); err != nil {
			logger.Error("Failed to delete safe auto-created folders", "error", err)
		} else if deleted > 0 {
			logger.Info("Deleted safe auto-created folders", "count", deleted)
		}
	}

	return nil
}

// processSingleBatch processes a single batch upload
func processSingleBatch(
	ctx context.Context,
	logger *slog.Logger,
	batchService batchWorkerService,
	s3Client *s3.Client,
	bucket string,
	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
	batchInfo *batch.BatchUploadDetails,
	cleanupService *foldercleanup.Service,
) error {
	logger.Info("Processing batch",
		"batch_id", batchInfo.ID,
		"client_id", batchInfo.ClientID,
		"filename", batchInfo.OriginalFilename)

	// Batch is already in 'processing' status by default when created

	// Download the zip file from S3
	zipData, err := downloadZipFromS3(ctx, s3Client, bucket, batchInfo.ArchiveKey)
	if err != nil {
		// Mark batch as failed if we can't download the zip
		_ = batchService.MarkFailed(ctx, batchInfo.ID)
		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
		return fmt.Errorf("failed to download zip from S3: %w", err)
	}

	// Extract and process files from the zip
	processedCount := 0
	failedCount := 0
	invalidCount := 0

	// Create a temporary S3 directory for extracted files
	tempDir := fmt.Sprintf("%s_burst", batchInfo.ID.String())
	basePath := path.Dir(batchInfo.ArchiveKey)
	extractPath := path.Join(basePath, tempDir)

	// Open the zip archive
	zipReader, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
	if err != nil {
		_ = batchService.MarkFailed(ctx, batchInfo.ID)
		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
		return fmt.Errorf("failed to open zip archive: %w", err)
	}

	// Process each file in the zip
	for _, file := range zipReader.File {
		processed, failed, invalid, err := processZipFile(ctx, logger, batchService, s3Client, bucket, extractPath, uploadHandler, batchInfo, file)
		if err != nil {
			_ = batchService.MarkFailed(ctx, batchInfo.ID)
			tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
			return fmt.Errorf("failed to process zip file %q: %w", file.Name, err)
		}
		processedCount += processed
		failedCount += failed
		invalidCount += invalid
	}

	totalProcessed := processedCount + failedCount + invalidCount
	if totalProcessed == 0 {
		if err := recordNoProcessableFilesOutcome(ctx, batchService, batchInfo); err != nil {
			_ = batchService.MarkFailed(ctx, batchInfo.ID)
			tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
			return fmt.Errorf("record no-processable-files outcome: %w", err)
		}
		failedCount = 1
		totalProcessed = 1
	}

	// Update batch progress - convert safely to int32
	const maxInt32 = 2147483647
	var processedInt32, failedInt32, invalidInt32 int32
	if processedCount <= maxInt32 {
		processedInt32 = int32(processedCount) // #nosec G115 - checked bounds above
	} else {
		processedInt32 = maxInt32 // Max int32 value
	}
	if failedCount <= maxInt32 {
		failedInt32 = int32(failedCount) // #nosec G115 - checked bounds above
	} else {
		failedInt32 = maxInt32 // Max int32 value
	}
	if invalidCount <= maxInt32 {
		invalidInt32 = int32(invalidCount) // #nosec G115 - checked bounds above
	} else {
		invalidInt32 = maxInt32 // Max int32 value
	}

	// Set total document count now that all files have been extracted
	totalInt32 := processedInt32 + failedInt32 + invalidInt32
	if err := batchService.SetTotalDocuments(ctx, batchInfo.ID, totalInt32); err != nil {
		logger.Error("Failed to set batch total documents", "batch_id", batchInfo.ID, "error", err)
	}

	if err := batchService.UpdateProgress(ctx, batchInfo.ClientID, batchInfo.ID,
		processedInt32, failedInt32, invalidInt32); err != nil {
		logger.Error("Failed to update batch progress", "batch_id", batchInfo.ID, "error", err)
	}

	if totalProcessed > 0 {
		if processedCount > 0 {
			// Mark as completed if at least one document was processed successfully
			if err := batchService.MarkCompleted(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as completed", "batch_id", batchInfo.ID, "error", err)
			}
		} else {
			// Mark as failed only if no documents were processed successfully
			if err := batchService.MarkFailed(ctx, batchInfo.ID); err != nil {
				logger.Error("Failed to mark batch as failed", "batch_id", batchInfo.ID, "error", err)
			}
		}
		tryFinalizeBatchFolders(ctx, logger, cleanupService, batchInfo.ID)
	}

	logger.Info("Batch processing completed",
		"batch_id", batchInfo.ID,
		"processed", processedCount,
		"failed", failedCount,
		"invalid", invalidCount)

	return nil
}

func tryFinalizeBatchFolders(ctx context.Context, logger *slog.Logger, cleanupService *foldercleanup.Service, batchID uuid.UUID) {
	if cleanupService == nil {
		return
	}
	if _, err := cleanupService.TryFinalizeBatch(ctx, batchID); err != nil {
		logger.Error("Failed to finalize batch folder cleanup", "batch_id", batchID, "error", err)
	}
}

// downloadZipFromS3 downloads a zip file from S3
func downloadZipFromS3(ctx context.Context, s3Client *s3.Client, bucket, key string) ([]byte, error) {
	result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to get object from S3: %w", err)
	}
	defer result.Body.Close()

	data, err := io.ReadAll(result.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read S3 object body: %w", err)
	}

	return data, nil
}

// uploadToS3 uploads data to S3
func uploadToS3(ctx context.Context, s3Client *s3.Client, bucket, key string, data []byte) error {
	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(data),
	})
	if err != nil {
		return fmt.Errorf("failed to put object to S3: %w", err)
	}
	return nil
}

// processZipFile processes a single file from the zip archive
// Returns (processedCount, failedCount, invalidCount)
func processZipFile(
	ctx context.Context,
	logger *slog.Logger,
	batchService batchWorkerService,
	s3Client *s3.Client,
	bucket, extractPath string,
	uploadHandler func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error,
	batchInfo *batch.BatchUploadDetails,
	file *zip.File,
) (int, int, int, error) {
	if file.FileInfo().IsDir() {
		return 0, 0, 0, nil
	}

	filename, skip := batchZipFilename(file.Name)
	if skip {
		if filename == "" {
			logger.Warn("Skipping file with invalid path", "original", file.Name)
		}
		return 0, 0, 0, nil
	}

	// Extract file content
	fileReader, err := file.Open()
	if err != nil {
		logger.Error("Failed to open file in zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedOpen, &errMsg, nil); recordErr != nil {
			return 0, 0, 0, fmt.Errorf("record failed_open outcome: %w", recordErr)
		}
		return 0, 1, 0, nil
	}

	fileData, err := io.ReadAll(fileReader)
	fileReader.Close()
	if err != nil {
		logger.Error("Failed to read file from zip", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedRead, &errMsg, nil); recordErr != nil {
			return 0, 0, 0, fmt.Errorf("record failed_read outcome: %w", recordErr)
		}
		return 0, 1, 0, nil
	}

	// Store extracted file in S3 temporary directory (filename now includes sanitized path)
	extractedKey := extractPath + "/" + filename
	if err := uploadToS3(ctx, s3Client, bucket, extractedKey, fileData); err != nil {
		logger.Error("Failed to upload extracted file to S3", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedS3Upload, &errMsg, nil); recordErr != nil {
			return 0, 0, 0, fmt.Errorf("record failed_s3_upload outcome: %w", recordErr)
		}
		return 0, 1, 0, nil
	}

	// Check if file type is supported
	if !isAcceptedExtension(filename) {
		logger.Warn("Invalid file type", "filename", filename)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusInvalidType, nil, nil); recordErr != nil {
			return 0, 0, 0, fmt.Errorf("record invalid_type outcome: %w", recordErr)
		}
		return 0, 0, 1, nil
	}

	// Call the document upload handler with batch ID
	if err := uploadHandler(ctx, batchInfo.ClientID, bytes.NewReader(fileData), filename, &batchInfo.ID); err != nil {
		logger.Error("Failed to upload document", "filename", filename, "error", err)
		_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
		errMsg := err.Error()
		if recordErr := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); recordErr != nil {
			return 0, 0, 0, fmt.Errorf("record failed_upload outcome: %w", recordErr)
		}
		return 0, 1, 0, nil
	}

	if err := recordSuccessfulBatchOutcome(ctx, batchService, batchInfo, filename, fileData); err != nil {
		return 0, 0, 0, err
	}

	logger.Info("Successfully processed file", "filename", filename, "batch_id", batchInfo.ID)
	return 1, 0, 0, nil
}

func recordNoProcessableFilesOutcome(ctx context.Context, batchService batchWorkerService, batchInfo *batch.BatchUploadDetails) error {
	filename := batchInfo.OriginalFilename
	if filename == "" {
		filename = batchInfo.ID.String() + ".zip"
	}
	errMsg := "archive contained no processable files"
	_ = batchService.AddFailedFilename(ctx, batchInfo.ID, filename)
	if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusFailedUpload, &errMsg, nil); err != nil {
		return fmt.Errorf("record failed_upload outcome: %w", err)
	}
	return nil
}

func batchZipFilename(filename string) (string, bool) {
	if strings.HasPrefix(filename, "__MACOSX/") {
		return filename, true
	}
	if strings.HasPrefix(filename, ".") && !strings.HasPrefix(filename, "./") {
		return filename, true
	}
	cleanFilename := sanitizeZipPath(filename)
	return cleanFilename, cleanFilename == ""
}

func recordSuccessfulBatchOutcome(
	ctx context.Context,
	batchService batchWorkerService,
	batchInfo *batch.BatchUploadDetails,
	filename string,
	fileData []byte,
) error {
	hash := sha256.Sum256(fileData)
	hashStr := hex.EncodeToString(hash[:])
	existingID, _ := batchService.CheckDuplicate(ctx, batchInfo.ClientID, hashStr)
	if existingID != nil {
		if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusDuplicate, nil, existingID); err != nil {
			return fmt.Errorf("record duplicate outcome: %w", err)
		}
		return nil
	}
	if err := batchService.RecordOutcome(ctx, batchInfo.ID, filename, repository.BatchOutcomeStatusSubmitted, nil, nil); err != nil {
		return fmt.Errorf("record submitted outcome: %w", err)
	}
	return nil
}