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

package api

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

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document/batch"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/test"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestProcessBatchWork tests the main batch processing work function
func TestProcessBatchWork(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_batch_worker"
	test.CreateTestClient(t, cfg, clientID, "Test Batch Worker Client")

	// Create a test batch service
	batchService := batch.New(cfg)

	// Upload a test ZIP file to S3
	zipContent := createTestZIPForWorker(t, 2)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()
	archiveKey := fmt.Sprintf("test/%s/batch.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create a batch record with totalDocs=0 to match production (documents.go:236).
	// The worker must call SetTotalDocuments after extraction for progress to work.
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "batch.zip", 0, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Track uploaded documents
	uploadedDocs := make(map[string][]byte)
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		data, err := io.ReadAll(docData)
		if err != nil {
			return err
		}
		uploadedDocs[filename] = data
		return nil
	}

	// Create worker config
	workerConfig := map[string]any{
		ConfigKeyBatchService:  batchService,
		ConfigKeyS3Client:      s3Client,
		ConfigKeyBucket:        bucket,
		ConfigKeyUploadHandler: uploadHandler,
	}

	// Execute the batch work
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processBatchWork(ctx, logger, workerConfig)
	require.NoError(t, err)

	// Verify that documents were uploaded
	assert.Len(t, uploadedDocs, 2)
	assert.Contains(t, uploadedDocs, "document_1.pdf")
	assert.Contains(t, uploadedDocs, "document_2.pdf")

	// Verify batch status was updated
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status)
	assert.Equal(t, int32(2), batchInfo.TotalDocuments, "TotalDocuments should be set after processing")

	// Summary counts are derived from outcome states. After extraction, outcomes
	// are in 'submitted' (non-terminal) state -- pipeline hasn't run yet.
	// Only terminal outcomes (clean_passed, clean_failed, etc.) are counted.
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	assert.Len(t, outcomes, 2, "should have 2 outcome rows")
	for _, o := range outcomes {
		assert.Equal(t, "submitted", o.Outcome)
	}
}

// TestProcessSingleBatch tests processing a single batch
func TestProcessSingleBatch(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_single_batch"
	test.CreateTestClient(t, cfg, clientID, "Test Single Batch Client")

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP with mixed content
	zipContent := createMixedContentZIP(t)
	archiveKey := fmt.Sprintf("test/%s/mixed.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed.zip", 3, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Mock upload handler
	uploadCount := 0
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		uploadCount++
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "mixed.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// All three file types (PDF, TXT, DOCX) are now accepted
	assert.Equal(t, 3, uploadCount)

	// Verify batch status
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status)
	assert.Equal(t, int32(3), batchInfo.TotalDocuments)

	// After extraction, outcomes are 'submitted' (non-terminal).
	// Derived counts only reflect terminal outcomes.
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	assert.Len(t, outcomes, 3)
	for _, o := range outcomes {
		assert.Equal(t, "submitted", o.Outcome)
	}
}

// TestProcessSingleBatchWithFailure tests handling of upload failures
func TestProcessSingleBatchWithFailure(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_batch_failure"
	test.CreateTestClient(t, cfg, clientID, "Test Batch Failure Client")

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP
	zipContent := createTestZIPForWorker(t, 3)
	archiveKey := fmt.Sprintf("test/%s/fail.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "fail.zip", 3, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Upload handler that fails on second document
	uploadCount := 0
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		uploadCount++
		if uploadCount == 2 {
			return fmt.Errorf("simulated upload failure")
		}
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "fail.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify results
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status) // Status is still completed since some succeeded
	assert.Contains(t, batchInfo.FailedFilenames, "document_2.pdf")

	// Derived counts: 2 submitted (non-terminal), 1 failed_upload (terminal)
	assert.Equal(t, int32(0), batchInfo.ProcessedDocuments, "submitted is non-terminal")
	assert.Equal(t, int32(1), batchInfo.FailedDocuments, "2nd file failed upload")
}

