package extractors

import (
	"context"
	"errors"
	"slices"
	"strings"
	"sync"
	"sync/atomic"

	"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"
	"github.com/sirupsen/logrus"
)

func (ex *GenericExtractor) Consume(
	ctx context.Context,
	tableInfo config.SourceTableInfo,
	columns []models.ColumnType,
	batchSize int,
	retryConfig config.RetryConfig,
	chPartitionsIn <-chan models.Partition,
	chBatchesOut chan<- models.Batch,
	chErrorsOut chan<- custom_errors.JobError,
	wgActivePartitions *sync.WaitGroup,
	rowsRead *int64,
	fromJsonColumns []config.FromJsonItem,
) {
	indexPrimaryKey := slices.IndexFunc(columns, func(col models.ColumnType) bool {
		return strings.EqualFold(col.Name(), tableInfo.PrimaryKey)
	})

	if indexPrimaryKey == -1 {
		select {
		case <-ctx.Done():
			return
		case chErrorsOut <- custom_errors.JobError{
			ShouldCancelJob: true,
			Msg:             "Primary key not found in provided columns",
		}:
		}

		return
	}

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

		select {
		case <-ctx.Done():
			return
		case partition, ok := <-chPartitionsIn:
			if !ok {
				return
			}

			rowsReadResult, err := ex.ProcessPartitionWithRetries(
				ctx,
				tableInfo,
				columns,
				batchSize,
				partition,
				indexPrimaryKey,
				retryConfig,
				chBatchesOut,
				fromJsonColumns,
			)
			wgActivePartitions.Done()

			if rowsReadResult > 0 {
				current := atomic.LoadInt64(rowsRead)
				logrus.Debugf("Rows read: +%v [current=%v] (%s.%s)", rowsReadResult, current, tableInfo.Schema, tableInfo.Table)
				atomic.AddInt64(rowsRead, int64(rowsReadResult))
			}

			if err != nil {
				if jobError, ok := errors.AsType[*custom_errors.JobError](err); ok {
					select {
					case <-ctx.Done():
						return
					case chErrorsOut <- *jobError:
					}
				} else {
					select {
					case <-ctx.Done():
						return
					case chErrorsOut <- custom_errors.JobError{ShouldCancelJob: false, Msg: err.Error(), Prev: err}:
					}
				}
			}
		}
	}
}