query-orchestration/internal/document/batch/service_test.go

package batch_test

import (
	"archive/zip"
	"bytes"
	"context"
	"fmt"
	"io"
	"testing"

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

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

type TestConfig struct {
	serviceconfig.BaseConfig
	objectstore.ObjectStoreConfig
}

func TestUpdateProgress(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_progress"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Progress",
	})
	require.NoError(t, err)

	// Create test batch with total documents
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test updating progress
	err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
	require.NoError(t, err)

	// Verify progress was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, int32(5), batch.ProcessedDocuments)
	assert.Equal(t, int32(2), batch.FailedDocuments)
	assert.Equal(t, int32(1), batch.InvalidTypeDocuments)
	assert.Equal(t, int32(80), batch.ProgressPercent) // (5+2+1)/10 * 100 = 80%
}

func TestUpdateProgressZeroTotalDocuments(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_zero"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Zero",
	})
	require.NoError(t, err)

	// Create test batch with zero total documents
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   0,
	})
	require.NoError(t, err)

	// Test updating progress should fail with zero total documents
	err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "batch has zero total documents")
}

func TestMarkCompleted(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_complete"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Complete",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test marking as completed
	err = service.MarkCompleted(t.Context(), batchID)
	require.NoError(t, err)

	// Verify status was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusCompleted, batch.Status)
}

func TestMarkFailed(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_failed"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Failed",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test marking as failed
	err = service.MarkFailed(t.Context(), batchID)
	require.NoError(t, err)

	// Verify status was updated
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, repository.BatchStatusFailed, batch.Status)
}

func TestAddFailedFilename(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_filename"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Filename",
	})
	require.NoError(t, err)

	// Create test batch
	batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: "test.zip",
		TotalDocuments:   10,
	})
	require.NoError(t, err)

	// Test adding failed filename
	err = service.AddFailedFilename(t.Context(), batchID, "failed_file.pdf")
	require.NoError(t, err)

	// Get the full details to check failed filenames
	details, err := service.Get(t.Context(), clientID, batchID)
	require.NoError(t, err)
	assert.Contains(t, details.FailedFilenames, "failed_file.pdf")
}

func TestCreateWithStorage_Success(t *testing.T) {
	// Setup: Create database testcontainer
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	service := documentbatch.New(cfg)

	// Create test client
	clientID := "test_client_storage"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Storage",
	})
	require.NoError(t, err)

	// Test: Create batch with storage metadata
	batchID, err := service.CreateWithStorage(
		t.Context(),
		clientID,
		"test.zip",
		0, // totalDocs unknown in Part 1
		"test-bucket",
		"batches/test-client/2025/01/01/test.zip",
		1024*1024, // 1MB
	)

	// Verify: Batch created successfully
	require.NoError(t, err)
	assert.NotEmpty(t, batchID)

	// Verify: Batch can be retrieved with storage metadata
	batch, err := service.Get(t.Context(), clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "test.zip", batch.OriginalFilename)
	assert.Equal(t, documentbatch.StatusProcessing, batch.Status)
	assert.Equal(t, "test-bucket", batch.ArchiveBucket)
	assert.Equal(t, "batches/test-client/2025/01/01/test.zip", batch.ArchiveKey)
	assert.Equal(t, int64(1024*1024), batch.FileSizeBytes)
}

func TestCreateWithStorage_DatabaseFailure(t *testing.T) {
	// Setup: Create database
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	_ = documentbatch.New(cfg) // service unused in skipped test

	// Create test client
	clientID := "test_client_storage_fail"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Storage Fail",
	})
	require.NoError(t, err)

	// Note: Skipping database connection close test as Pool interface doesn't expose Close method
	// This test would require mocking the database layer to simulate connection failures
	t.Skip("Database connection failure testing requires mocking")
}