// TestDownloadZipFromS3 tests downloading a ZIP file from S3
func TestDownloadZipFromS3(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateAWSResources(t, cfg)

	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test data
	testData := []byte("test zip content")
	key := "test/download.zip"

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(testData),
	})
	require.NoError(t, err)

	// Download the data
	downloaded, err := downloadZipFromS3(ctx, s3Client, bucket, key)
	require.NoError(t, err)
	assert.Equal(t, testData, downloaded)
}

// TestUploadToS3 tests uploading data to S3
func TestUploadToS3(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateAWSResources(t, cfg)

	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test data
	testData := []byte("test upload content")
	key := "test/upload.txt"

	err := uploadToS3(ctx, s3Client, bucket, key, testData)
	require.NoError(t, err)

	// Verify the upload
	result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	require.NoError(t, err)
	defer result.Body.Close()

	downloaded, err := io.ReadAll(result.Body)
	require.NoError(t, err)
	assert.Equal(t, testData, downloaded)
}

// Helper function to create a test ZIP file with PDFs
func createTestZIPForWorker(t *testing.T, numFiles int) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	for i := 1; i <= numFiles; i++ {
		filename := fmt.Sprintf("document_%d.pdf", i)
		fileWriter, err := zipWriter.Create(filename)
		require.NoError(t, err)

		pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
		_, err = fileWriter.Write([]byte(pdfContent))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}

// Helper function to create a ZIP with mixed content types
func createMixedContentZIP(t *testing.T) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Add a PDF
	pdfWriter, err := zipWriter.Create("document.pdf")
	require.NoError(t, err)
	_, err = pdfWriter.Write([]byte("%%PDF-1.4\n%%Test PDF\n%%%%EOF"))
	require.NoError(t, err)

	// Add a TXT file
	txtWriter, err := zipWriter.Create("notes.txt")
	require.NoError(t, err)
	_, err = txtWriter.Write([]byte("This is a text file"))
	require.NoError(t, err)

	// Add a DOCX file
	docxWriter, err := zipWriter.Create("report.docx")
	require.NoError(t, err)
	_, err = docxWriter.Write([]byte("PK")) // Simplified DOCX header
	require.NoError(t, err)

	err = zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}

// Helper function to create a ZIP with nested folder structure
func createNestedFolderZIP(t *testing.T) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create files in various nested folder structures
	testFiles := []struct {
		path    string
		content string
	}{
		{"document1.pdf", "%%PDF-1.4\n%%Root level PDF\n%%%%EOF"},
		{"foldera/file1.pdf", "%%PDF-1.4\n%%FolderA File1\n%%%%EOF"},
		{"foldera/file2.pdf", "%%PDF-1.4\n%%FolderA File2\n%%%%EOF"},
		{"folderb/folderc/file3.pdf", "%%PDF-1.4\n%%Nested File3\n%%%%EOF"},
		{"folderb/folderc/file4.pdf", "%%PDF-1.4\n%%Nested File4\n%%%%EOF"},
		{"folderb/file5.pdf", "%%PDF-1.4\n%%FolderB File5\n%%%%EOF"},
		{"deep/nested/structure/test/file6.pdf", "%%PDF-1.4\n%%Deep nested\n%%%%EOF"},
		// Add a non-PDF to verify it's skipped
		{"foldera/readme.txt", "This should be skipped"},
	}

	for _, tf := range testFiles {
		fileWriter, err := zipWriter.Create(tf.path)
		require.NoError(t, err)
		_, err = fileWriter.Write([]byte(tf.content))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}

// Helper function to create a ZIP with path traversal attempts
func createMaliciousPathZIP(t *testing.T) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create files with various path traversal attempts
	testFiles := []struct {
		path    string
		content string
	}{
		{"../escape.pdf", "%%PDF-1.4\n%%Escape attempt\n%%%%EOF"},
		{"../../etc/passwd", "malicious content"},
		{"./valid.pdf", "%%PDF-1.4\n%%Valid PDF\n%%%%EOF"},
		{"folder/../sibling.pdf", "%%PDF-1.4\n%%Sibling\n%%%%EOF"},
		{"normal/file.pdf", "%%PDF-1.4\n%%Normal\n%%%%EOF"},
	}

	for _, tf := range testFiles {
		// Note: zip.Writer doesn't validate paths, so we can create these entries
		header := &zip.FileHeader{
			Name:   tf.path,
			Method: zip.Deflate,
		}
		fileWriter, err := zipWriter.CreateHeader(header)
		require.NoError(t, err)
		_, err = fileWriter.Write([]byte(tf.content))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}

