package router

import (
	"context"
	"errors"
	"net/netip"
	"sync"
	"testing"
	"time"

	"buf.build/go/protovalidate"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/nettool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/redistool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_redis"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/proto"
	clocktesting "k8s.io/utils/clock/testing"
)

var (
	_ Registerer = &RedisTracker{}
	_ Tracker    = &RedisTracker{}
	_ Querier    = &RedisTracker{}
)

const (
	ttl = time.Minute
)

var (
	selfURLMultiURL = nettool.NewMultiURLForAddresses("grpc", "", 10, []netip.Addr{netip.AddrFrom4([4]byte{1, 1, 1, 1})})
	// protobuf panics with NPE if we try to marshal the proto during init() in the same package :facepalm:
	selfURLProto = sync.OnceValue(func() string {
		return string(mustToProto(selfURLMultiURL))
	})
)

func TestRegisterConnection(t *testing.T) {
	ctrl := gomock.NewController(t)
	b := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, tm := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b),
		b.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()).
			Do(func(keys []api.AgentKey, ttl time.Duration, kvs ...redistool.BuilderKV[string]) {
				require.Len(t, kvs, 1)
				assert.Equal(t, selfURLProto(), kvs[0].HashKey)
				assert.Equal(t, tm.Add(ttl).Unix(), kvs[0].Value.ExpiresAt)
			}),
		b.EXPECT().
			Do(gomock.Any()),
	)
	bld := r.RegistrationBuilder()
	bld.Register(ttl, testhelpers.AgentkKey1)
	assert.NoError(t, bld.Do(context.Background()))
}

func TestUnregisterConnection(t *testing.T) {
	ctrl := gomock.NewController(t)
	b1 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	b2 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, _ := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b1),
		b1.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()),
		b1.EXPECT().
			Do(gomock.Any()),
		hash.EXPECT().
			IOBuilder().
			Return(b2),
		b2.EXPECT().
			Unset([]api.AgentKey{testhelpers.AgentkKey1}, selfURLProto()),
		b2.EXPECT().
			Do(gomock.Any()),
	)

	bld := r.RegistrationBuilder()
	bld.Register(ttl, testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))

	bld = r.RegistrationBuilder()
	bld.Unregister(testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))
}

func TestUnregisterConnection_Batched(t *testing.T) {
	ctrl := gomock.NewController(t)
	b := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, _ := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b),
		b.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()),
		b.EXPECT().
			Unset([]api.AgentKey{testhelpers.AgentkKey1}, selfURLProto()),
		b.EXPECT().
			Do(gomock.Any()),
	)

	bld := r.RegistrationBuilder()
	bld.Register(ttl, testhelpers.AgentkKey1)
	bld.Unregister(testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))
}

func TestUnregisterConnection_TwoConnections(t *testing.T) {
	ctrl := gomock.NewController(t)
	b1 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	b2 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	b3 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	b4 := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, _ := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b1),
		b1.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()),
		b1.EXPECT().
			Do(gomock.Any()),
		hash.EXPECT().
			IOBuilder().
			Return(b2),
		b2.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()),
		b2.EXPECT().
			Do(gomock.Any()),
		hash.EXPECT().
			IOBuilder().
			Return(b3),
		b3.EXPECT().
			Unset([]api.AgentKey{testhelpers.AgentkKey1}, selfURLProto()),
		b3.EXPECT().
			Do(gomock.Any()),
		hash.EXPECT().
			IOBuilder().
			Return(b4),
		b4.EXPECT().
			Unset([]api.AgentKey{testhelpers.AgentkKey1}, selfURLProto()),
		b4.EXPECT().
			Do(gomock.Any()),
	)

	bld := r.RegistrationBuilder()
	bld.Register(ttl, testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))

	bld = r.RegistrationBuilder()
	bld.Register(ttl, testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))

	bld = r.RegistrationBuilder()
	bld.Unregister(testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))

	bld = r.RegistrationBuilder()
	bld.Unregister(testhelpers.AgentkKey1)
	require.NoError(t, bld.Do(context.Background()))
}

