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

package batch

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgtype"
)

// Batch status constants
const (
	StatusProcessing = "processing"
	StatusCompleted  = "completed"
	StatusFailed     = "failed"
	StatusCanceled   = "canceled"
)

// BatchUploadSummary represents a summary of a batch upload
type BatchUploadSummary struct {
	ID                   uuid.UUID  `json:"batch_id"`
	ClientID             string     `json:"client_id"`
	OriginalFilename     string     `json:"original_filename"`
	Status               string     `json:"status"`
	TotalDocuments       int32      `json:"total_documents"`
	ProcessedDocuments   int32      `json:"processed_documents"`
	FailedDocuments      int32      `json:"failed_documents"`
	InvalidTypeDocuments int32      `json:"invalid_type_documents"`
	ProgressPercent      int32      `json:"progress_percent"`
	CreatedAt            time.Time  `json:"created_at"`
	CompletedAt          *time.Time `json:"completed_at,omitempty"`
}

// BatchDocumentOutcome represents the per-file result from batch extraction
// and pipeline processing. This is the single source of truth for a file's
// status within a batch.
//
// Fields:
//   - ID: unique outcome record ID
//   - BatchID: the batch this file belongs to
//   - Filename: original filename from the ZIP
//   - Outcome: current status (extraction outcome or pipeline stage)
//   - ErrorDetail: human-readable error/failure message, nil for success states
//   - DocumentID: resolved document ID (nil if document not yet created)
//   - CleanFail: specific clean failure reason from documentCleans table (nil if N/A)
//   - CreatedAt: when the outcome row was created (extraction time)
//   - UpdatedAt: when the outcome was last changed (pipeline progression)
type BatchDocumentOutcome struct {
	ID          uuid.UUID
	BatchID     uuid.UUID
	Filename    string
	Outcome     string
	ErrorDetail *string
	DocumentID  *uuid.UUID
	CleanFail   *string
	CreatedAt   time.Time
	UpdatedAt   time.Time
}

// BatchUploadDetails includes full details with failed filenames and document outcomes
type BatchUploadDetails struct {
	BatchUploadSummary
	FailedFilenames  []string                `json:"failed_filenames"`
	DocumentOutcomes []*BatchDocumentOutcome `json:"document_outcomes,omitempty"`
	ArchiveBucket    string                  `json:"archive_bucket,omitempty"`
	ArchiveKey       string                  `json:"archive_key,omitempty"`
	FileSizeBytes    int64                   `json:"file_size_bytes,omitempty"`
}

// Service handles batch upload operations
type Service struct {
	cfg serviceconfig.ConfigProvider
}

// New creates a new batch service
func New(cfg serviceconfig.ConfigProvider) *Service {
	return &Service{
		cfg: cfg,
	}
}

// Create creates a new batch upload record
func (s *Service) Create(ctx context.Context, clientID string, filename string, totalDocs int32) (uuid.UUID, error) {
	return s.cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
		ClientID:         clientID,
		OriginalFilename: filename,
		TotalDocuments:   totalDocs,
	})
}

// CreateWithStorage creates a batch record with S3 storage metadata
func (s *Service) CreateWithStorage(ctx context.Context, clientID string, filename string, totalDocs int32, bucket string, key string, sizeBytes int64) (uuid.UUID, error) {
	result, err := s.cfg.GetDBQueries().CreateBatchUploadWithStorage(ctx, &repository.CreateBatchUploadWithStorageParams{
		ClientID:         clientID,
		OriginalFilename: filename,
		TotalDocuments:   totalDocs,
		Column4:          repository.BatchStatusProcessing,
		ArchiveBucket:    &bucket,
		ArchiveKey:       &key,
		FileSizeBytes:    &sizeBytes,
	})
	if err != nil {
		return uuid.Nil, fmt.Errorf("failed to create batch upload record: %w", err)
	}

	return result.ID, nil
}

// Get retrieves batch upload details including per-file document outcomes.
//
// Parameters:
//   - ctx: request context
//   - clientID: the client identifier
//   - batchID: the batch UUID
//
// Returns:
//   - *BatchUploadDetails: full batch details with document outcomes
//   - error: database error, or nil
func (s *Service) Get(ctx context.Context, clientID string, batchID uuid.UUID) (*BatchUploadDetails, error) {
	batchRow, err := s.cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{
		ID:       batchID,
		ClientID: clientID,
	})
	if err != nil {
		return nil, err
	}

	details := s.convertToDetailsWithStorage(batchRow)

	// Populate document outcomes
	outcomes, err := s.ListOutcomes(ctx, batchID)
	if err != nil {
		return nil, fmt.Errorf("failed to list batch outcomes: %w", err)
	}
	details.DocumentOutcomes = outcomes

	// Derive summary counts from actual outcome states rather than
	// the extraction-phase counters stored in batch_uploads.
	s.deriveSummaryCounts(details)

	return details, nil
}