// TestSanitizeZipPath tests the path sanitization function
func TestSanitizeZipPath(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{"simple file", "document.pdf", "document.pdf"},
		{"file in folder", "folder/document.pdf", "folder/document.pdf"},
		{"nested folders", "folder/subfolder/document.pdf", "folder/subfolder/document.pdf"},
		{"deep nesting", "a/b/c/d/e/f/document.pdf", "a/b/c/d/e/f/document.pdf"},
		{"dot prefix cleaned", "./document.pdf", "document.pdf"},
		{"dot in path cleaned", "folder/./document.pdf", "folder/document.pdf"},
		{"parent traversal rejected", "../document.pdf", ""},
		{"parent in middle cleaned", "folder/../document.pdf", "document.pdf"},
		{"absolute path rejected", "/etc/passwd", ""},
		{"backslash rejected", "folder\\document.pdf", ""},
		{"double parent rejected", "../../document.pdf", ""},
		{"complex traversal attempt", "folder/../../etc/passwd", ""},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := sanitizeZipPath(tt.input)
			assert.Equal(t, tt.expected, result, "sanitizeZipPath(%q) = %q, want %q", tt.input, result, tt.expected)
		})
	}
}

// TestProcessBatchWithNestedFolders tests processing a ZIP with nested folder structure
func TestProcessBatchWithNestedFolders(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_nested_folders"
	test.CreateTestClient(t, cfg, clientID, "Test Nested Folders Client")

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP with nested folders
	zipContent := createNestedFolderZIP(t)
	archiveKey := fmt.Sprintf("test/%s/nested.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "nested.zip", 8, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Use the real upload handler that creates documents in the database
	uploadHandler := createDocumentUploadHandler(cfg)

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "nested.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify batch status
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status)
	assert.Equal(t, int32(8), batchInfo.TotalDocuments, "7 PDFs + 1 TXT")
	// After extraction, all outcomes are 'submitted' (non-terminal).
	// Derived ProcessedDocuments only counts terminal outcomes.
	assert.Equal(t, int32(0), batchInfo.InvalidTypeDocuments, "all types accepted")

	// Manually trigger document initialization for uploaded files since S3 notifications
	// don't work in test environment. Get all document uploads for this batch.
	uploads, err := cfg.GetDBQueries().ListDocumentUploads(ctx, clientID)
	require.NoError(t, err)

	// Process each uploaded file through the document initialization pipeline
	for _, upload := range uploads {
		if upload.BatchID != nil && *upload.BatchID == batchID {
			// Parse the bucket key
			key, err := objectstore.ParseBucketKey(upload.Key)
			require.NoError(t, err)

			// Get file hash from S3 object metadata (normally provided by S3 event)
			headResult, err := s3Client.HeadObject(ctx, &s3.HeadObjectInput{
				Bucket: aws.String(upload.Bucket),
				Key:    aws.String(upload.Key),
			})
			require.NoError(t, err)

			// Create document directly in database since full pipeline requires additional setup
			_, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
				Clientid: key.ClientID,
				Hash:     *headResult.ETag, // Use ETag as hash for testing
				BatchID:  upload.BatchID,
				Filename: upload.Filename,
			})
			require.NoError(t, err)
		}
	}

	// Get documents for this specific batch from database
	batchDocs, err := cfg.GetDBQueries().ListDocumentsByBatch(ctx, &batchID)
	require.NoError(t, err)

	// Verify we have the expected number of documents (7 PDFs + 1 TXT)
	assert.Equal(t, 8, len(batchDocs), "Should have exactly 8 documents in database")

	// Verify filenames are preserved with their paths
	foundFilenames := make(map[string]bool)
	for _, doc := range batchDocs {
		if doc.Filename != nil {
			foundFilenames[*doc.Filename] = true
		}
	}

	// Check that we found all expected filenames with their paths
	expectedFilenames := []string{
		"document1.pdf",
		"foldera/file1.pdf",
		"foldera/file2.pdf",
		"folderb/folderc/file3.pdf",
		"folderb/folderc/file4.pdf",
		"folderb/file5.pdf",
		"deep/nested/structure/test/file6.pdf",
		"foldera/readme.txt",
	}

	for _, expectedFilename := range expectedFilenames {
		assert.True(t, foundFilenames[expectedFilename],
			"Should have preserved filename with path: %s", expectedFilename)
	}
}