// TestCreate_Success tests the Create function with a valid ZIP file uploaded to S3
func TestCreate_Success(t *testing.T) {
	// Setup: Create database and AWS resources with testcontainers
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create batch service and storage service
	batchService := documentbatch.New(cfg)
	storageService := storage.New(cfg)

	// Create test client
	clientID := "test_client_create"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Create",
	})
	require.NoError(t, err)

	// Create test ZIP content with 3 PDFs
	zipContent := createTestZIPContent(t, 3)

	// First upload the ZIP to S3 (simulating what would happen in the full flow)
	storageResult, err := storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)
	require.NoError(t, err)

	// Test: Use the Create function to create a batch record
	batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 3)

	// Verify: Batch created successfully
	require.NoError(t, err)
	assert.NotEmpty(t, batchID)

	// Verify: Batch can be retrieved
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, "test.zip", batch.OriginalFilename)
	assert.Equal(t, int32(3), batch.TotalDocuments)
	assert.Equal(t, clientID, batch.ClientID)
	assert.Equal(t, repository.BatchStatusProcessing, batch.Status)

	// Verify: File still exists in S3
	s3Client := cfg.GetStoreClient()
	_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
		Bucket: &storageResult.Bucket,
		Key:    &storageResult.Key,
	})
	require.NoError(t, err, "ZIP file should exist in S3")
}

// TestCreate_InvalidZIP tests Create function with invalid ZIP data
func TestCreate_InvalidZIP(t *testing.T) {
	// Setup: Create database and AWS resources
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	batchService := documentbatch.New(cfg)
	storageService := storage.New(cfg)

	// Create test client
	clientID := "test_client_invalid_zip"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Invalid ZIP",
	})
	require.NoError(t, err)

	// Create invalid ZIP content (just plain text, not a ZIP)
	invalidContent := bytes.NewReader([]byte("This is not a valid ZIP file"))

	// Test: Try to validate before storing (storage service should validate)
	err = storageService.ValidateZIPFile("notazip.txt", int64(invalidContent.Len()))
	require.Error(t, err)
	assert.Contains(t, err.Error(), "invalid file type")

	// Even if we bypass validation and try with .zip extension, the content is invalid
	// The Create function itself doesn't validate ZIP content, it just creates a DB record
	// So we test that Create works but processing would fail later
	batchID, err := batchService.Create(t.Context(), clientID, "fake.zip", 1)

	// The Create function should succeed (it only creates DB record)
	require.NoError(t, err)
	assert.NotEmpty(t, batchID)

	// Verify the batch exists but would fail during actual processing
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, "fake.zip", batch.OriginalFilename)
}

// TestCreate_S3UploadFailure tests Create function when S3 upload fails
func TestCreate_S3UploadFailure(t *testing.T) {
	// Setup: Create database and AWS resources
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	batchService := documentbatch.New(cfg)
	storageService := storage.New(cfg)

	// Create test client
	clientID := "test_client_s3_fail"
	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client S3 Fail",
	})
	require.NoError(t, err)

	// Create valid ZIP content
	zipContent := createTestZIPContent(t, 2)

	// Temporarily break S3 configuration to simulate upload failure
	// Save original bucket and set invalid one
	originalBucket := cfg.GetBucket()
	cfg.SetBucket("non-existent-bucket-that-should-not-exist")

	// Test: Try to upload to S3 with invalid bucket
	_, err = storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)

	// Verify: S3 upload should fail
	require.Error(t, err)
	assert.Contains(t, err.Error(), "failed to upload ZIP to S3")

	// Restore original bucket for cleanup
	cfg.SetBucket(originalBucket)

	// The Create function itself is independent of S3, so it would still work
	// This demonstrates separation of concerns
	batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 2)
	require.NoError(t, err)
	assert.NotEmpty(t, batchID)

	// Verify batch exists in DB even though S3 upload failed
	batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, "test.zip", batch.OriginalFilename)
	assert.Equal(t, int32(2), batch.TotalDocuments)
}

// Helper function to create test ZIP content with specified number of PDF files
func createTestZIPContent(t *testing.T, numFiles int) io.Reader {
	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)

		// Write minimal PDF content
		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 bytes.NewReader(buf.Bytes())
}

