package agent_tracker //nolint:staticcheck

import (
	"context"
	"errors"
	"runtime"
	"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/redistool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/matcher"
	"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/mock_tool"
	"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/gitlab-agent/v18/pkg/entity"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

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

func TestGC_HappyPath(t *testing.T) {
	r, connectedAgents, agentkConnectionsByID, agentVersions, agentwConnectionsByID, connectionsByAgentVersion, _, _, _ := setupTracker(t)

	wasCalled1 := false
	wasCalled2 := false
	wasCalled3 := false
	wasCalled4 := false
	wasCalled5 := false
	wasCalled6 := false

	connectedAgents.EXPECT().
		GC(gomock.Any()).
		DoAndReturn(func(_ context.Context) (int, error) {
			wasCalled3 = true
			return 3, nil
		})

	agentkConnectionsByID.EXPECT().
		GC(gomock.Any()).
		DoAndReturn(func(_ context.Context) (int, error) {
			wasCalled2 = true
			return 2, nil
		})

	agentwConnectionsByID.EXPECT().
		GC(gomock.Any()).
		DoAndReturn(func(_ context.Context) (int, error) {
			wasCalled1 = true
			return 1, nil
		})

	gomock.InOrder(
		agentVersions.EXPECT().
			GC(gomock.Any()).
			DoAndReturn(func(_ context.Context) (int, error) {
				wasCalled6 = true
				return 4, nil
			}),
		agentVersions.EXPECT().
			Scan(gomock.Any(), agentVersionKey).
			Return(func(yield func(redistool.ScanEntry, error) bool) {
				wasCalled4 = true
				yield(redistool.ScanEntry{RawHashKey: "v16.9.0"}, nil)
			}),
		connectionsByAgentVersion.EXPECT().
			GCFor([]string{"v16.9.0"}).
			Return(func(_ context.Context) (int, error) {
				wasCalled5 = true
				return 5, nil
			}),
	)

	assert.Equal(t, 1+2+3+4+5, r.runGC(context.Background()))
	assert.True(t, wasCalled1)
	assert.True(t, wasCalled2)
	assert.True(t, wasCalled3)
	assert.True(t, wasCalled4)
	assert.True(t, wasCalled5)
	assert.True(t, wasCalled6)
}

func TestGC_AllCalledOnError(t *testing.T) {
	r, connectedAgents, agentkConnectionsByID, agentVersions, agentwConnectionsByID, connectionsByAgentVersion, rep, _, _ := setupTracker(t)

	wasCalled1 := false
	wasCalled2 := false
	wasCalled3 := false
	wasCalled4 := false
	wasCalled5 := false
	wasCalled6 := false

	gomock.InOrder(
		connectedAgents.EXPECT().
			GC(gomock.Any()).
			DoAndReturn(func(_ context.Context) (int, error) {
				wasCalled3 = true
				return 3, errors.New("err3")
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Failed to GC data in connected_agents Redis hash", matcher.ErrorEq("err3")),
	)

	gomock.InOrder(
		agentkConnectionsByID.EXPECT().
			GC(gomock.Any()).
			DoAndReturn(func(_ context.Context) (int, error) {
				wasCalled2 = true
				return 2, errors.New("err2")
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Failed to GC data in connections_by_agent_id Redis hash", matcher.ErrorEq("err2")),
	)

	gomock.InOrder(
		agentwConnectionsByID.EXPECT().
			GC(gomock.Any()).
			DoAndReturn(func(_ context.Context) (int, error) {
				wasCalled1 = true
				return 1, errors.New("err7")
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Failed to GC data in agentw_connections_by_id Redis hash", matcher.ErrorEq("err7")),
	)

	err4 := errors.New("err4")
	err5 := errors.New("err5")
	gomock.InOrder(
		agentVersions.EXPECT().
			GC(gomock.Any()).
			DoAndReturn(func(_ context.Context) (int, error) {
				wasCalled6 = true
				return 4, nil
			}),
		agentVersions.EXPECT().
			Scan(gomock.Any(), agentVersionKey).
			Return(func(yield func(redistool.ScanEntry, error) bool) {
				wasCalled4 = true
				if !yield(redistool.ScanEntry{RawHashKey: "v16.9.0"}, nil) {
					return
				}
				yield(redistool.ScanEntry{}, err4)
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "getAgentVersions: failed to scan redis hash", err4),
		connectionsByAgentVersion.EXPECT().
			GCFor([]string{"v16.9.0"}).
			Return(func(_ context.Context) (int, error) {
				wasCalled5 = true
				return 5, err5
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Failed to GC data in connections_by_agent_version Redis hash", err5),
	)

	assert.Equal(t, 1+2+3+4+5, r.runGC(context.Background()))
	assert.True(t, wasCalled1)
	assert.True(t, wasCalled2)
	assert.True(t, wasCalled3)
	assert.True(t, wasCalled4)
	assert.True(t, wasCalled5)
	assert.True(t, wasCalled6)
}

func TestGetAgentkConnectionsByID_HappyPath(t *testing.T) {
	r, _, agentkConnectionsByID, _, _, _, _, info, _ := setupTracker(t)
	infoBytes, err := proto.Marshal(info)
	require.NoError(t, err)
	agentkConnectionsByID.EXPECT().
		Scan(gomock.Any(), info.AgentId).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{RawHashKey: "k2", Value: infoBytes}, nil)
			assert.True(t, keepGoing)
		})
	var gotInfos int
	agentKey := api.AgentKey{
		ID:   info.AgentId,
		Type: api.AgentTypeKubernetes,
	}
	for infoActual := range r.GetAgentkConnectionsByID(context.Background(), agentKey.ID) {
		gotInfos++
		matcher.AssertProtoEqual(t, info, infoActual)
	}
	require.NoError(t, err)
	assert.Equal(t, 1, gotInfos)
}

func TestGetAgentkConnectionsByID_ScanError(t *testing.T) {
	r, _, agentkConnectionsByID, _, _, _, rep, info, _ := setupTracker(t)
	gomock.InOrder(
		agentkConnectionsByID.EXPECT().
			Scan(gomock.Any(), info.AgentId).
			Return(func(yield func(redistool.ScanEntry, error) bool) {
				keepGoing := yield(redistool.ScanEntry{}, errors.New("intended error"))
				assert.True(t, keepGoing) // ignores error to keep going
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Redis connections_by_agent_id hash scan", matcher.ErrorEq("intended error")),
	)
	agentKey := api.AgentKey{
		ID:   info.AgentId,
		Type: api.AgentTypeKubernetes,
	}
	for range r.GetAgentkConnectionsByID(context.Background(), agentKey.ID) {
		require.FailNow(t, "unexpected call")
	}
}

func TestGetAgentkConnectionsByID_UnmarshalError(t *testing.T) {
	r, _, agentkConnectionsByID, _, _, _, rep, info, _ := setupTracker(t)
	agentkConnectionsByID.EXPECT().
		Scan(gomock.Any(), info.AgentId).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{RawHashKey: "k2", Value: []byte{1, 2, 3}}, nil) // invalid bytes
			assert.True(t, keepGoing)                                                              // ignores error to keep going
		})
	rep.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), "Redis connections_by_agent_id hash scan: proto.Unmarshal(ConnectedAgentkInfo)", matcher.ErrorIs(proto.Error))
	agentKey := api.AgentKey{
		ID:   info.AgentId,
		Type: api.AgentTypeKubernetes,
	}
	for range r.GetAgentkConnectionsByID(context.Background(), agentKey.ID) {
		require.FailNow(t, "unexpected call")
	}
}

func TestGetAgentwConnectionsByID_HappyPath(t *testing.T) {
	r, _, _, _, agentwConnectionsByID, _, _, _, agentwInfo := setupTracker(t)
	infoBytes, err := proto.Marshal(agentwInfo)
	require.NoError(t, err)
	agentwConnectionsByID.EXPECT().
		Scan(gomock.Any(), agentwInfo.AgentId).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{RawHashKey: "k2", Value: infoBytes}, nil)
			assert.True(t, keepGoing)
		})
	var gotInfos int
	for infoActual := range r.GetAgentwConnectionsByID(context.Background(), agentwInfo.AgentId) {
		gotInfos++
		matcher.AssertProtoEqual(t, agentwInfo, infoActual)
	}
	assert.Equal(t, 1, gotInfos)
}

func TestGetAgentwConnectionsByID_ScanError(t *testing.T) {
	r, _, _, _, agentwConnectionsByID, _, rep, _, agentwInfo := setupTracker(t)
	gomock.InOrder(
		agentwConnectionsByID.EXPECT().
			Scan(gomock.Any(), agentwInfo.AgentId).
			Return(func(yield func(redistool.ScanEntry, error) bool) {
				keepGoing := yield(redistool.ScanEntry{}, errors.New("intended error"))
				assert.True(t, keepGoing) // ignores error to keep going
			}),
		rep.EXPECT().
			HandleProcessingError(gomock.Any(), gomock.Any(), "Redis agentw_connections_by_id hash scan", matcher.ErrorEq("intended error")),
	)
	for range r.GetAgentwConnectionsByID(context.Background(), agentwInfo.AgentId) {
		require.FailNow(t, "unexpected call")
	}
}

func TestGetAgentwConnectionsByID_UnmarshalError(t *testing.T) {
	r, _, _, _, agentwConnectionsByID, _, rep, _, agentwInfo := setupTracker(t)
	agentwConnectionsByID.EXPECT().
		Scan(gomock.Any(), agentwInfo.AgentId).
		Return(func(yield func(redistool.ScanEntry, error) bool) {
			keepGoing := yield(redistool.ScanEntry{RawHashKey: "k2", Value: []byte{1, 2, 3}}, nil) // invalid bytes
			assert.True(t, keepGoing)                                                              // ignores error to keep going
		})
	rep.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), "Redis agentw_connections_by_id hash scan: proto.Unmarshal(ConnectedAgentwInfo)", matcher.ErrorIs(proto.Error))
	for range r.GetAgentwConnectionsByID(context.Background(), agentwInfo.AgentId) {
		require.FailNow(t, "unexpected call")
	}
}

