// Copyright 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 inhibit

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/prometheus/common/model"
	"github.com/prometheus/common/promslog"
	"github.com/stretchr/testify/require"

	"github.com/prometheus/alertmanager/alert"
	amcommoncfg "github.com/prometheus/alertmanager/config/common"
	"github.com/prometheus/alertmanager/eventrecorder"
	"github.com/prometheus/alertmanager/marker"
	"github.com/prometheus/alertmanager/pkg/labels"
	"github.com/prometheus/alertmanager/provider"
	"github.com/prometheus/alertmanager/store"
)

var nopLogger = promslog.NewNopLogger()

// checkMutes calls ih.Mutes with a fresh AlertMarker in the context,
// asserts the mute result matches wantMuted, and verifies the marker status.
func checkMutes(t *testing.T, ih *Inhibitor, target model.LabelSet, wantMuted bool, msgAndArgs ...any) {
	t.Helper()
	m := marker.NewAlertMarker()
	ctx := marker.WithContext(context.Background(), m)
	got := ih.Mutes(ctx, target)
	require.Equal(t, wantMuted, got, msgAndArgs...)
	fp := target.Fingerprint()
	status := m.Status(fp)
	if wantMuted {
		require.Equal(t, alert.AlertStateSuppressed, status.State, msgAndArgs...)
		require.NotEmpty(t, status.InhibitedBy, msgAndArgs...)
	} else {
		require.Equal(t, alert.AlertStateActive, status.State, msgAndArgs...)
	}
}

func TestInhibitRuleHasEqual(t *testing.T) {
	t.Parallel()

	now := time.Now()
	cases := []struct {
		name    string
		initial map[model.Fingerprint]*alert.Alert
		equal   model.LabelNames
		input   model.LabelSet
		result  bool
	}{
		{
			name:    "no source alerts",
			initial: map[model.Fingerprint]*alert.Alert{},
			input:   model.LabelSet{"a": "b"},
			result:  false,
		},
		{
			name:    "no equal labels, any source alerts satisfies the requirement",
			initial: map[model.Fingerprint]*alert.Alert{1: {}},
			input:   model.LabelSet{"a": "b"},
			result:  true,
		},
		{
			name: "matching but already resolved",
			initial: map[model.Fingerprint]*alert.Alert{
				1: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "b", "b": "f"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(-time.Second),
					},
				},
				2: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "b", "b": "c"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(-time.Second),
					},
				},
			},
			equal:  model.LabelNames{"a", "b"},
			input:  model.LabelSet{"a": "b", "b": "c"},
			result: false,
		},
		{
			name: "matching and unresolved",
			initial: map[model.Fingerprint]*alert.Alert{
				1: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "b", "c": "d"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(-time.Second),
					},
				},
				2: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "b", "c": "f"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(time.Hour),
					},
				},
			},
			equal:  model.LabelNames{"a"},
			input:  model.LabelSet{"a": "b"},
			result: true,
		},
		{
			name: "equal label does not match",
			initial: map[model.Fingerprint]*alert.Alert{
				1: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "c", "c": "d"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(-time.Second),
					},
				},
				2: {
					Alert: model.Alert{
						Labels:   model.LabelSet{"a": "c", "c": "f"},
						StartsAt: now.Add(-time.Minute),
						EndsAt:   now.Add(-time.Second),
					},
				},
			},
			equal:  model.LabelNames{"a"},
			input:  model.LabelSet{"a": "b"},
			result: false,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			r := &InhibitRule{
				Equal:  map[model.LabelName]struct{}{},
				scache: store.NewAlerts(),
				sindex: newIndex(),
			}
			for _, ln := range c.equal {
				r.Equal[ln] = struct{}{}
			}
			for _, v := range c.initial {
				r.scache.Set(v)
				r.updateIndex(v)
			}

			if _, have := r.hasEqual(c.input, false, time.Now()); have != c.result {
				t.Errorf("Unexpected result %t, expected %t", have, c.result)
			}
		})
	}
}