func TestListUnprocessed(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)
	ctx := t.Context()

	// Create multiple test clients
	clientID1 := "test_client_unprocessed1"
	clientID2 := "test_client_unprocessed2"

	for _, clientID := range []string{clientID1, clientID2} {
		err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
			Clientid: clientID,
			Name:     fmt.Sprintf("Test Client %s", clientID),
		})
		require.NoError(t, err)
	}

	// Create processing batches (should be returned)
	var expectedBatchIDs []string

	// Client 1: 2 processing batches
	for i := 0; i < 2; i++ {
		batchID, err := service.Create(ctx, clientID1, fmt.Sprintf("processing%d.zip", i), 10)
		require.NoError(t, err)
		expectedBatchIDs = append(expectedBatchIDs, batchID.String())
	}

	// Client 2: 1 processing batch
	batchID, err := service.Create(ctx, clientID2, "processing_client2.zip", 5)
	require.NoError(t, err)
	expectedBatchIDs = append(expectedBatchIDs, batchID.String())

	// Create completed batch (should NOT be returned)
	completedBatchID, err := service.Create(ctx, clientID1, "completed.zip", 8)
	require.NoError(t, err)
	err = service.MarkCompleted(ctx, completedBatchID)
	require.NoError(t, err)

	// Create failed batch (should NOT be returned)
	failedBatchID, err := service.Create(ctx, clientID2, "failed.zip", 3)
	require.NoError(t, err)
	err = service.MarkFailed(ctx, failedBatchID)
	require.NoError(t, err)

	// Test: Get unprocessed batches
	unprocessedBatches, err := service.ListUnprocessed(ctx)
	require.NoError(t, err)

	// Verify: Should return exactly 3 processing batches
	assert.Len(t, unprocessedBatches, 3)

	// Verify: All returned batches have processing status
	actualBatchIDs := make([]string, len(unprocessedBatches))
	for i, batch := range unprocessedBatches {
		assert.Equal(t, documentbatch.StatusProcessing, batch.Status)
		actualBatchIDs[i] = batch.ID.String()
	}

	// Verify: All processing batches are returned
	for _, expectedID := range expectedBatchIDs {
		assert.Contains(t, actualBatchIDs, expectedID)
	}

	// Verify: Completed and failed batches are NOT returned
	assert.NotContains(t, actualBatchIDs, completedBatchID.String())
	assert.NotContains(t, actualBatchIDs, failedBatchID.String())

	// Verify: Batches from both clients are returned
	clientIDs := make(map[string]bool)
	for _, batch := range unprocessedBatches {
		clientIDs[batch.ClientID] = true
	}
	assert.True(t, clientIDs[clientID1], "should include batches from client1")
	assert.True(t, clientIDs[clientID2], "should include batches from client2")

	// Verify: Batches are ordered by created_at ASC (oldest first)
	if len(unprocessedBatches) > 1 {
		for i := 0; i < len(unprocessedBatches)-1; i++ {
			assert.True(t,
				unprocessedBatches[i].CreatedAt.Before(unprocessedBatches[i+1].CreatedAt) ||
					unprocessedBatches[i].CreatedAt.Equal(unprocessedBatches[i+1].CreatedAt),
				"batches should be ordered by created_at ASC")
		}
	}
}

func TestListUnprocessed_EmptyResult(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)
	ctx := t.Context()

	// Create test client
	clientID := "test_client_empty"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Empty",
	})
	require.NoError(t, err)

	// Create only completed batches (none processing)
	for i := 0; i < 2; i++ {
		batchID, err := service.Create(ctx, clientID, fmt.Sprintf("completed%d.zip", i), 10)
		require.NoError(t, err)
		err = service.MarkCompleted(ctx, batchID)
		require.NoError(t, err)
	}

	// Test: Get unprocessed batches when none exist
	unprocessedBatches, err := service.ListUnprocessed(ctx)
	require.NoError(t, err)

	// Verify: Should return empty list
	assert.Empty(t, unprocessedBatches)
}

func TestListUnprocessed_DatabaseError(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)

	// Use a canceled context to simulate database error
	ctx, cancel := context.WithCancel(t.Context())
	cancel() // Cancel immediately

	// Test: Try to get unprocessed batches with canceled context
	unprocessedBatches, err := service.ListUnprocessed(ctx)

	// Verify: Should return error
	require.Error(t, err)
	assert.Nil(t, unprocessedBatches)
	assert.Contains(t, err.Error(), "failed to query unprocessed batches")
}

