package router

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

	"github.com/stretchr/testify/assert"
	"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/testing/mock_modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tunserver"
	"go.uber.org/mock/gomock"
	"k8s.io/apimachinery/pkg/util/wait"
)

var (
	_    tunserver.PollingGatewayURLQuerier[nettool.MultiURL, api.AgentKey] = (*AggregatingQuerier)(nil)
	url1                                                                    = nettool.NewMultiURLForAddresses("grpc", "", 123, []netip.Addr{netip.AddrFrom4([4]byte{127, 0, 0, 1})})
	url2                                                                    = nettool.NewMultiURLForAddresses("grpcs", "", 321, []netip.Addr{netip.AddrFrom4([4]byte{127, 0, 0, 2})})
)

func TestPollGatewayURLs_OnlyStartsSinglePoll(t *testing.T) {
	ctrl := gomock.NewController(t)
	q := NewMockQuerier(ctrl)
	q.EXPECT().
		KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1)
	api := mock_modshared.NewMockAPI(ctrl)
	aq := NewAggregatingQuerier(testlogger.New(t), q, api, nt(), time.Minute, time.Minute)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	go aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		assert.Fail(t, "unexpected call")
	})
	aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		assert.Fail(t, "unexpected call")
	})
}

func TestPollGatewayURLs_PollingCycle(t *testing.T) {
	ctrl := gomock.NewController(t)
	q := NewMockQuerier(ctrl)
	q.EXPECT().
		KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
		Return([]nettool.MultiURL{url1, url2}, nil)
	api := mock_modshared.NewMockAPI(ctrl)
	aq := NewAggregatingQuerier(testlogger.New(t), q, api, nt(), time.Minute, time.Minute)
	call := 0
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		switch call {
		case 0:
			assert.Equal(t, []nettool.MultiURL{url1, url2}, kasURLs)
		default:
			assert.FailNow(t, "unexpected invocation")
		}
		call++
	})
}

func TestPollGatewayURLs_CacheAfterStopped(t *testing.T) {
	ctrl := gomock.NewController(t)
	q := NewMockQuerier(ctrl)
	gomock.InOrder(
		q.EXPECT().
			KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
			Return([]nettool.MultiURL{url1}, nil),
		q.EXPECT().
			KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
			Return([]nettool.MultiURL{url2}, nil),
	)
	api := mock_modshared.NewMockAPI(ctrl)
	aq := NewAggregatingQuerier(testlogger.New(t), q, api, nt(), time.Minute, time.Minute)
	ctx, cancel := context.WithCancel(context.Background())
	aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		assert.Equal(t, []nettool.MultiURL{url1}, kasURLs)
		cancel()
	})
	kasURLs := aq.CachedGatewayURLs(testhelpers.AgentkKey1) // from cache
	assert.Equal(t, []nettool.MultiURL{url1}, kasURLs)
	ctx, cancel = context.WithCancel(context.Background())
	aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		assert.Equal(t, []nettool.MultiURL{url2}, kasURLs) // from redis
		cancel()
	})
}

func TestPollGatewayURLs_CacheWhenRunning(t *testing.T) {
	ctrl := gomock.NewController(t)
	q := NewMockQuerier(ctrl)
	modAPI := mock_modshared.NewMockAPI(ctrl)
	aq := NewAggregatingQuerier(testlogger.New(t), q, modAPI, nt(), time.Second, time.Minute)
	start1 := make(chan struct{})
	gomock.InOrder(
		q.EXPECT().
			KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
			Return([]nettool.MultiURL{url1}, nil),
		q.EXPECT().
			KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
			DoAndReturn(func(ctx context.Context, agentKey api.AgentKey) ([]nettool.MultiURL, error) {
				close(start1)                      // start concurrent query
				assert.Eventually(t, func() bool { // wait for aq.PollGatewayURLs() to register second callback
					aq.mu.Lock()
					defer aq.mu.Unlock()
					return aq.listeners[agentKey].subs.Len() == 2
				}, time.Second, 10*time.Millisecond)
				return []nettool.MultiURL{url2}, nil
			}),
	)
	ctx := t.Context()
	count1 := 0
	go aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		switch count1 {
		case 0:
			assert.Equal(t, []nettool.MultiURL{url1}, kasURLs) // first call
		case 1:
			assert.Equal(t, []nettool.MultiURL{url2}, kasURLs) // second call
		default:
			assert.Fail(t, "unexpected invocation")
		}
		count1++
	})
	<-start1
	kasURLs := aq.CachedGatewayURLs(testhelpers.AgentkKey1)
	assert.Equal(t, []nettool.MultiURL{url1}, kasURLs) // from cache
	count2 := 0
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()
	aq.PollGatewayURLs(ctx2, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		switch count2 {
		case 0:
			assert.Equal(t, []nettool.MultiURL{url2}, kasURLs) // from redis
			cancel2()
		default:
			assert.FailNow(t, "unexpected invocation")
		}
		count2++
	})
	assert.Equal(t, 1, count2)
}

func TestPollGatewayURLs_GCRemovesExpiredCache(t *testing.T) {
	ctrl := gomock.NewController(t)
	q := NewMockQuerier(ctrl)
	q.EXPECT().
		KASURLsByAgentKey(gomock.Any(), testhelpers.AgentkKey1).
		Return([]nettool.MultiURL{url1}, nil)
	api := mock_modshared.NewMockAPI(ctrl)
	gcPeriod := time.Second
	aq := NewAggregatingQuerier(testlogger.New(t), q, api, nt(), time.Minute, gcPeriod)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	aq.PollGatewayURLs(ctx, testhelpers.AgentkKey1, func(kasURLs []nettool.MultiURL) {
		cancel()
	})
	ctx, cancel = context.WithCancel(context.Background())
	var wg wait.Group
	defer wg.Wait()
	defer cancel()
	wg.Start(func() {
		_ = aq.Run(ctx)
	})
	time.Sleep(gcPeriod * 2)
	kasURLs := aq.CachedGatewayURLs(testhelpers.AgentkKey1)
	assert.Empty(t, kasURLs)
}
