package transformers

import (
	"context"
	"errors"
	"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/models"
)

func (mssqlTr *MssqlTransformer) Consume(
	ctx context.Context,
	columns []models.ColumnType,
	retryConfig config.RetryConfig,
	chBatchesIn <-chan models.Batch,
	chBatchesOut chan<- models.Batch,
	chJobErrorsOut chan<- custom_errors.JobError,
	wgActiveBatches *sync.WaitGroup,
) {
	transformationPlan := computeTransformationPlan(columns)
	storagePlan := computeStorageTransformationPlan(ctx, mssqlTr.azureClient, mssqlTr.toStorage, columns, mssqlTr.sourceTable)
	transformationPlan = append(transformationPlan, storagePlan...)

	for {
		if ctx.Err() != nil {
			return
		}

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

		case batch, ok := <-chBatchesIn:
			if !ok {
				return
			}

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

			err := ProcessBatchWithRetries(ctx, &batch, transformationPlan, retryConfig)
			if err != nil {
				if errors.Is(err, ctx.Err()) {
					return
				}

				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
					select {
					case chJobErrorsOut <- *jobError:
					case <-ctx.Done():
						return
					}
				} else {
					select {
					case chJobErrorsOut <- custom_errors.JobError{ShouldCancelJob: true, Msg: "Transformation failed", Prev: err}:
					case <-ctx.Done():
						return
					}
				}

				return
			}

			select {
			case chBatchesOut <- batch:
			case <-ctx.Done():
				return
			}

			wgActiveBatches.Add(1)
		}
	}
}