package transformers

import (
	"context"
	"sync"

	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/custom_errors"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)

type PostgresTransformer struct {
	sourceTable config.SourceTableInfo
}

func NewPostgresTransformer(sourceTable config.SourceTableInfo) etl.Transformer {
	return &PostgresTransformer{sourceTable: sourceTable}
}

func (pgTr *PostgresTransformer) Consume(
	ctx context.Context,
	columns []models.ColumnType,
	retryConfig config.RetryConfig,
	batchSize int,
	chBatchesIn <-chan models.Batch,
	chBatchesOut chan<- models.Batch,
	chJobErrorsOut chan<- custom_errors.JobError,
	wgActiveBatches *sync.WaitGroup,
) {
	transformationPlan := computePostgresTransformationPlan(columns)

	acc := &batchAccumulator{batchSize: batchSize}

	for {
		select {
		case <-ctx.Done():
			return

		case batch, ok := <-chBatchesIn:
			if !ok {
				acc.flush(ctx, chBatchesOut, wgActiveBatches)
				return
			}

			if len(transformationPlan) > 0 {
				if err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig); err != nil {
					sendTransformError(ctx, err, chJobErrorsOut)
					return
				}
			}

			if batchSize <= 0 {
				wgActiveBatches.Add(1)
				select {
				case chBatchesOut <- batch:
				case <-ctx.Done():
					wgActiveBatches.Done()
					return
				}
				continue
			}

			acc.add(batch)
			if acc.ready() {
				if !acc.flush(ctx, chBatchesOut, wgActiveBatches) {
					return
				}
			}
		}
	}
}