package prom

import (
	"github.com/pingcap/test-infra/sdk/pkg/prom/model/labels"
	"github.com/pingcap/test-infra/sdk/pkg/prom/parser"
)

// InjectLabel are the simplified version of https://github.com/cherti/promql-labelinjector
func InjectLabel(query string, labelKey, labelValue string) string {
	expr, err := parser.ParseExpr(query)
	if err != nil {
		// fails to parse promql, prometheus server will throw same error
		// we don't handle it here to simplify functionality of modifyQuery
		return query
	}

	parser.Inspect(expr, func(node parser.Node, _ []parser.Node) error {
		// see: https://www.cnblogs.com/charlieroro/p/15261026.html
		// only VectorSelector needs injection
		switch n := node.(type) {
		case *parser.VectorSelector:
			n.LabelMatchers = rewriteLabelset(n.LabelMatchers, labelKey, labelValue)
		}
		return nil
	})

	return expr.String()
}

// rewiteLabelset tasek a set of labelmatchers and ensures the injection target label is
// present in them in the desired form. It returns the labelset where this is ensured.
func rewriteLabelset(labelMatchers []*labels.Matcher, labelKey string, labelValue string) []*labels.Matcher {
	// label is already presented, use that label without modification (to keep things simple)
	for _, matcher := range labelMatchers {
		if matcher.Name == labelKey {
			return labelMatchers
		}
	}

	// label is not presented, inject it
	joblabel, err := labels.NewMatcher(labels.MatchEqual, labelKey, labelValue)
	if err != nil {
		return labelMatchers
	}
	labelMatchers = append(labelMatchers, joblabel)

	return labelMatchers
}