func TestKASURLsByAgentKey_HappyPath(t *testing.T) {
	r, hash, _ := setupTracker(t)
	hash.EXPECT().
		Scan(gomock.Any(), testhelpers.AgentkKey1).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{RawHashKey: selfURLProto()}, nil)
			assert.True(t, keepGoing)
		})
	kasURLs, err := r.KASURLsByAgentKey(context.Background(), testhelpers.AgentkKey1)
	require.NoError(t, err)
	assert.Equal(t, []nettool.MultiURL{selfURLMultiURL}, kasURLs)
}

func TestKASURLsByAgentKey_ScanError(t *testing.T) {
	r, hash, _ := setupTracker(t)
	hash.EXPECT().
		Scan(gomock.Any(), testhelpers.AgentkKey1).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{}, errors.New("intended error2"))
			assert.True(t, keepGoing)
		})
	kasURLs, err := r.KASURLsByAgentKey(context.Background(), testhelpers.AgentkKey1)
	assert.EqualError(t, err, "intended error2")
	assert.Empty(t, kasURLs)
}

func TestRefresh_NoAgents(t *testing.T) {
	ctrl := gomock.NewController(t)
	b := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, _ := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b),
		b.EXPECT().
			Set(gomock.Len(0), ttl, gomock.Any()),
		b.EXPECT().
			Do(gomock.Any()),
	)
	builder := r.RegistrationBuilder()
	builder.Refresh(ttl)
	assert.NoError(t, builder.Do(context.Background()))
}

func TestRefresh_OneAgent(t *testing.T) {
	ctrl := gomock.NewController(t)
	b := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, tm := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b),
		b.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1}, ttl, gomock.Any()).
			Do(func(keys []api.AgentKey, ttl time.Duration, kvs ...redistool.BuilderKV[string]) {
				require.Len(t, kvs, 1)
				assert.Equal(t, selfURLProto(), kvs[0].HashKey)
				assert.Equal(t, tm.Add(ttl).Unix(), kvs[0].Value.ExpiresAt)
			}),
		b.EXPECT().
			Do(gomock.Any()),
	)
	builder := r.RegistrationBuilder()
	builder.Refresh(ttl, testhelpers.AgentkKey1)
	assert.NoError(t, builder.Do(context.Background()))
}

func TestRefresh_TwoAgents(t *testing.T) {
	ctrl := gomock.NewController(t)
	b := mock_redis.NewMockIOBuilder[api.AgentKey, string](ctrl)
	r, hash, tm := setupTracker(t)
	gomock.InOrder(
		hash.EXPECT().
			IOBuilder().
			Return(b),
		b.EXPECT().
			Set([]api.AgentKey{testhelpers.AgentkKey1, testhelpers.AgentkKey2}, ttl, gomock.Any()).
			Do(func(keys []api.AgentKey, ttl time.Duration, kvs ...redistool.BuilderKV[string]) {
				require.Len(t, kvs, 1)
				assert.Equal(t, selfURLProto(), kvs[0].HashKey)
				assert.Equal(t, tm.Add(ttl).Unix(), kvs[0].Value.ExpiresAt)
			}),
		b.EXPECT().
			Do(gomock.Any()),
	)
	builder := r.RegistrationBuilder()
	builder.Refresh(ttl, testhelpers.AgentkKey1, testhelpers.AgentkKey2)
	assert.NoError(t, builder.Do(context.Background()))
}

func setupTracker(t *testing.T) (*RedisTracker, *mock_redis.MockExpiringHashAPI[api.AgentKey, string], time.Time) {
	ctrl := gomock.NewController(t)
	hash := mock_redis.NewMockExpiringHashAPI[api.AgentKey, string](ctrl)
	tm := time.Now()
	v, err := protovalidate.New()
	require.NoError(t, err)
	return &RedisTracker{
		validator:         v,
		ownPrivateAPIURL:  selfURLProto(),
		clock:             clocktesting.NewFakePassiveClock(tm),
		tunnelsByAgentKey: hash,
	}, hash, tm
}

func mustToProto(u nettool.MultiURL) []byte {
	data, err := proto.Marshal(fromMultiURL(u))
	if err != nil {
		panic(err)
	}
	return data
}