// List retrieves all batch uploads for a client with pagination
func (s *Service) List(ctx context.Context, clientID string, limit int32, offset int32) ([]*BatchUploadSummary, error) {
	if limit <= 0 {
		limit = 20 // default limit
	}
	if limit > 100 {
		limit = 100 // max limit
	}

	batches, err := s.cfg.GetDBQueries().ListBatchUploads(ctx, &repository.ListBatchUploadsParams{
		ClientID: clientID,
		Limit:    int64(limit),
		Offset:   int64(offset),
	})
	if err != nil {
		return nil, err
	}

	summaries := make([]*BatchUploadSummary, len(batches))
	for i, batch := range batches {
		summaries[i] = s.convertToSummary(batch)
	}

	return summaries, nil
}

// UpdateProgress updates batch processing progress
func (s *Service) UpdateProgress(ctx context.Context, clientID string, batchID uuid.UUID, processed, failed, invalidType int32) error {
	// Calculate progress percentage
	batch, err := s.cfg.GetDBQueries().GetBatchUpload(ctx, &repository.GetBatchUploadParams{
		ID:       batchID,
		ClientID: clientID,
	})
	if err != nil {
		return fmt.Errorf("failed to get batch for progress update: %w", err)
	}

	total := batch.TotalDocuments
	if total == 0 {
		return fmt.Errorf("batch has zero total documents")
	}

	progressPercent := int32(float64(processed+failed+invalidType) / float64(total) * 100)

	return s.cfg.GetDBQueries().UpdateBatchProgress(ctx, &repository.UpdateBatchProgressParams{
		ID:                   batchID,
		ProcessedDocuments:   processed,
		FailedDocuments:      failed,
		InvalidTypeDocuments: invalidType,
		ProgressPercent:      progressPercent,
	})
}

// SetTotalDocuments updates the total document count for a batch.
// Called after ZIP extraction when the actual file count is known.
//
// Parameters:
//   - ctx: request context
//   - batchID: the batch UUID
//   - total: the total number of files extracted from the ZIP
//
// Returns:
//   - error: database error, or nil
func (s *Service) SetTotalDocuments(ctx context.Context, batchID uuid.UUID, total int32) error {
	return s.cfg.GetDBQueries().SetBatchTotalDocuments(ctx, &repository.SetBatchTotalDocumentsParams{
		ID:             batchID,
		TotalDocuments: total,
	})
}

// UpdateStatus updates the status of a batch
func (s *Service) UpdateStatus(ctx context.Context, batchID uuid.UUID, status string) error {
	// Convert string status to repository.BatchStatus
	var batchStatus repository.BatchStatus
	switch status {
	case StatusProcessing:
		batchStatus = repository.BatchStatusProcessing
	case StatusCompleted:
		batchStatus = repository.BatchStatusCompleted
	case StatusFailed:
		batchStatus = repository.BatchStatusFailed
	case StatusCanceled:
		batchStatus = repository.BatchStatusCancelled
	default:
		return fmt.Errorf("invalid status: %s", status)
	}

	return s.cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{
		ID:      batchID,
		Column2: batchStatus,
	})
}

// MarkCompleted marks a batch as completed
func (s *Service) MarkCompleted(ctx context.Context, batchID uuid.UUID) error {
	return s.UpdateStatus(ctx, batchID, StatusCompleted)
}

// MarkFailed marks a batch as failed
func (s *Service) MarkFailed(ctx context.Context, batchID uuid.UUID) error {
	return s.UpdateStatus(ctx, batchID, StatusFailed)
}

// AddFailedFilename adds a failed filename to the batch
func (s *Service) AddFailedFilename(ctx context.Context, batchID uuid.UUID, filename string) error {
	return s.cfg.GetDBQueries().AddFailedFilename(ctx, &repository.AddFailedFilenameParams{
		ID:      batchID,
		Column2: filename,
	})
}