// TestProcessBatchWithPathTraversal tests that path traversal attempts are handled safely
func TestProcessBatchWithPathTraversal(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_path_traversal"
	test.CreateTestClient(t, cfg, clientID, "Test Path Traversal Client")

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP with path traversal attempts
	zipContent := createMaliciousPathZIP(t)
	archiveKey := fmt.Sprintf("test/%s/malicious.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "malicious.zip", 5, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Track uploaded documents
	uploadedDocs := make(map[string][]byte)
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		data, err := io.ReadAll(docData)
		if err != nil {
			return err
		}
		uploadedDocs[filename] = data
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "malicious.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify that path traversal attempts were blocked
	assert.NotContains(t, uploadedDocs, "../escape.pdf", "Should reject parent traversal")
	assert.NotContains(t, uploadedDocs, "../../etc/passwd", "Should reject multiple parent traversal")
	assert.NotContains(t, uploadedDocs, "/etc/passwd", "Should reject absolute paths")

	// Verify that safe files were processed correctly
	assert.Contains(t, uploadedDocs, "valid.pdf", "Should process ./valid.pdf as valid.pdf")
	assert.Contains(t, uploadedDocs, "sibling.pdf", "Should clean folder/../sibling.pdf to sibling.pdf")
	assert.Contains(t, uploadedDocs, "normal/file.pdf", "Should process normal paths correctly")

	// Verify only safe PDFs were processed
	assert.Equal(t, 3, len(uploadedDocs), "Should have processed exactly 3 safe PDFs")
}

// TestProcessZipFile_RecordsOutcomes tests that processZipFile records per-file
// outcomes for each extracted file.
func TestProcessZipFile_RecordsOutcomes(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	clientID := "test_outcomes_record"
	test.CreateTestClient(t, cfg, clientID, "Test Outcomes Record Client")

	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok)
	bucket := cfg.GetBucket()

	// Create ZIP with 2 PDFs
	zipContent := createTestZIPForWorker(t, 2)
	archiveKey := fmt.Sprintf("test/%s/outcomes.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	batchID, err := batchService.CreateWithStorage(ctx, clientID, "outcomes.zip", 2, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		_, _ = io.ReadAll(docData)
		return nil
	}

	workerConfig := map[string]any{
		ConfigKeyBatchService:  batchService,
		ConfigKeyS3Client:      s3Client,
		ConfigKeyBucket:        bucket,
		ConfigKeyUploadHandler: uploadHandler,
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processBatchWork(ctx, logger, workerConfig)
	require.NoError(t, err)

	// Verify outcomes were recorded for each file
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	assert.Len(t, outcomes, 2)

	for _, o := range outcomes {
		assert.Equal(t, "submitted", o.Outcome, "each PDF should be recorded as submitted")
		assert.Contains(t, o.Filename, ".pdf")
	}
}

// TestProcessZipFile_DetectsDuplicates verifies that when a file in the batch
// has the same content hash as an existing document in the system, it is
// recorded as a duplicate with the existing document's ID.
func TestProcessZipFile_DetectsDuplicates(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	clientID := "test_outcomes_dup"
	test.CreateTestClient(t, cfg, clientID, "Test Outcomes Dup Client")

	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok)
	bucket := cfg.GetBucket()

	// Pre-create an existing document with a known hash
	knownContent := []byte("%%PDF-1.4\n%%Already existing content\n%%%%EOF")
	knownHash := sha256Hash(knownContent)
	existingDocID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     knownHash,
	})
	require.NoError(t, err)

	// Create a ZIP with one file that has the same content as the existing doc
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)
	w, err := zipWriter.Create("duplicate_of_existing.pdf")
	require.NoError(t, err)
	_, err = w.Write(knownContent)
	require.NoError(t, err)
	require.NoError(t, zipWriter.Close())
	zipContent := buf.Bytes()

	archiveKey := fmt.Sprintf("test/%s/dup.zip", clientID)
	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	batchID, err := batchService.CreateWithStorage(ctx, clientID, "dup.zip", 1, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
		_, _ = io.ReadAll(docData)
		return nil
	}

	workerConfig := map[string]any{
		ConfigKeyBatchService:  batchService,
		ConfigKeyS3Client:      s3Client,
		ConfigKeyBucket:        bucket,
		ConfigKeyUploadHandler: uploadHandler,
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processBatchWork(ctx, logger, workerConfig)
	require.NoError(t, err)

	// Verify outcome is "duplicate" with the existing document's ID
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	require.Len(t, outcomes, 1)

	assert.Equal(t, "duplicate", outcomes[0].Outcome)
	require.NotNil(t, outcomes[0].DocumentID)
	assert.Equal(t, existingDocID, *outcomes[0].DocumentID)
}