func TestInhibitRuleMatches(t *testing.T) {
	t.Parallel()

	rule1 := amcommoncfg.InhibitRule{
		SourceMatch: map[string]string{"s1": "1"},
		TargetMatch: map[string]string{"t1": "1"},
		Equal:       []string{"e"},
	}
	rule2 := amcommoncfg.InhibitRule{
		SourceMatch: map[string]string{"s2": "1"},
		TargetMatch: map[string]string{"t2": "1"},
		Equal:       []string{"e"},
	}

	ih := NewInhibitor(nil, []amcommoncfg.InhibitRule{rule1, rule2}, nopLogger, eventrecorder.NopRecorder())
	now := time.Now()
	// Active alert that matches the source filter of rule1.
	sourceAlert1 := &alert.Alert{
		Alert: model.Alert{
			Labels:   model.LabelSet{"s1": "1", "t1": "2", "e": "1"},
			StartsAt: now.Add(-time.Minute),
			EndsAt:   now.Add(time.Hour),
		},
	}
	// Active alert that matches the source filter _and_ the target filter of rule2.
	sourceAlert2 := &alert.Alert{
		Alert: model.Alert{
			Labels:   model.LabelSet{"s2": "1", "t2": "1", "e": "1"},
			StartsAt: now.Add(-time.Minute),
			EndsAt:   now.Add(time.Hour),
		},
	}

	ih.rules[0].scache = store.NewAlerts()
	ih.rules[0].scache.Set(sourceAlert1)
	ih.rules[0].sindex = newIndex()
	ih.rules[0].updateIndex(sourceAlert1)

	ih.rules[1].scache = store.NewAlerts()
	ih.rules[1].scache.Set(sourceAlert2)
	ih.rules[1].sindex = newIndex()
	ih.rules[1].updateIndex(sourceAlert2)

	cases := []struct {
		target   model.LabelSet
		expected bool
	}{
		{
			// Matches target filter of rule1, inhibited.
			target:   model.LabelSet{"t1": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches target filter of rule2, inhibited.
			target:   model.LabelSet{"t2": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches target filter of rule1 (plus noise), inhibited.
			target:   model.LabelSet{"t1": "1", "t3": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches target filter of rule1 plus rule2, inhibited.
			target:   model.LabelSet{"t1": "1", "t2": "1", "e": "1"},
			expected: true,
		},
		{
			// Doesn't match target filter, not inhibited.
			target:   model.LabelSet{"t1": "0", "e": "1"},
			expected: false,
		},
		{
			// Matches both source and target filters of rule1,
			// inhibited because sourceAlert1 matches only the
			// source filter of rule1.
			target:   model.LabelSet{"s1": "1", "t1": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches both source and target filters of rule2,
			// not inhibited because sourceAlert2 matches also both the
			// source and target filter of rule2.
			target:   model.LabelSet{"s2": "1", "t2": "1", "e": "1"},
			expected: false,
		},
		{
			// Matches target filter, equal label doesn't match, not inhibited
			target:   model.LabelSet{"t1": "1", "e": "0"},
			expected: false,
		},
	}

	for _, c := range cases {
		checkMutes(t, ih, c.target, c.expected, "target %v", c.target)
	}
}

func TestInhibitRuleMatchers(t *testing.T) {
	t.Parallel()

	rule1 := amcommoncfg.InhibitRule{
		SourceMatchers: amcommoncfg.Matchers{&labels.Matcher{Type: labels.MatchEqual, Name: "s1", Value: "1"}},
		TargetMatchers: amcommoncfg.Matchers{&labels.Matcher{Type: labels.MatchNotEqual, Name: "t1", Value: "1"}},
		Equal:          []string{"e"},
	}
	rule2 := amcommoncfg.InhibitRule{
		SourceMatchers: amcommoncfg.Matchers{&labels.Matcher{Type: labels.MatchEqual, Name: "s2", Value: "1"}},
		TargetMatchers: amcommoncfg.Matchers{&labels.Matcher{Type: labels.MatchEqual, Name: "t2", Value: "1"}},
		Equal:          []string{"e"},
	}

	ih := NewInhibitor(nil, []amcommoncfg.InhibitRule{rule1, rule2}, nopLogger, eventrecorder.NopRecorder())
	now := time.Now()
	// Active alert that matches the source filter of rule1.
	sourceAlert1 := &alert.Alert{
		Alert: model.Alert{
			Labels:   model.LabelSet{"s1": "1", "t1": "2", "e": "1"},
			StartsAt: now.Add(-time.Minute),
			EndsAt:   now.Add(time.Hour),
		},
	}
	// Active alert that matches the source filter _and_ the target filter of rule2.
	sourceAlert2 := &alert.Alert{
		Alert: model.Alert{
			Labels:   model.LabelSet{"s2": "1", "t2": "1", "e": "1"},
			StartsAt: now.Add(-time.Minute),
			EndsAt:   now.Add(time.Hour),
		},
	}

	ih.rules[0].scache = store.NewAlerts()
	ih.rules[0].scache.Set(sourceAlert1)
	ih.rules[0].sindex = newIndex()
	ih.rules[0].updateIndex(sourceAlert1)

	ih.rules[1].scache = store.NewAlerts()
	ih.rules[1].scache.Set(sourceAlert2)
	ih.rules[1].sindex = newIndex()
	ih.rules[1].updateIndex(sourceAlert2)

	cases := []struct {
		target   model.LabelSet
		expected bool
	}{
		{
			// Matches target filter of rule1, inhibited.
			target:   model.LabelSet{"t1": "1", "e": "1"},
			expected: false,
		},
		{
			// Matches target filter of rule2, inhibited.
			target:   model.LabelSet{"t2": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches target filter of rule1 (plus noise), inhibited.
			target:   model.LabelSet{"t1": "1", "t3": "1", "e": "1"},
			expected: false,
		},
		{
			// Matches target filter of rule1 plus rule2, inhibited.
			target:   model.LabelSet{"t1": "1", "t2": "1", "e": "1"},
			expected: true,
		},
		{
			// Doesn't match target filter, not inhibited.
			target:   model.LabelSet{"t1": "0", "e": "1"},
			expected: true,
		},
		{
			// Matches both source and target filters of rule1,
			// inhibited because sourceAlert1 matches only the
			// source filter of rule1.
			target:   model.LabelSet{"s1": "1", "t1": "1", "e": "1"},
			expected: false,
		},
		{
			// Matches both source and target filters of rule2,
			// not inhibited because sourceAlert2 matches also both the
			// source and target filter of rule2.
			target:   model.LabelSet{"s2": "1", "t2": "1", "e": "1"},
			expected: true,
		},
		{
			// Matches target filter, equal label doesn't match, not inhibited
			target:   model.LabelSet{"t1": "1", "e": "0"},
			expected: false,
		},
	}

	for _, c := range cases {
		checkMutes(t, ih, c.target, c.expected, "target %v", c.target)
	}
}

func TestInhibitRuleName(t *testing.T) {
	t.Parallel()

	config1 := amcommoncfg.InhibitRule{
		Name: "test-rule",
		SourceMatchers: []*labels.Matcher{
			{Type: labels.MatchEqual, Name: "severity", Value: "critical"},
		},
		TargetMatchers: []*labels.Matcher{
			{Type: labels.MatchEqual, Name: "severity", Value: "warning"},
		},
		Equal: []string{"instance"},
	}
	config2 := amcommoncfg.InhibitRule{
		SourceMatchers: []*labels.Matcher{
			{Type: labels.MatchEqual, Name: "severity", Value: "critical"},
		},
		TargetMatchers: []*labels.Matcher{
			{Type: labels.MatchEqual, Name: "severity", Value: "warning"},
		},
		Equal: []string{"instance"},
	}

	rule1 := NewInhibitRule(config1)
	rule2 := NewInhibitRule(config2)

	require.Equal(t, "test-rule", rule1.Name, "Expected named rule to have adopt name from config")
	require.Empty(t, rule2.Name, "Expected unnamed rule to have empty name")
}

type fakeAlerts struct {
	alerts   []*alert.Alert
	finished chan struct{}
}

func newFakeAlerts(alerts []*alert.Alert) *fakeAlerts {
	return &fakeAlerts{
		alerts:   alerts,
		finished: make(chan struct{}),
	}
}

func (f *fakeAlerts) GetPending() provider.AlertIterator          { return nil }
func (f *fakeAlerts) Get(model.Fingerprint) (*alert.Alert, error) { return nil, nil }
func (f *fakeAlerts) Put(context.Context, ...*alert.Alert) error  { return nil }
func (f *fakeAlerts) Subscribe(name string) provider.AlertIterator {
	ch := make(chan *provider.Alert)
	done := make(chan struct{})
	go func() {
		for _, a := range f.alerts {
			ch <- &provider.Alert{
				Data:   a,
				Header: map[string]string{},
			}
		}
		// Send another (meaningless) alert to make sure that the inhibitor has
		// processed everything.
		ch <- &provider.Alert{
			Data: &alert.Alert{
				Alert: model.Alert{
					Labels:   model.LabelSet{},
					StartsAt: time.Now(),
				},
			},
			Header: map[string]string{},
		}
		close(f.finished)
		<-done
	}()
	return provider.NewAlertIterator(ch, done, nil)
}

func (f *fakeAlerts) SlurpAndSubscribe(name string) ([]*alert.Alert, provider.AlertIterator) {
	ch := make(chan *provider.Alert)
	done := make(chan struct{})
	go func() {
		for _, a := range f.alerts {
			ch <- &provider.Alert{
				Data:   a,
				Header: map[string]string{},
			}
		}
		// Send another (meaningless) alert to make sure that the inhibitor has
		// processed everything.
		ch <- &provider.Alert{
			Data: &alert.Alert{
				Alert: model.Alert{
					Labels:   model.LabelSet{},
					StartsAt: time.Now(),
				},
			},
			Header: map[string]string{},
		}
		close(f.finished)
		<-done
	}()
	return []*alert.Alert{}, provider.NewAlertIterator(ch, done, nil)
}

func TestInhibit(t *testing.T) {
	t.Parallel()

	now := time.Now()
	inhibitRule := func() amcommoncfg.InhibitRule {
		return amcommoncfg.InhibitRule{
			SourceMatch: map[string]string{"s": "1"},
			TargetMatch: map[string]string{"t": "1"},
			Equal:       []string{"e"},
		}
	}
	// alertOne is muted by alertTwo when it is active.
	alertOne := func() *alert.Alert {
		return &alert.Alert{
			Alert: model.Alert{
				Labels:   model.LabelSet{"t": "1", "e": "f"},
				StartsAt: now.Add(-time.Minute),
				EndsAt:   now.Add(time.Hour),
			},
		}
	}
	alertTwo := func(resolved bool) *alert.Alert {
		var end time.Time
		if resolved {
			end = now.Add(-time.Second)
		} else {
			end = now.Add(time.Hour)
		}
		return &alert.Alert{
			Alert: model.Alert{
				Labels:   model.LabelSet{"s": "1", "e": "f"},
				StartsAt: now.Add(-time.Minute),
				EndsAt:   end,
			},
		}
	}

	type exp struct {
		lbls  model.LabelSet
		muted bool
	}
	for i, tc := range []struct {
		alerts   []*alert.Alert
		expected []exp
	}{
		{
			// alertOne shouldn't be muted since alertTwo hasn't fired.
			alerts: []*alert.Alert{alertOne()},
			expected: []exp{
				{
					lbls:  model.LabelSet{"t": "1", "e": "f"},
					muted: false,
				},
			},
		},
		{
			// alertOne should be muted by alertTwo which is active.
			alerts: []*alert.Alert{alertOne(), alertTwo(false)},
			expected: []exp{
				{
					lbls:  model.LabelSet{"t": "1", "e": "f"},
					muted: true,
				},
				{
					lbls:  model.LabelSet{"s": "1", "e": "f"},
					muted: false,
				},
			},
		},
		{
			// alertOne shouldn't be muted since alertTwo is resolved.
			alerts: []*alert.Alert{alertOne(), alertTwo(false), alertTwo(true)},
			expected: []exp{
				{
					lbls:  model.LabelSet{"t": "1", "e": "f"},
					muted: false,
				},
				{
					lbls:  model.LabelSet{"s": "1", "e": "f"},
					muted: false,
				},
			},
		},
	} {
		ap := newFakeAlerts(tc.alerts)
		inhibitor := NewInhibitor(ap, []amcommoncfg.InhibitRule{inhibitRule()}, nopLogger, eventrecorder.NopRecorder())

		go func() {
			for ap.finished != nil {
				select {
				case <-ap.finished:
					ap.finished = nil
				default:
				}
			}
			inhibitor.Stop()
		}()
		inhibitor.Run()

		for _, expected := range tc.expected {
			checkMutes(t, inhibitor, expected.lbls, expected.muted, "tc: %d, labels %q", i, expected.lbls)
		}
	}
}

func TestInhibitRule_fingerprintEquals(t *testing.T) {
	rule := &InhibitRule{
		Equal: map[model.LabelName]struct{}{
			"cluster": {},
			"service": {},
		},
	}

	lset := model.LabelSet{
		"cluster":  "prod",
		"service":  "api",
		"instance": "host1",
	}

	fp := rule.fingerprintEquals(lset)

	// Same equal labels should produce same fingerprint
	lset2 := model.LabelSet{
		"cluster":  "prod",
		"service":  "api",
		"instance": "host2", // different non-equal label
	}
	require.Equal(t, fp, rule.fingerprintEquals(lset2))

	// Different equal label value should produce different fingerprint
	lset3 := model.LabelSet{
		"cluster": "staging",
		"service": "api",
	}
	require.NotEqual(t, fp, rule.fingerprintEquals(lset3))
}

func BenchmarkFingerprintEquals(b *testing.B) {
	// Test fingerprintEquals with varying number of equal labels
	for _, numLabels := range []int{1, 3, 5, 10} {
		b.Run(fmt.Sprintf("%d_equal_labels", numLabels), func(b *testing.B) {
			equalLabels := make(map[model.LabelName]struct{}, numLabels)
			for i := range numLabels {
				equalLabels[model.LabelName(fmt.Sprintf("label_%d", i))] = struct{}{}
			}

			rule := &InhibitRule{Equal: equalLabels}

			// Create a label set with matching values
			lset := make(model.LabelSet, numLabels+2)
			lset["source"] = "true"
			lset["target"] = "true"
			for i := range numLabels {
				lset[model.LabelName(fmt.Sprintf("label_%d", i))] = model.LabelValue(fmt.Sprintf("value_%d", i))
			}

			b.ResetTimer()
			b.ReportAllocs()

			for b.Loop() {
				_ = rule.fingerprintEquals(lset)
			}
		})
	}
}