// TestRecordOutcome verifies that RecordOutcome inserts a per-file outcome row
// and that it can be read back via ListOutcomes.
func TestRecordOutcome(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_record_outcome_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Record Outcome",
	})
	require.NoError(t, err)

	batchID, err := service.Create(ctx, clientID, "outcomes.zip", 3)
	require.NoError(t, err)

	// Insert multiple outcomes
	err = service.RecordOutcome(ctx, batchID, "good.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
	require.NoError(t, err)

	errMsg := "file is not a PDF"
	err = service.RecordOutcome(ctx, batchID, "notes.txt", repository.BatchOutcomeStatusInvalidType, &errMsg, nil)
	require.NoError(t, err)

	// Verify persistence via ListOutcomes
	outcomes, err := service.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	require.Len(t, outcomes, 2)

	// Verify the outcome details
	found := map[string]string{}
	for _, o := range outcomes {
		found[o.Filename] = o.Outcome
	}
	assert.Equal(t, "submitted", found["good.pdf"])
	assert.Equal(t, "invalid_type", found["notes.txt"])
}

// TestListOutcomes inserts outcomes for various states and verifies the listing.
func TestListOutcomes(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_list_outcomes_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client List Outcomes",
	})
	require.NoError(t, err)

	batchID, err := service.Create(ctx, clientID, "list.zip", 5)
	require.NoError(t, err)

	// Insert a variety of outcomes
	err = service.RecordOutcome(ctx, batchID, "a.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
	require.NoError(t, err)
	err = service.RecordOutcome(ctx, batchID, "b.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
	require.NoError(t, err)

	errOpen := "cannot open"
	err = service.RecordOutcome(ctx, batchID, "c.pdf", repository.BatchOutcomeStatusFailedOpen, &errOpen, nil)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "d.txt", repository.BatchOutcomeStatusInvalidType, nil, nil)
	require.NoError(t, err)

	dupID := uuid.New()
	// Create a real document for the duplicate reference
	realDocID, createErr := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     "hash_dup",
		BatchID:  &batchID,
	})
	require.NoError(t, createErr)
	_ = dupID // not needed since we use the real ID
	err = service.RecordOutcome(ctx, batchID, "e.pdf", repository.BatchOutcomeStatusDuplicate, nil, &realDocID)
	require.NoError(t, err)

	outcomes, err := service.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	require.Len(t, outcomes, 5)

	// Verify ordering is by created_at ASC
	expectedFilenames := []string{"a.pdf", "b.pdf", "c.pdf", "d.txt", "e.pdf"}
	for i, o := range outcomes {
		assert.Equal(t, expectedFilenames[i], o.Filename)
	}

	// Verify the duplicate has document_id set
	assert.NotNil(t, outcomes[4].DocumentID)
	assert.Equal(t, realDocID, *outcomes[4].DocumentID)
}

// TestListOutcomes_WithCleanFail verifies that the clean_fail column from
// the currentCleanEntries view is joined when outcome is clean_failed.
func TestListOutcomes_WithCleanFail(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_clean_fail_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Clean Fail",
	})
	require.NoError(t, err)

	batchID, err := service.Create(ctx, clientID, "cleanfail.zip", 1)
	require.NoError(t, err)

	// Create a real document
	docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     "hash_clean_fail",
		BatchID:  &batchID,
	})
	require.NoError(t, err)

	// Insert a clean_failed outcome
	err = service.RecordOutcome(ctx, batchID, "corrupt.pdf", repository.BatchOutcomeStatusCleanFailed, nil, &docID)
	require.NoError(t, err)

	// Create a documentClean with fail type and a documentCleanEntry to populate
	// the currentCleanEntries view
	cleanID, err := cfg.GetDBQueries().AddDocumentClean(ctx, &repository.AddDocumentCleanParams{
		Documentid: docID,
		Fail: repository.NullCleanfailtype{
			Cleanfailtype: repository.CleanfailtypeInvalidMimetype,
			Valid:         true,
		},
	})
	require.NoError(t, err)
	err = cfg.GetDBQueries().AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{
		Cleanid: cleanID,
		Version: 1,
	})
	require.NoError(t, err)

	outcomes, err := service.ListOutcomes(ctx, batchID)
	require.NoError(t, err)
	require.Len(t, outcomes, 1)

	assert.Equal(t, "clean_failed", outcomes[0].Outcome)
	require.NotNil(t, outcomes[0].CleanFail)
	assert.Equal(t, "invalid_mimetype", *outcomes[0].CleanFail)
}