func TestGetConnectedAgentsCount_HappyPath(t *testing.T) {
	r, connectedAgents, _, _, _, _, _, _, _ := setupTracker(t)
	connectedAgents.EXPECT().
		Len(gomock.Any(), connectedAgentsKey).
		Return(int64(1), nil)
	size, err := r.GetConnectedAgentsCount(context.Background())
	require.NoError(t, err)
	assert.EqualValues(t, 1, size)
}

func TestGetConnectedAgentsCount_LenError(t *testing.T) {
	r, connectedAgents, _, _, _, _, _, _, _ := setupTracker(t)
	connectedAgents.EXPECT().
		Len(gomock.Any(), connectedAgentsKey).
		Return(int64(0), errors.New("intended error"))
	size, err := r.GetConnectedAgentsCount(context.Background())
	require.Error(t, err)
	assert.Zero(t, size)
}

func TestRedisTracker_RegisterAgentkExpiring(t *testing.T) {
	// GIVEN
	r, connectedAgents, agentkConnectionsByID, agentVersions, _, connectionsByAgentVersion, _, agentkInfo, _ := setupTracker(t)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	// EXPECTS
	connectedAgents.EXPECT().SetEX(gomock.Any(), connectedAgentsKey, agentkInfo.AgentId, gomock.Any(), gomock.Any())
	agentkConnectionsByID.EXPECT().SetEX(gomock.Any(), agentkInfo.AgentId, agentkInfo.ConnectionId, gomock.Any(), gomock.Any())
	connectionsByAgentVersion.EXPECT().SetEX(gomock.Any(), agentkInfo.AgentMeta.Version, agentkInfo.ConnectionId, gomock.Any(), gomock.Any())
	agentVersions.EXPECT().SetEX(gomock.Any(), agentVersionKey, agentkInfo.AgentMeta.Version, gomock.Any(), gomock.Any())

	// WHEN
	err := r.RegisterAgentkExpiring(ctx, agentkInfo)

	// THEN
	require.NoError(t, err)
}

