// Copyright 2025 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package index

import (
	"context"
	"slices"

	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
)

// LookupPlanner plans how to execute index lookups by deciding which matchers
// to apply during index lookup versus after series retrieval.
type LookupPlanner interface {
	PlanIndexLookup(ctx context.Context, plan LookupPlan, hints *storage.SelectHints) (LookupPlan, error)
}

// ChainLookupPlanners is useful to break up the planning logic into multiple independent planners and connect them in sequence.
// The returned LookupPlanner calls the planners in sequence and feeds the output of one planner as the input of the next.
func ChainLookupPlanners(planners ...LookupPlanner) LookupPlanner {
	return chainedLookupPlanner{planners: planners}
}

type chainedLookupPlanner struct {
	planners []LookupPlanner
}

func (c chainedLookupPlanner) PlanIndexLookup(ctx context.Context, plan LookupPlan, hints *storage.SelectHints) (LookupPlan, error) {
	var err error
	for _, p := range c.planners {
		plan, err = p.PlanIndexLookup(ctx, plan, hints)
		if err != nil {
			return nil, err
		}
	}
	return plan, nil
}

// LookupPlan represents the decision of which matchers to apply during
// index lookup versus during series scanning.
type LookupPlan interface {
	// ScanMatchers returns matchers that should be applied during series scanning
	ScanMatchers() []*labels.Matcher
	// IndexMatchers returns matchers that should be applied during index lookup
	IndexMatchers() []*labels.Matcher
}

// NewIndexOnlyLookupPlan creates a new LookupPlan which uses the provided matchers for looking up the index only.
func NewIndexOnlyLookupPlan(matchers []*labels.Matcher) LookupPlan {
	return indexOnlyLookupPlan(matchers)
}

type indexOnlyLookupPlan []*labels.Matcher

func (indexOnlyLookupPlan) ScanMatchers() []*labels.Matcher {
	return nil
}

func (l indexOnlyLookupPlan) IndexMatchers() []*labels.Matcher {
	return l
}

// ScanEmptyMatchersLookupPlanner implements LookupPlanner by deferring empty matchers such as l="", l=~".+" to scan matchers.
type ScanEmptyMatchersLookupPlanner struct{}

func (p *ScanEmptyMatchersLookupPlanner) PlanIndexLookup(_ context.Context, plan LookupPlan, _ *storage.SelectHints) (LookupPlan, error) {
	indexMatchers := plan.IndexMatchers()
	if len(indexMatchers) <= 1 || !p.canOptimizeIndexMatchers(indexMatchers) {
		// If there is only one matcher, then using the index is usually more efficient.
		// This also covers test cases which use matchers such as {""=""} or {""=~".*"} to mean "match all series"
		// Also avoid allocating a plan if we're not going to change any of the index matchers.
		return plan, nil
	}

	// Clone so that we don't overwrite the matchers of clients when we resize the slice
	indexMatchers = slices.Clone(indexMatchers)
	scanMatchers := slices.Clone(plan.ScanMatchers())

	for i := 0; i < len(indexMatchers); i++ {
		matcher := indexMatchers[i]
		if p.shouldDropMatcher(matcher) {
			indexMatchers = append(indexMatchers[:i], indexMatchers[i+1:]...)
			i--
		} else if p.shouldScanMatcher(matcher) {
			indexMatchers = append(indexMatchers[:i], indexMatchers[i+1:]...)
			i--
			scanMatchers = append(scanMatchers, matcher)
		}
	}

	if len(indexMatchers) == 0 && len(scanMatchers) > 0 {
		// Zero index matchers match no series. We retain one index matcher so that we have a base set of series which we can scan.
		indexMatchers = scanMatchers[:1]
		scanMatchers = scanMatchers[1:]
	}

	if len(scanMatchers) == 0 {
		// Avoid an allocation if the plan is simple.
		return NewIndexOnlyLookupPlan(indexMatchers), nil
	}

	return &concreteLookupPlan{
		indexMatchers: indexMatchers,
		scanMatchers:  scanMatchers,
	}, nil
}

func (*ScanEmptyMatchersLookupPlanner) shouldDropMatcher(matcher *labels.Matcher) bool {
	// This matches everything (empty and arbitrary values), so it doesn't reduce the selectivity of the whole set of matchers.
	return matcher.Type == labels.MatchRegexp && matcher.Value == ".*"
}

func (*ScanEmptyMatchersLookupPlanner) shouldScanMatcher(matcher *labels.Matcher) bool {
	// Put empty string matchers and regex matchers that match everything in scan matchers.
	// These matchers are unlikely to reduce the selectivity of the matchers, but can still filter out some series, so we can't ignore them.
	return (matcher.Type == labels.MatchEqual && matcher.Value == "") ||
		(matcher.Type == labels.MatchRegexp && matcher.Value == ".+")
}

func (p *ScanEmptyMatchersLookupPlanner) canOptimizeIndexMatchers(matchers []*labels.Matcher) bool {
	for _, matcher := range matchers {
		if p.shouldScanMatcher(matcher) || p.shouldDropMatcher(matcher) {
			return true
		}
	}
	return false
}

// concreteLookupPlan implements LookupPlan by storing pre-computed index and scan matchers.
type concreteLookupPlan struct {
	indexMatchers []*labels.Matcher
	scanMatchers  []*labels.Matcher
}

func (p *concreteLookupPlan) ScanMatchers() []*labels.Matcher {
	return p.scanMatchers
}

func (p *concreteLookupPlan) IndexMatchers() []*labels.Matcher {
	return p.indexMatchers
}