func TestSetTotalDocuments(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)
	ctx := t.Context()

	// Create test client
	clientID := "test_client_set_total"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Set Total",
	})
	require.NoError(t, err)

	// Create batch with totalDocs=0
	batchID, err := service.Create(ctx, clientID, "test.zip", 0)
	require.NoError(t, err)

	// Set total documents to 15
	err = service.SetTotalDocuments(ctx, batchID, 15)
	require.NoError(t, err)

	// Fetch and assert TotalDocuments == 15
	batch, err := cfg.GetDBQueries().GetBatchUpload(ctx, &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, int32(15), batch.TotalDocuments)

	// Update progress and verify all fields
	err = service.UpdateProgress(ctx, clientID, batchID, 10, 3, 2)
	require.NoError(t, err)

	batch, err = cfg.GetDBQueries().GetBatchUpload(ctx, &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, int32(10), batch.ProcessedDocuments)
	assert.Equal(t, int32(3), batch.FailedDocuments)
	assert.Equal(t, int32(2), batch.InvalidTypeDocuments)
	assert.Equal(t, int32(100), batch.ProgressPercent) // (10+3+2)/15 * 100 = 100%
}

func TestSetTotalDocuments_ThenUpdateProgress(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)

	service := documentbatch.New(cfg)
	ctx := t.Context()

	// Create test client
	clientID := "test_client_set_total_flow"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Set Total Flow",
	})
	require.NoError(t, err)

	// Create batch with totalDocs=0
	batchID, err := service.Create(ctx, clientID, "flow.zip", 0)
	require.NoError(t, err)

	// Set total to 20
	err = service.SetTotalDocuments(ctx, batchID, 20)
	require.NoError(t, err)

	// Update progress (10 processed, 5 failed, 3 invalid)
	err = service.UpdateProgress(ctx, clientID, batchID, 10, 5, 3)
	require.NoError(t, err)

	// Assert ProgressPercent == 90 (18/20 * 100)
	batch, err := cfg.GetDBQueries().GetBatchUpload(ctx, &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	require.NoError(t, err)
	assert.Equal(t, int32(90), batch.ProgressPercent)
}

// TestCheckDuplicate creates a document, then verifies that CheckDuplicate
// returns its ID when the hash matches.
func TestCheckDuplicate(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_check_dup_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Check Dup",
	})
	require.NoError(t, err)

	// Create a document with a known hash
	knownHash := "sha256_known_hash_abcdef"
	docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     knownHash,
	})
	require.NoError(t, err)

	// Check duplicate with the same hash -- should find it
	existingID, err := service.CheckDuplicate(ctx, clientID, knownHash)
	require.NoError(t, err)
	require.NotNil(t, existingID)
	assert.Equal(t, docID, *existingID)

	// Check duplicate with a different hash -- should return nil
	noMatch, err := service.CheckDuplicate(ctx, clientID, "totally_different_hash")
	require.NoError(t, err)
	assert.Nil(t, noMatch)
}