func TestRedisTracker_RegisterAgentwExpiring(t *testing.T) {
	// GIVEN
	r, _, _, _, agentwConnectionsByID, _, _, _, agentwInfo := setupTracker(t)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	// EXPECTS
	agentwConnectionsByID.EXPECT().SetEX(gomock.Any(), agentwInfo.AgentId, agentwInfo.ConnectionId, gomock.Any(), gomock.Any())

	// WHEN
	err := r.RegisterAgentwExpiring(ctx, agentwInfo)

	// THEN
	require.NoError(t, err)
}

func TestRedisTracker_UnregisterAgentk(t *testing.T) {
	// GIVEN
	r, connectedAgents, connectionsByAgentID, agentVersions, _, connectionsByAgentVersion, _, _, _ := setupTracker(t)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	testAgentVersion := "any"
	testConnectionID := int64(123456789)
	agentkInfo := &DisconnectAgentkInfo{
		AgentMeta: &entity.AgentkMeta{
			Version: testAgentVersion,
		},
		ConnectionId: testConnectionID,
		AgentId:      testhelpers.AgentkKey1.ID,
		ProjectId:    testhelpers.ProjectID,
	}

	// EXPECTS
	connectedAgents.EXPECT().DelEX(gomock.Any(), connectedAgentsKey, testhelpers.AgentkKey1.ID)
	connectionsByAgentID.EXPECT().DelEX(gomock.Any(), testhelpers.AgentkKey1.ID, testConnectionID)
	agentVersions.EXPECT().DelEX(gomock.Any(), agentVersionKey, testAgentVersion)
	connectionsByAgentVersion.EXPECT().DelEX(gomock.Any(), testAgentVersion, testConnectionID)

	// WHEN
	err := r.UnregisterAgentk(ctx, agentkInfo)

	// THEN
	require.NoError(t, err)
}