// convertToSummary converts repository list row to service summary model
func (s *Service) convertToSummary(batch *repository.ListBatchUploadsRow) *BatchUploadSummary {
	summary := &BatchUploadSummary{
		ID:                   batch.ID,
		ClientID:             batch.ClientID,
		OriginalFilename:     batch.OriginalFilename,
		Status:               string(batch.Status),
		TotalDocuments:       batch.TotalDocuments,
		ProcessedDocuments:   batch.ProcessedDocuments,
		FailedDocuments:      batch.FailedDocuments,
		InvalidTypeDocuments: batch.InvalidTypeDocuments,
		ProgressPercent:      batch.ProgressPercent,
		CreatedAt:            s.pgTimestampToTime(batch.CreatedAt),
	}

	if batch.CompletedAt.Valid {
		t := s.pgTimestampToTime(batch.CompletedAt)
		summary.CompletedAt = &t
	}

	return summary
}

// pgTimestampToTime converts pgtype.Timestamp to time.Time
func (s *Service) pgTimestampToTime(ts pgtype.Timestamp) time.Time {
	if ts.Valid {
		return ts.Time
	}
	return time.Time{}
}

// convertToDetailsWithStorage converts repository model with storage to service details model
func (s *Service) convertToDetailsWithStorage(batch *repository.GetBatchUploadWithStorageRow) *BatchUploadDetails {
	details := &BatchUploadDetails{
		BatchUploadSummary: BatchUploadSummary{
			ID:                   batch.ID,
			ClientID:             batch.ClientID,
			OriginalFilename:     batch.OriginalFilename,
			Status:               string(batch.Status),
			TotalDocuments:       batch.TotalDocuments,
			ProcessedDocuments:   batch.ProcessedDocuments,
			FailedDocuments:      batch.FailedDocuments,
			InvalidTypeDocuments: batch.InvalidTypeDocuments,
			ProgressPercent:      0, // Calculate if needed
			CreatedAt:            batch.CreatedAt.Time,
		},
		FailedFilenames: make([]string, 0),
	}

	if batch.CompletedAt.Valid {
		details.CompletedAt = &batch.CompletedAt.Time
	}

	// Add storage fields
	if batch.ArchiveBucket != nil {
		details.ArchiveBucket = *batch.ArchiveBucket
	}
	if batch.ArchiveKey != nil {
		details.ArchiveKey = *batch.ArchiveKey
	}
	if batch.FileSizeBytes != nil {
		details.FileSizeBytes = *batch.FileSizeBytes
	}

	// Parse failed filenames from JSONB
	if len(batch.FailedFilenames) > 0 {
		var filenames []string
		if err := json.Unmarshal(batch.FailedFilenames, &filenames); err == nil {
			details.FailedFilenames = filenames
		}
	}

	// Calculate progress percent
	if batch.TotalDocuments > 0 {
		processed := batch.ProcessedDocuments + batch.FailedDocuments + batch.InvalidTypeDocuments
		details.ProgressPercent = int32(float64(processed) / float64(batch.TotalDocuments) * 100)
	}

	return details
}

// RecordOutcome inserts or upserts a per-file outcome row for a batch document.
// Called by the batch worker during ZIP extraction to record the initial outcome
// for each file.
//
// Parameters:
//   - ctx: request context
//   - batchID: the batch this file belongs to
//   - filename: original filename from the ZIP
//   - outcome: the typed outcome status (e.g. repository.BatchOutcomeStatusSubmitted)
//   - errorDetail: optional error message for failure outcomes
//   - documentID: optional document ID (set for duplicates)
//
// Returns:
//   - error: database error, or nil
func (s *Service) RecordOutcome(ctx context.Context, batchID uuid.UUID, filename string,
	outcome repository.BatchOutcomeStatus, errorDetail *string, documentID *uuid.UUID) error {
	return s.cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{
		BatchID:     batchID,
		Filename:    filename,
		Column3:     outcome,
		ErrorDetail: errorDetail,
		DocumentID:  documentID,
	})
}