// TestProcessZipFile_MixedOutcomes tests a ZIP with PDF + non-PDF + paths to
// verify all outcome types (submitted, invalid_type) are correctly recorded.
func TestProcessZipFile_MixedOutcomes(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	clientID := "test_outcomes_mixed"
	test.CreateTestClient(t, cfg, clientID, "Test Outcomes Mixed Client")

	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok)
	bucket := cfg.GetBucket()

	// Create mixed ZIP
	zipContent := createMixedContentZIP(t)
	archiveKey := fmt.Sprintf("test/%s/mixed_outcomes.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed_outcomes.zip", 3, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error {
		_, _ = io.ReadAll(docData)
		return nil
	}

	workerConfig := map[string]any{
		ConfigKeyBatchService:  batchService,
		ConfigKeyS3Client:      s3Client,
		ConfigKeyBucket:        bucket,
		ConfigKeyUploadHandler: uploadHandler,
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processBatchWork(ctx, logger, workerConfig)
	require.NoError(t, err)

	// Verify outcomes
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	assert.Len(t, outcomes, 3) // 1 PDF + 1 txt + 1 docx

	outcomeMap := map[string]string{}
	for _, o := range outcomes {
		outcomeMap[o.Filename] = o.Outcome
	}

	// All three types are now accepted
	assert.Equal(t, "submitted", outcomeMap["document.pdf"])
	assert.Equal(t, "submitted", outcomeMap["notes.txt"])
	assert.Equal(t, "submitted", outcomeMap["report.docx"])
}

// sha256Hash computes the hex-encoded SHA-256 hash of data.
func sha256Hash(data []byte) string {
	h := sha256.Sum256(data)
	return hex.EncodeToString(h[:])
}

// TestProcessBatchWithDuplicateFilenames tests handling of duplicate filenames in different folders
func TestProcessBatchWithDuplicateFilenames(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client with root folder
	clientID := "test_duplicate_names"
	test.CreateTestClient(t, cfg, clientID, "Test Duplicate Names Client")

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Create ZIP with duplicate filenames in different folders
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create files with same names in different folders
	testFiles := []struct {
		path    string
		content string
	}{
		{"document.pdf", "%%PDF-1.4\n%%Root document\n%%%%EOF"},
		{"folder1/document.pdf", "%%PDF-1.4\n%%Folder1 document\n%%%%EOF"},
		{"folder2/document.pdf", "%%PDF-1.4\n%%Folder2 document\n%%%%EOF"},
		{"folder1/subfolder/document.pdf", "%%PDF-1.4\n%%Subfolder document\n%%%%EOF"},
	}

	for _, tf := range testFiles {
		fileWriter, err := zipWriter.Create(tf.path)
		require.NoError(t, err)
		_, err = fileWriter.Write([]byte(tf.content))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)
	zipContent := buf.Bytes()

	// Upload the ZIP
	archiveKey := fmt.Sprintf("test/%s/duplicates.zip", clientID)
	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "duplicates.zip", 4, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Track uploaded documents
	uploadedDocs := make(map[string]string)
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		data, err := io.ReadAll(docData)
		if err != nil {
			return err
		}
		// Extract the content identifier from the PDF
		content := string(data)
		uploadedDocs[filename] = content
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "duplicates.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify all files were uploaded with correct paths
	assert.Equal(t, 4, len(uploadedDocs), "Should have uploaded all 4 PDFs")

	// Verify each file has unique path and correct content
	assert.Contains(t, uploadedDocs["document.pdf"], "Root document")
	assert.Contains(t, uploadedDocs["folder1/document.pdf"], "Folder1 document")
	assert.Contains(t, uploadedDocs["folder2/document.pdf"], "Folder2 document")
	assert.Contains(t, uploadedDocs["folder1/subfolder/document.pdf"], "Subfolder document")
}

// TestProcessBatchWithAllSupportedFileTypes tests the full multi-format batch upload pipeline.
// A batch ZIP containing one file of each supported type is uploaded and every document
// ends up in the system as a first-class document with an ID. This is the critical
// integration test for multi-format support.
func TestProcessBatchWithAllSupportedFileTypes(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	clientID := "test_multiformat_batch"
	test.CreateTestClient(t, cfg, clientID, "Test Multi-Format Batch Client")

	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok)
	bucket := cfg.GetBucket()

	// Build a ZIP containing one file of each supported type in folder structure
	zipContent := createMultiFormatZIP(t)
	archiveKey := fmt.Sprintf("test/%s/multiformat.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// 10 files total: PDF, TIFF, JPEG, PNG, BMP, DOCX, XLSX, PPTX, TXT, EML
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "multiformat.zip", 10, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Use the real upload handler that persists to the database
	uploadHandler := createDocumentUploadHandler(cfg)

	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "multiformat.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify batch completed successfully
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status, "batch should complete")
	assert.Equal(t, int32(10), batchInfo.TotalDocuments, "all 10 files should be extracted")
	// After extraction, outcomes are 'submitted' (non-terminal). Derived counts
	// only reflect terminal outcomes so ProcessedDocuments is 0 at this stage.
	assert.Equal(t, int32(0), batchInfo.FailedDocuments, "no files should fail")
	assert.Equal(t, int32(0), batchInfo.InvalidTypeDocuments, "no files should be invalid type")

	// Verify outcomes: every file should be submitted
	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	assert.Len(t, outcomes, 10, "should have 10 outcomes")

	for _, o := range outcomes {
		assert.Equal(t, "submitted", o.Outcome,
			"file %s should be submitted, got %s", o.Filename, o.Outcome)
	}

	// Verify folder structure is preserved in outcomes
	outcomeFilenames := make(map[string]bool)
	for _, o := range outcomes {
		outcomeFilenames[o.Filename] = true
	}

	expectedFiles := []string{
		"documents/report.pdf",
		"images/scan.tiff",
		"images/photo.jpeg",
		"images/diagram.png",
		"images/legacy.bmp",
		"office/proposal.docx",
		"office/data.xlsx",
		"office/slides.pptx",
		"text/notes.txt",
		"email/message.eml",
	}
	for _, f := range expectedFiles {
		assert.True(t, outcomeFilenames[f], "should have outcome for %s", f)
	}
}