func TestRedisTracker_UnregisterAgentw(t *testing.T) {
	// GIVEN
	r, _, _, _, agentwConnectionsByID, _, _, _, _ := setupTracker(t)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	testAgentVersion := "any"
	testConnectionID := int64(123456789)
	agentwInfo := &DisconnectAgentwInfo{
		AgentMeta: &entity.AgentwMeta{
			Version: testAgentVersion,
		},
		ConnectionId: testConnectionID,
		AgentId:      testhelpers.AgentwKey1.ID,
	}

	// EXPECTS
	agentwConnectionsByID.EXPECT().DelEX(gomock.Any(), testhelpers.AgentwKey1.ID, testConnectionID)

	// WHEN
	err := r.UnregisterAgentw(ctx, agentwInfo)

	// THEN
	require.NoError(t, err)
}

func setupTracker(t *testing.T) (*RedisTracker, *mock_redis.MockExpiringHash[int64, int64],
	*mock_redis.MockExpiringHash[int64, int64], *mock_redis.MockExpiringHash[int64, string],
	*mock_redis.MockExpiringHash[int64, int64], *mock_redis.MockExpiringHash[string, int64],
	*mock_tool.MockErrReporter, *ConnectedAgentkInfo, *ConnectedAgentwInfo) {
	ctrl := gomock.NewController(t)
	rep := mock_tool.NewMockErrReporter(ctrl)
	connectedAgents := mock_redis.NewMockExpiringHash[int64, int64](ctrl)
	agentkConnectionsByID := mock_redis.NewMockExpiringHash[int64, int64](ctrl)
	agentVersions := mock_redis.NewMockExpiringHash[int64, string](ctrl)
	agentwConnectionsByID := mock_redis.NewMockExpiringHash[int64, int64](ctrl)
	connectionsByAgentVersion := mock_redis.NewMockExpiringHash[string, int64](ctrl)
	v, err := protovalidate.New()
	require.NoError(t, err)

	tr := &RedisTracker{
		log:                       testlogger.New(t),
		errRep:                    rep,
		validator:                 v,
		gcPeriod:                  time.Minute,
		agentkConnectionsByID:     agentkConnectionsByID,
		connectedAgents:           connectedAgents,
		agentwConnectionsByID:     agentwConnectionsByID,
		agentVersions:             agentVersions,
		connectionsByAgentVersion: connectionsByAgentVersion,
	}
	return tr, connectedAgents, agentkConnectionsByID, agentVersions, agentwConnectionsByID, connectionsByAgentVersion, rep, agentkConnInfo(), agentwConnInfo()
}

func agentkConnInfo() *ConnectedAgentkInfo {
	return &ConnectedAgentkInfo{
		AgentMeta: &entity.AgentkMeta{
			Version:      "v1.2.3",
			GitRef:       "0123456789abcdef0123456789abcdef00000000",
			PodNamespace: "ns",
			PodName:      "name",
		},
		ConnectedAt:  timestamppb.Now(),
		ConnectionId: 123,
		AgentId:      345,
		ProjectId:    456,
	}
}

func agentwConnInfo() *ConnectedAgentwInfo {
	return &ConnectedAgentwInfo{
		AgentMeta: &entity.AgentwMeta{
			Version:      "v1.2.3",
			GitRef:       "9876543210fedcba9876543210fedcba00000000",
			Architecture: runtime.GOARCH,
		},
		ConnectedAt:  timestamppb.Now(),
		ConnectionId: 123,
		AgentId:      345,
	}
}
