// SPDX-License-Identifier: AGPL-3.0-only

package rangevectorsplitting

import (
	"github.com/go-kit/log"
	"github.com/prometheus/prometheus/model/histogram"
	"github.com/prometheus/prometheus/promql/parser/posrange"
	"github.com/prometheus/prometheus/util/annotations"

	"github.com/grafana/mimir/pkg/streamingpromql/operators/functions"
	"github.com/grafana/mimir/pkg/streamingpromql/optimize/plan/rangevectorsplitting/cache"
	"github.com/grafana/mimir/pkg/streamingpromql/planning"
	"github.com/grafana/mimir/pkg/streamingpromql/types"
	"github.com/grafana/mimir/pkg/util/limiter"
)

// Range represents a time range within a query split.
// Start is exclusive, End is inclusive.
type Range struct {
	Start     int64
	End       int64
	Cacheable bool
}

// SplitGenerateFunc generates an intermediate result for a single split range.
// hasValue indicates whether the step produced any data for the current series.
// If false, the series is excluded from the cached results for this range in order to reduce cache entry size
type SplitGenerateFunc[T any] func(
	step *types.RangeVectorStepData,
	emitAnnotation types.EmitAnnotationFunc,
	memoryConsumptionTracker *limiter.MemoryConsumptionTracker,
) (t T, hasValue bool, err error)

// SplitCombineFunc combines intermediate results from multiple split ranges.
// Histograms within the input ranges must not be modified in place. These histograms can share references with
// histogram protos that are waiting to be cached.
type SplitCombineFunc[T any] func(
	ranges []T,
	rangeStart int64,
	rangeEnd int64,
	emitAnnotation types.EmitAnnotationFunc,
	memoryConsumptionTracker *limiter.MemoryConsumptionTracker,
) (f float64, hasFloat bool, h *histogram.FloatHistogram, err error)

type SplitOperatorFactory func(
	innerNode planning.Node,
	materializer *planning.Materializer,
	timeRange types.QueryTimeRange,
	ranges []Range,
	cacheKey string,
	irCache *cache.CacheFactory,
	expressionPosition posrange.PositionRange,
	annotations *annotations.Annotations,
	memoryConsumptionTracker *limiter.MemoryConsumptionTracker,
	enableDelayedNameRemoval bool,
	logger log.Logger,
) (types.Operator, error)

func NewSplitOperatorFactory[T any](
	generate SplitGenerateFunc[T],
	combine SplitCombineFunc[T],
	codec cache.SplitCodec[T],
	funcDef functions.FunctionOverRangeVectorDefinition,
	funcId functions.Function,
) SplitOperatorFactory {
	return func(
		innerNode planning.Node,
		materializer *planning.Materializer,
		timeRange types.QueryTimeRange,
		ranges []Range,
		cacheKey string,
		irCache *cache.CacheFactory,
		expressionPosition posrange.PositionRange,
		annotations *annotations.Annotations,
		memoryConsumptionTracker *limiter.MemoryConsumptionTracker,
		enableDelayedNameRemoval bool,
		logger log.Logger,
	) (types.Operator, error) {
		return NewSplittingFunctionOverRangeVector[T](
			innerNode,
			materializer,
			timeRange,
			ranges,
			cacheKey,
			irCache,
			funcId,
			funcDef,
			generate,
			combine,
			codec,
			expressionPosition,
			annotations,
			memoryConsumptionTracker,
			enableDelayedNameRemoval,
			logger,
		)
	}
}