// TestProcessBatchWithMixedValidInvalidFileTypes tests a batch with valid new-format
// files and unsupported files, verifying that valid files succeed and unsupported
// files get appropriate rejection without aborting the batch.
func TestProcessBatchWithMixedValidInvalidFileTypes(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &BaseConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	clientID := "test_mixed_valid_invalid"
	test.CreateTestClient(t, cfg, clientID, "Test Mixed Valid Invalid Client")

	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok)
	bucket := cfg.GetBucket()

	// Create ZIP with mix of valid and unsupported files
	zipContent := createMixedValidInvalidZIP(t)
	archiveKey := fmt.Sprintf("test/%s/mixed_valid_invalid.zip", clientID)

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed_valid_invalid.zip", 5, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	uploadHandler := func(_ context.Context, _ string, docData io.Reader, _ string, _ *uuid.UUID) error {
		_, _ = io.ReadAll(docData)
		return nil
	}

	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "mixed_valid_invalid.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Batch should still complete (not fail) because some files succeeded
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status)
	// Derived counts: 3 submitted (non-terminal), 2 invalid_type (terminal)
	assert.Equal(t, int32(0), batchInfo.ProcessedDocuments, "submitted is non-terminal")
	assert.Equal(t, int32(2), batchInfo.InvalidTypeDocuments, "2 unsupported files")

	outcomes, err := batchService.ListOutcomes(ctx, batchID)
	require.NoError(t, err)

	outcomeMap := map[string]string{}
	for _, o := range outcomes {
		outcomeMap[o.Filename] = o.Outcome
	}

	// Valid files submitted
	assert.Equal(t, "submitted", outcomeMap["report.pdf"])
	assert.Equal(t, "submitted", outcomeMap["notes.txt"])
	assert.Equal(t, "submitted", outcomeMap["photo.png"])

	// Unsupported files rejected
	assert.Equal(t, "invalid_type", outcomeMap["malware.exe"])
	assert.Equal(t, "invalid_type", outcomeMap["archive.zip"])
}