// ListOutcomes retrieves all per-file outcome rows for a batch, with clean
// failure detail joined from the currentCleanEntries view.
//
// Parameters:
//   - ctx: request context
//   - batchID: the batch UUID
//
// Returns:
//   - []*BatchDocumentOutcome: list of outcome rows ordered by creation time
//   - error: database error, or nil
func (s *Service) ListOutcomes(ctx context.Context, batchID uuid.UUID) ([]*BatchDocumentOutcome, error) {
	rows, err := s.cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID)
	if err != nil {
		return nil, err
	}

	outcomes := make([]*BatchDocumentOutcome, len(rows))
	for i, row := range rows {
		o := &BatchDocumentOutcome{
			ID:          row.ID,
			BatchID:     row.BatchID,
			Filename:    row.Filename,
			Outcome:     string(row.Outcome),
			ErrorDetail: row.ErrorDetail,
			DocumentID:  row.DocumentID,
			CreatedAt:   s.pgTimestampToTime(row.CreatedAt),
			UpdatedAt:   s.pgTimestampToTime(row.UpdatedAt),
		}
		if row.CleanFail.Valid {
			failStr := string(row.CleanFail.Cleanfailtype)
			o.CleanFail = &failStr
		}
		outcomes[i] = o
	}

	return outcomes, nil
}

// CheckDuplicate checks if a document with the given hash already exists for
// the specified client. Returns the existing document's UUID if found, nil otherwise.
//
// Parameters:
//   - ctx: request context
//   - clientID: the client identifier
//   - hash: the SHA-256 hash of the file content
//
// Returns:
//   - *uuid.UUID: existing document ID if a duplicate exists, nil otherwise
//   - error: database error (other than not-found), or nil
func (s *Service) CheckDuplicate(ctx context.Context, clientID string, hash string) (*uuid.UUID, error) {
	id, err := s.cfg.GetDBQueries().GetDocumentIDByHash(ctx, &repository.GetDocumentIDByHashParams{
		Clientid: clientID,
		Hash:     hash,
	})
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	return &id, nil
}

// deriveSummaryCounts recomputes ProcessedDocuments, FailedDocuments, and
// InvalidTypeDocuments from the actual document outcome states. This ensures the
// summary reflects final pipeline results rather than extraction-phase counters.
//
// Terminal outcome bucketing:
//
//	processed:    clean_passed, sync_skipped, duplicate, init_duplicate
//	failed:       clean_failed, failed_open, failed_read, failed_s3_upload, failed_upload
//	invalid_type: invalid_type
//	in-progress:  submitted, init_complete, sync_complete (not expected at completion)
func (s *Service) deriveSummaryCounts(details *BatchUploadDetails) {
	if len(details.DocumentOutcomes) == 0 {
		return
	}

	var processed, failed, invalidType int32
	for _, o := range details.DocumentOutcomes {
		switch o.Outcome {
		case "clean_passed", "sync_skipped", "duplicate", "init_duplicate":
			processed++
		case "clean_failed", "failed_open", "failed_read",
			"failed_s3_upload", "failed_upload":
			failed++
		case "invalid_type":
			invalidType++
		default:
			// Non-terminal states (submitted, init_complete, sync_complete)
			// are in-progress; don't count them in any bucket.
		}
	}

	details.ProcessedDocuments = processed
	details.FailedDocuments = failed
	details.InvalidTypeDocuments = invalidType

	// Recalculate progress percent from derived counts
	total := details.TotalDocuments
	if total > 0 {
		done := processed + failed + invalidType
		details.ProgressPercent = int32(float64(done) / float64(total) * 100)
	}
}

// ListUnprocessed retrieves all unprocessed batch uploads across all clients
func (s *Service) ListUnprocessed(ctx context.Context) ([]*BatchUploadSummary, error) {
	batches, err := s.cfg.GetDBQueries().GetUnprocessedBatches(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to query unprocessed batches: %w", err)
	}

	summaries := make([]*BatchUploadSummary, len(batches))
	for i, batch := range batches {
		summaries[i] = &BatchUploadSummary{
			ID:                   batch.ID,
			ClientID:             batch.ClientID,
			OriginalFilename:     batch.OriginalFilename,
			Status:               string(batch.Status),
			TotalDocuments:       batch.TotalDocuments,
			ProcessedDocuments:   batch.ProcessedDocuments,
			FailedDocuments:      batch.FailedDocuments,
			InvalidTypeDocuments: batch.InvalidTypeDocuments,
			ProgressPercent:      batch.ProgressPercent,
			CreatedAt:            s.pgTimestampToTime(batch.CreatedAt),
		}

		if batch.CompletedAt.Valid {
			t := s.pgTimestampToTime(batch.CompletedAt)
			summaries[i].CompletedAt = &t
		}
	}

	return summaries, nil
}