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

package batch

import (
	"context"
	"encoding/json"
	"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"`
}

// BatchUploadDetails includes full details with failed filenames
type BatchUploadDetails struct {
	BatchUploadSummary
	FailedFilenames []string `json:"failed_filenames"`
	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
func (s *Service) Get(ctx context.Context, clientID string, batchID uuid.UUID) (*BatchUploadDetails, error) {
	batch, err := s.cfg.GetDBQueries().GetBatchUploadWithStorage(ctx, &repository.GetBatchUploadWithStorageParams{
		ID:       batchID,
		ClientID: clientID,
	})
	if err != nil {
		return nil, err
	}

	return s.convertToDetailsWithStorage(batch), 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,
	})
}

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

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