// Copyright 2019 Antrea 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 networkpolicy

import (
	"fmt"
	"math"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/utils/clock"
	clocktesting "k8s.io/utils/clock/testing"

	"antrea.io/antrea/pkg/agent/types"
	"antrea.io/antrea/pkg/apis/controlplane/v1beta2"
)

func TestNewIDAllocator(t *testing.T) {
	tests := []struct {
		name                    string
		args                    []uint32
		expectedLastAllocatedID uint32
		expectedAvailableSets   map[uint32]struct{}
		expectedAvailableSlice  []uint32
	}{
		{
			"zero-allocated-ids",
			nil,
			0,
			map[uint32]struct{}{},
			nil,
		},
		{
			"consecutive-allocated-ids",
			[]uint32{1, 2},
			2,
			map[uint32]struct{}{},
			nil,
		},
		{
			"inconsecutive-allocated-ids",
			[]uint32{2, 7, 5},
			7,
			map[uint32]struct{}{1: {}, 3: {}, 4: {}, 6: {}},
			[]uint32{1, 3, 4, 6},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := newIDAllocator(MinAllocatorAsyncDeleteInterval, tt.args...)
			assert.Equalf(t, tt.expectedLastAllocatedID, got.lastAllocatedID, "Got lastAllocatedID %v, expected %v", got.lastAllocatedID, tt.expectedLastAllocatedID)
			assert.Equalf(t, tt.expectedAvailableSets, got.availableSet, "Got availableSet %v, expected %v", got.availableSet, tt.expectedAvailableSets)
			assert.Equalf(t, tt.expectedAvailableSlice, got.availableSlice, "Got availableSlice %v, expected %v", got.availableSlice, tt.expectedAvailableSlice)
		})
	}
}

func TestAllocateForRule(t *testing.T) {
	rule := &types.PolicyRule{
		Direction: v1beta2.DirectionIn,
		From:      []types.Address{},
		To:        ofPortsToOFAddresses(sets.New[int32](1)),
		Service:   nil,
	}
	tests := []struct {
		name       string
		args       []uint32
		rule       *types.PolicyRule
		expectedID uint32
	}{
		{
			"zero-allocated-ids",
			nil,
			rule,
			1,
		},
		{
			"consecutive-allocated-ids",
			[]uint32{1, 2},
			rule,
			3,
		},
		{
			"inconsecutive-allocated-ids",
			[]uint32{1, 7, 5},
			rule,
			2,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := newIDAllocator(MinAllocatorAsyncDeleteInterval, tt.args...)
			err := a.allocateForRule(tt.rule)
			require.NoError(t, err)
			assert.Equalf(t, tt.expectedID, tt.rule.FlowID, "Got id %v, expected %v", tt.rule.FlowID, tt.expectedID)
			ruleFromCache, exists, err := a.getRuleFromAsyncCache(tt.expectedID)
			assert.Truef(t, exists, "Rule with id %d should present in the async rule cache", tt.expectedID)
			assert.NoErrorf(t, err, "getRuleFromAsyncCache should return valid rule with id: %v", tt.expectedID)
			assert.Equalf(t, tt.rule, ruleFromCache, "getRuleFromAsyncCache should return expected rule")
		})
	}
}

func TestAllocateForRuleOutOfIDs(t *testing.T) {
	a := newIDAllocator(MinAllocatorAsyncDeleteInterval)
	// Overwrite lastAllocatedID to trigger the error in allocateForRule
	a.lastAllocatedID = math.MaxUint32
	require.EqualError(t, a.allocateForRule(&types.PolicyRule{}), "no ID available")
}

func TestRelease(t *testing.T) {
	tests := []struct {
		name                   string
		newArgs                []uint32
		releaseArgs            uint32
		expectedErr            error
		expectedAvailableSets  map[uint32]struct{}
		expectedAvailableSlice []uint32
	}{
		{
			"duplicate-release",
			[]uint32{2},
			1,
			fmt.Errorf("ID %d has been released, duplicate release is not allowed", 1),
			map[uint32]struct{}{1: {}},
			[]uint32{1},
		},
		{
			"invalid-release",
			[]uint32{1, 2},
			5,
			fmt.Errorf("ID %d was not allocated, can't be released", 5),
			map[uint32]struct{}{},
			nil,
		},
		{
			"valid-release",
			[]uint32{1, 2, 3},
			2,
			nil,
			map[uint32]struct{}{2: {}},
			[]uint32{2},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := newIDAllocator(MinAllocatorAsyncDeleteInterval, tt.newArgs...)
			actualErr := a.release(tt.releaseArgs)
			assert.Equalf(t, tt.expectedErr, actualErr, "Got error %v, expected %v", actualErr, tt.expectedErr)
			assert.Equalf(t, tt.expectedAvailableSets, a.availableSet, "Got availableSet %v, expected %v", a.availableSet, tt.expectedAvailableSets)
			assert.Equalf(t, tt.expectedAvailableSlice, a.availableSlice, "Got availableSlice %v, expected %v", a.availableSlice, tt.expectedAvailableSlice)
		})
	}
}