// createMultiFormatZIP creates a ZIP with one valid file of each supported type
// organized in folders.
func createMultiFormatZIP(t *testing.T) []byte {
	t.Helper()

	// Read fixture files from the types testdata directory
	fixtureDir := "../../document/types/testdata/multiformat"

	files := []struct {
		zipPath     string
		fixtureName string
	}{
		{"documents/report.pdf", ""}, // PDF needs to be inline (no fixture)
		{"images/scan.tiff", "valid.tiff"},
		{"images/photo.jpeg", "valid.jpeg"},
		{"images/diagram.png", "valid.png"},
		{"images/legacy.bmp", "valid.bmp"},
		{"office/proposal.docx", "valid.docx"},
		{"office/data.xlsx", "valid.xlsx"},
		{"office/slides.pptx", "valid.pptx"},
		{"text/notes.txt", "valid.txt"},
		{"email/message.eml", "valid.eml"},
	}

	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	for _, f := range files {
		var data []byte
		var err error

		if f.fixtureName == "" {
			// Inline PDF content
			data = []byte("%%PDF-1.4\n%%Multi-format test PDF\n%%%%EOF")
		} else {
			data, err = os.ReadFile(fmt.Sprintf("%s/%s", fixtureDir, f.fixtureName))
			require.NoError(t, err, "failed to read fixture %s", f.fixtureName)
		}

		w, err := zipWriter.Create(f.zipPath)
		require.NoError(t, err)
		_, err = w.Write(data)
		require.NoError(t, err)
	}

	require.NoError(t, zipWriter.Close())
	return buf.Bytes()
}

// createMixedValidInvalidZIP creates a ZIP with valid supported files and unsupported files.
func createMixedValidInvalidZIP(t *testing.T) []byte {
	t.Helper()

	files := []struct {
		name    string
		content []byte
	}{
		{"report.pdf", []byte("%%PDF-1.4\n%%Test\n%%%%EOF")},
		{"notes.txt", []byte("This is a valid text file")},
		{"photo.png", nil}, // will be loaded from fixture
		{"malware.exe", []byte("MZ\x90\x00\x03\x00")}, // PE signature
		{"archive.zip", []byte("Not a real zip but has .zip extension")},
	}

	// Load PNG fixture
	pngData, err := os.ReadFile("../../document/types/testdata/multiformat/valid.png")
	require.NoError(t, err)
	files[2].content = pngData

	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	for _, f := range files {
		w, err := zipWriter.Create(f.name)
		require.NoError(t, err)
		_, err = w.Write(f.content)
		require.NoError(t, err)
	}

	require.NoError(t, zipWriter.Close())
	return buf.Bytes()
}