// TestGet_DerivesSummaryCountsFromOutcomes verifies that the Get endpoint
// derives ProcessedDocuments, FailedDocuments, and InvalidTypeDocuments from
// actual document outcome states rather than the extraction-phase counters
// stored in batch_uploads.
func TestGet_DerivesSummaryCountsFromOutcomes(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_derive_counts_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client Derive Counts",
	})
	require.NoError(t, err)

	// Create batch -- extraction-phase counters will all be zero
	batchID, err := service.Create(ctx, clientID, "mixed.zip", 0)
	require.NoError(t, err)

	// Set total so progress percent can be calculated
	err = service.SetTotalDocuments(ctx, batchID, 7)
	require.NoError(t, err)

	// Create real documents for outcomes that require document_id
	docA, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID, Hash: "hash_a", BatchID: &batchID,
	})
	require.NoError(t, err)
	docB, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID, Hash: "hash_b", BatchID: &batchID,
	})
	require.NoError(t, err)
	docC, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID, Hash: "hash_c", BatchID: &batchID,
	})
	require.NoError(t, err)

	// Record outcomes: 3 processed, 3 failed, 1 invalid_type
	err = service.RecordOutcome(ctx, batchID, "a.pdf",
		repository.BatchOutcomeStatusCleanPassed, nil, &docA)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "b.pdf",
		repository.BatchOutcomeStatusDuplicate, nil, &docB)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "c.pdf",
		repository.BatchOutcomeStatusInitDuplicate, nil, &docC)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "d.pdf",
		repository.BatchOutcomeStatusCleanFailed, nil, nil)
	require.NoError(t, err)

	errOpen := "cannot open"
	err = service.RecordOutcome(ctx, batchID, "e.pdf",
		repository.BatchOutcomeStatusFailedOpen, &errOpen, nil)
	require.NoError(t, err)

	errRead := "truncated"
	err = service.RecordOutcome(ctx, batchID, "f.pdf",
		repository.BatchOutcomeStatusFailedRead, &errRead, nil)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "g.exe",
		repository.BatchOutcomeStatusInvalidType, nil, nil)
	require.NoError(t, err)

	// Call Get and assert derived summary counts
	details, err := service.Get(ctx, clientID, batchID)
	require.NoError(t, err)

	assert.Equal(t, int32(3), details.ProcessedDocuments,
		"processed should count clean_passed + duplicate + init_duplicate")
	assert.Equal(t, int32(3), details.FailedDocuments,
		"failed should count clean_failed + failed_open + failed_read")
	assert.Equal(t, int32(1), details.InvalidTypeDocuments,
		"invalid_type should count invalid_type outcomes")
	assert.Equal(t, int32(100), details.ProgressPercent,
		"progress should be 100%% when all 7 outcomes are terminal")
	assert.Equal(t, int32(7), details.TotalDocuments)
}

// TestGet_DerivedCountsWithInProgressOutcomes verifies that non-terminal
// outcomes (submitted, init_complete, sync_complete) are not counted in
// any summary bucket.
func TestGet_DerivedCountsWithInProgressOutcomes(t *testing.T) {
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	ctx := t.Context()

	service := documentbatch.New(cfg)

	clientID := "test_inprog_counts_" + uuid.New().String()[:8]
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Client InProgress Counts",
	})
	require.NoError(t, err)

	batchID, err := service.Create(ctx, clientID, "partial.zip", 0)
	require.NoError(t, err)
	err = service.SetTotalDocuments(ctx, batchID, 3)
	require.NoError(t, err)

	// 1 terminal, 2 non-terminal
	docA, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID, Hash: "hash_term", BatchID: &batchID,
	})
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "a.pdf",
		repository.BatchOutcomeStatusCleanPassed, nil, &docA)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "b.pdf",
		repository.BatchOutcomeStatusSubmitted, nil, nil)
	require.NoError(t, err)

	err = service.RecordOutcome(ctx, batchID, "c.pdf",
		repository.BatchOutcomeStatusInitComplete, nil, nil)
	require.NoError(t, err)

	details, err := service.Get(ctx, clientID, batchID)
	require.NoError(t, err)

	assert.Equal(t, int32(1), details.ProcessedDocuments,
		"only terminal processed outcomes should be counted")
	assert.Equal(t, int32(0), details.FailedDocuments,
		"no failed outcomes recorded")
	assert.Equal(t, int32(0), details.InvalidTypeDocuments)
	assert.Equal(t, int32(33), details.ProgressPercent,
		"only 1 of 3 outcomes is terminal: 33%%")
}