// fakeClock is a wrapper around clocktesting.FakeClock that tracks the number
// of times NewTimer has been called, so we can write a race-free test.
type fakeClock struct {
	*clocktesting.FakeClock
	timersAdded atomic.Int32
}

func newFakeClock(t time.Time) *fakeClock {
	return &fakeClock{
		FakeClock: clocktesting.NewFakeClock(t),
	}
}

func (c *fakeClock) TimersAdded() int32 {
	return c.timersAdded.Load()
}

func (c *fakeClock) NewTimer(d time.Duration) clock.Timer {
	defer c.timersAdded.Add(1)
	return c.FakeClock.NewTimer(d)
}

func TestIdAllocatorWorker(t *testing.T) {
	rule := &types.PolicyRule{
		Direction: v1beta2.DirectionIn,
		From:      []types.Address{},
		To:        ofPortsToOFAddresses(sets.New[int32](1)),
		Service:   nil,
	}
	tests := []struct {
		name                    string
		args                    []uint32
		testAsyncDeleteInterval time.Duration
		expectedDeleteInterval  time.Duration
		rule                    *types.PolicyRule
		expectedID              uint32
	}{
		{
			// Provided interval is less than MinAllocatorAsyncDeleteInterval, so the
			// rule should take MinAllocatorAsyncDeleteInterval to be deleted.
			"delete-rule-with-test-min-async-delete-interval",
			nil,
			MinAllocatorAsyncDeleteInterval / 2,
			MinAllocatorAsyncDeleteInterval,
			rule,
			1,
		},
		{
			// Provided interval is more than MinAllocatorAsyncDeleteInterval, so the
			// rule should be deleted after the provided interval.
			"delete-rule-with-test-async-delete-interval",
			nil,
			2 * MinAllocatorAsyncDeleteInterval,
			2 * MinAllocatorAsyncDeleteInterval,
			rule,
			1,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			startTime := time.Now()
			expectedDeleteTime := startTime.Add(tt.expectedDeleteInterval)
			fakeClock := newFakeClock(startTime)
			a := newIDAllocatorWithClock(tt.testAsyncDeleteInterval, fakeClock, tt.args...)
			require.NoError(t, a.allocateForRule(tt.rule), "Error allocating ID for rule")
			stopCh := make(chan struct{})
			defer close(stopCh)
			go a.runWorker(stopCh)

			require.Zero(t, fakeClock.TimersAdded())
			a.forgetRule(tt.rule.FlowID)

			ruleHasBeenDeleted := func() bool {
				a.Lock()
				defer a.Unlock()
				return len(a.availableSlice) > 0
			}

			// Wait for the DelayingQueue to create the timer which signals that the
			// rule is ready to be deleted.
			require.Eventually(t, func() bool {
				return fakeClock.TimersAdded() > 0
			}, 1*time.Second, 10*time.Millisecond)

			fakeClock.SetTime(expectedDeleteTime.Add(-10 * time.Millisecond))

			// We wait for a small duration and ensure that the rule is not deleted.
			require.Never(t, func() bool {
				return ruleHasBeenDeleted()
			}, 100*time.Millisecond, 10*time.Millisecond, "Rule ID was unexpectedly released")
			_, exists, err := a.getRuleFromAsyncCache(tt.expectedID)
			require.NoError(t, err)
			assert.True(t, exists, "Rule should be present in asyncRuleCache")

			fakeClock.SetTime(expectedDeleteTime.Add(10 * time.Millisecond))

			require.Eventually(t, func() bool {
				return ruleHasBeenDeleted()
			}, 1*time.Second, 10*time.Millisecond, "Rule ID was not released")
			_, exists, err = a.getRuleFromAsyncCache(tt.expectedID)
			require.NoError(t, err)
			assert.False(t, exists, "Rule should not be present in asyncRuleCache")
		})
	}
}

func TestVlanIDAllocator(t *testing.T) {
	vlanIDAllocator := newL7VlanIDAllocator()
	ruleID1 := "rule1"
	ruleID2 := "rule2"
	ruleID3 := "rule3"
	ruleID4 := "rule4"

	vlanID1 := vlanIDAllocator.allocate(ruleID1)
	assert.Equal(t, vlanID1, vlanIDAllocator.query(ruleID1))

	vlanID2 := vlanIDAllocator.allocate(ruleID2)
	assert.Equal(t, vlanID2, vlanIDAllocator.query(ruleID2))

	vlanIDAllocator.release(ruleID1)
	assert.Equal(t, uint32(0), vlanIDAllocator.query(ruleID1))

	vlanIDAllocator.release(ruleID2)
	assert.Equal(t, uint32(0), vlanIDAllocator.query(ruleID2))

	vlanID3 := vlanIDAllocator.allocate(ruleID3)
	assert.Equal(t, vlanID3, vlanIDAllocator.query(ruleID3))
	assert.Equal(t, vlanID2, vlanID3)

	vlanID4 := vlanIDAllocator.allocate(ruleID4)
	assert.Equal(t, vlanID4, vlanIDAllocator.query(ruleID4))
	assert.Equal(t, vlanID1, vlanID4)
}
